diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 48b6730b0c6..c20454c7058 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4423,6 +4423,7 @@ dependencies = [ "pyo3-async-runtimes", "reqwest", "rstest", + "rustls 0.23.42", "serde", "serde_json", "serde_with", diff --git a/litellm-rust/crates/callbacks-legacy-python/python_contract.json b/litellm-rust/crates/callbacks-legacy-python/python_contract.json index ab887bbc014..8a212223523 100644 --- a/litellm-rust/crates/callbacks-legacy-python/python_contract.json +++ b/litellm-rust/crates/callbacks-legacy-python/python_contract.json @@ -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", diff --git a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs index 1bb84184e4c..f68541c5c7b 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs @@ -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>> { + match self.call.lookup(py, name)? { + Some(value) if from_py::(&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 { + 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> { + 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, context: &RequestContext, ) -> PyResult>> { + 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::()?; for (name, sent) in wire.body.as_object().into_iter().flatten() { - if let Some(value) = self.call.lookup(py, name)? - && from_py::(&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::()? .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 = 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 = 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, + #[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 = 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"); }); } diff --git a/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs b/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs index 1951c9ebfb4..093e86c43dc 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/callbacks.rs @@ -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; + + fn defer_success( + &self, + py: Python<'_>, + pending: &Bound<'_, PyAny>, + response: &Option>, + start: &Py, + end: &Option>, + ) -> 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 { + Logging::ClaimAsyncSuccess + .call(py, (self.object(py),))? + .extract() + } + + fn defer_success( + &self, + py: Python<'_>, + pending: &Bound<'_, PyAny>, + response: &Option>, + start: &Py, + end: &Option>, + ) -> 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") diff --git a/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs b/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs index 292648ba69c..ea72191ddb2 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs @@ -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, + 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::(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]) ", ); diff --git a/litellm-rust/crates/callbacks-legacy-python/src/python.rs b/litellm-rust/crates/callbacks-legacy-python/src/python.rs index 6d910d35f53..5b6faa3dbc2 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/python.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/python.rs @@ -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")] diff --git a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs index 8acfa7a898b..ee16e0c0059 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs @@ -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) diff --git a/litellm-rust/crates/host-python/src/binding.rs b/litellm-rust/crates/host-python/src/binding.rs index e36b7d70c28..fe6173dd100 100644 --- a/litellm-rust/crates/host-python/src/binding.rs +++ b/litellm-rust/crates/host-python/src/binding.rs @@ -28,6 +28,11 @@ pub trait PythonBinding: PythonOwned { ) -> PyResult>; /// 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<'_>, diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs index dbf843a5295..052156cbc5d 100644 --- a/litellm-rust/crates/host-python/src/driver.rs +++ b/litellm-rust/crates/host-python/src/driver.rs @@ -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, } } diff --git a/litellm-rust/crates/http/src/error.rs b/litellm-rust/crates/http/src/error.rs index 7078d4aba2a..9505aacf44b 100644 --- a/litellm-rust/crates/http/src/error.rs +++ b/litellm-rust/crates/http/src/error.rs @@ -16,6 +16,15 @@ pub enum TlsSource { #[error(transparent)] pub struct Error(Arc); +// 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, + 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) -> 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::().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 { + match &*self.0 { + Repr::BodyCut { cut, .. } => Some(*cut), + _ => None, + } + } + pub fn computed_header(name: impl Into) -> 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, + ) { + 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"); + } } diff --git a/litellm-rust/crates/http/src/lib.rs b/litellm-rust/crates/http/src/lib.rs index 02129c62083..1aa11843b46 100644 --- a/litellm-rust/crates/http/src/lib.rs +++ b/litellm-rust/crates/http/src/lib.rs @@ -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}; diff --git a/litellm-rust/crates/inference-messages/src/common_utils.rs b/litellm-rust/crates/inference-messages/src/common_utils.rs index e5ece2837d8..d0f109c1e99 100644 --- a/litellm-rust/crates/inference-messages/src/common_utils.rs +++ b/litellm-rust/crates/inference-messages/src/common_utils.rs @@ -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>, ) -> Result, 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}") ); } } diff --git a/litellm-rust/crates/inference-messages/src/handler.rs b/litellm-rust/crates/inference-messages/src/handler.rs index e0a7aa5adb8..c7a7d93bf98 100644 --- a/litellm-rust/crates/inference-messages/src/handler.rs +++ b/litellm-rust/crates/inference-messages/src/handler.rs @@ -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> + 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 { - let response = serde_json::from_str(text).map_err(|err| { - Error::invalid_response(litellm_llms::ErrorDetail::invalid( - "messages response JSON", - err, - )) - })?; +) -> Result { + let response = match serde_json::from_str::(text) { + Ok(response) => response, + Err(err) => { + let Ok(mut body) = serde_json::from_str::>(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) } diff --git a/litellm-rust/crates/inference-messages/src/lib.rs b/litellm-rust/crates/inference-messages/src/lib.rs index 745442d0314..25ebeaacb49 100644 --- a/litellm-rust/crates/inference-messages/src/lib.rs +++ b/litellm-rust/crates/inference-messages/src/lib.rs @@ -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)] diff --git a/litellm-rust/crates/inference-messages/src/route.rs b/litellm-rust/crates/inference-messages/src/route.rs index 7ff886689da..5f06d16d49f 100644 --- a/litellm-rust/crates/inference-messages/src/route.rs +++ b/litellm-rust/crates/inference-messages/src/route.rs @@ -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>; +pub type MessagesOutput = HostedCompletion; /// 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; + type Response = MessagesBody; type Error = Error; type Request = MessagesCall; type HostCall = Infallible; diff --git a/litellm-rust/crates/inference-messages/src/types.rs b/litellm-rust/crates/inference-messages/src/types.rs index ece17b3380c..772ea6f7292 100644 --- a/litellm-rust/crates/inference-messages/src/types.rs +++ b/litellm-rust/crates/inference-messages/src/types.rs @@ -42,7 +42,16 @@ pub(super) fn invalid_request(err: serde_json::Error) -> Error { } pub type MessagesCallResponse = - CallOutput, super::route::MessagesStreamHead, Bytes, Error>; + CallOutput; + +// 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), + Unchecked(Map), +} #[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] pub struct MessagesShaping { diff --git a/litellm-rust/crates/inference-messages/tests/bedrock_native/main.rs b/litellm-rust/crates/inference-messages/tests/bedrock_native/main.rs index f91d5c79f7c..5207f4e9d58 100644 --- a/litellm-rust/crates/inference-messages/tests/bedrock_native/main.rs +++ b/litellm-rust/crates/inference-messages/tests/bedrock_native/main.rs @@ -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 { 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"), } } diff --git a/litellm-rust/crates/inference-messages/tests/caching.rs b/litellm-rust/crates/inference-messages/tests/caching.rs index 663c87453a9..787501a8fed 100644 --- a/litellm-rust/crates/inference-messages/tests/caching.rs +++ b/litellm-rust/crates/inference-messages/tests/caching.rs @@ -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::( - 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::(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::( + 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::(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, ) { - 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(); diff --git a/litellm-rust/crates/inference-messages/tests/live/support.rs b/litellm-rust/crates/inference-messages/tests/live/support.rs index 81e295e34ad..52684928f1f 100644 --- a/litellm-rust/crates/inference-messages/tests/live/support.rs +++ b/litellm-rust/crates/inference-messages/tests/live/support.rs @@ -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 { 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"), } } diff --git a/litellm-rust/crates/inference-messages/tests/messages/host.rs b/litellm-rust/crates/inference-messages/tests/messages/host.rs index c3d06027c32..063e9d18738 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/host.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/host.rs @@ -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() + ) + ); +} diff --git a/litellm-rust/crates/inference-messages/tests/messages/main.rs b/litellm-rust/crates/inference-messages/tests/messages/main.rs index a01e571ec75..fe203287766 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/main.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/main.rs @@ -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 { 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") } diff --git a/litellm-rust/crates/inference-messages/tests/messages/request.rs b/litellm-rust/crates/inference-messages/tests/messages/request.rs index 216aa96fbe6..3784eb6cc68 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/request.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/request.rs @@ -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"}))] diff --git a/litellm-rust/crates/inference-messages/tests/messages/response.rs b/litellm-rust/crates/inference-messages/tests/messages/response.rs index 76813fe0a34..f978da5e995 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/response.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/response.rs @@ -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); diff --git a/litellm-rust/crates/inference/src/error.rs b/litellm-rust/crates/inference/src/error.rs index 841e0c53799..4135424d6ef 100644 --- a/litellm-rust/crates/inference/src/error.rs +++ b/litellm-rust/crates/inference/src/error.rs @@ -164,6 +164,14 @@ impl Error { } } + pub fn body_cut(&self) -> Option { + 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 { diff --git a/litellm-rust/crates/llms/src/error.rs b/litellm-rust/crates/llms/src/error.rs index 2844fbf3e01..2975ceee3cb 100644 --- a/litellm-rust/crates/llms/src/error.rs +++ b/litellm-rust/crates/llms/src/error.rs @@ -103,6 +103,13 @@ impl Error { matches!(self, Self::Http(error) if error.is_body_read_timeout()) } + pub fn body_cut(&self) -> Option { + match self { + Self::Http(error) => error.body_cut(), + _ => None, + } + } + pub fn is_timeout(&self) -> Option { match self { Self::Http(error) if error.is_timeout() => Some(error.is_connect_timeout()), diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 6ecfebf0096..79773ca5559 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -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 diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index 4e3847951d8..fe9e06b3d23 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -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")] diff --git a/litellm-rust/crates/python-bridge/src/cache/python/AGENTS.md b/litellm-rust/crates/python-bridge/src/cache/python/AGENTS.md index 81a0997f02a..3b516251f43 100644 --- a/litellm-rust/crates/python-bridge/src/cache/python/AGENTS.md +++ b/litellm-rust/crates/python-bridge/src/cache/python/AGENTS.md @@ -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 diff --git a/litellm-rust/crates/python-bridge/src/cache/python/host.rs b/litellm-rust/crates/python-bridge/src/cache/python/host.rs index 82b7d881df6..a8cbd0ce607 100644 --- a/litellm-rust/crates/python-bridge/src/cache/python/host.rs +++ b/litellm-rust/crates/python-bridge/src/cache/python/host.rs @@ -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, Error>>), @@ -21,6 +21,7 @@ pub(crate) struct PythonCache { arguments: Option>, pending: Option, 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; } diff --git a/litellm-rust/crates/python-bridge/src/cache/python/service.rs b/litellm-rust/crates/python-bridge/src/cache/python/service.rs index 38872b85d30..018c8647043 100644 --- a/litellm-rust/crates/python-bridge/src/cache/python/service.rs +++ b/litellm-rust/crates/python-bridge/src/cache/python/service.rs @@ -45,6 +45,7 @@ pub(crate) enum CacheCall { }, Store { value: Value, + unary: bool, reply: Reply>, }, } @@ -101,8 +102,13 @@ where ) -> Pin> + 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)? }) diff --git a/litellm-rust/crates/python-bridge/src/cache/selection.rs b/litellm-rust/crates/python-bridge/src/cache/selection.rs index e96f079e5da..78c1606ea40 100644 --- a/litellm-rust/crates/python-bridge/src/cache/selection.rs +++ b/litellm-rust/crates/python-bridge/src/cache/selection.rs @@ -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

(std::marker::PhantomData

); diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index 426cd861729..db5696d52ce 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -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)?; } diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 09b787b2c94..7513ff6b87b 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -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()) + ); + } } diff --git a/litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs b/litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs index ba016821e48..9da6a0e00e5 100644 --- a/litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs +++ b/litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs @@ -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, 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, PyNumber)>, - ) -> bool { + ) -> Option { self.0 .increment_many_post_call(&request.0, &self::increments(increments)) + .map(PyPostCallSettlement) } fn set_max<'py>( diff --git a/litellm-rust/crates/python-bridge/src/router/host.rs b/litellm-rust/crates/python-bridge/src/router/host.rs index 0ded98f27c5..27f9fe347e0 100644 --- a/litellm-rust/crates/python-bridge/src/router/host.rs +++ b/litellm-rust/crates/python-bridge/src/router/host.rs @@ -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]) -> 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], ids: &[String]) -> Vec> { - let mut taken = vec![false; candidates.len()]; + let mut positions: std::collections::HashMap<&str, Vec> = + 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> = ["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 = 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"); diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs index af614f17360..2a5f79ca393 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -239,9 +239,13 @@ impl PythonBinding for MessagesPythonHost { fn encode_response( &mut self, py: Python<'_>, - response: Box, + response: litellm_inference_messages::MessagesBody, ) -> PyResult> { - 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( diff --git a/litellm-rust/crates/redis-connect/Cargo.toml b/litellm-rust/crates/redis-connect/Cargo.toml index 72a453d36de..29b3293b856 100644 --- a/litellm-rust/crates/redis-connect/Cargo.toml +++ b/litellm-rust/crates/redis-connect/Cargo.toml @@ -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 diff --git a/litellm-rust/crates/redis-connect/src/connect.rs b/litellm-rust/crates/redis-connect/src/connect.rs index cf7fe2334d5..d11f6ba47ff 100644 --- a/litellm-rust/crates/redis-connect/src/connect.rs +++ b/litellm-rust/crates/redis-connect/src/connect.rs @@ -42,9 +42,6 @@ pub fn connect( host: HostServices, breaker: Option, ) -> Result { - 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; diff --git a/litellm-rust/crates/redis-connect/src/sources/gcp.rs b/litellm-rust/crates/redis-connect/src/sources/gcp.rs index 9da60fe9c08..9fb238709a1 100644 --- a/litellm-rust/crates/redis-connect/src/sources/gcp.rs +++ b/litellm-rust/crates/redis-connect/src/sources/gcp.rs @@ -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() { diff --git a/litellm-rust/crates/redis-store/AGENTS.md b/litellm-rust/crates/redis-store/AGENTS.md index cdd2bd7ca64..963f7ff168f 100644 --- a/litellm-rust/crates/redis-store/AGENTS.md +++ b/litellm-rust/crates/redis-store/AGENTS.md @@ -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`, 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`, 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`. 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` diff --git a/litellm-rust/crates/redis-store/src/spend_counters/keys.rs b/litellm-rust/crates/redis-store/src/spend_counters/keys.rs index 061e19ed1de..5221f8ab66c 100644 --- a/litellm-rust/crates/redis-store/src/spend_counters/keys.rs +++ b/litellm-rust/crates/redis-store/src/spend_counters/keys.rs @@ -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), Key { token: String, window: Option, @@ -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 { 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())), diff --git a/litellm-rust/crates/redis-store/src/spend_counters/mod.rs b/litellm-rust/crates/redis-store/src/spend_counters/mod.rs index 98f619cc501..425f83ef85c 100644 --- a/litellm-rust/crates/redis-store/src/spend_counters/mod.rs +++ b/litellm-rust/crates/redis-store/src/spend_counters/mod.rs @@ -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, + 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::>()).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, request: &RequestRedis, increments: &[(SpendCounter, Number)], - ) -> bool { - let Some(remote) = self.cache.redis().get() else { - return false; - }; + ) -> Option { + 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::>(); 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 = 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?) } diff --git a/litellm-rust/crates/redis-store/src/spend_counters/types.rs b/litellm-rust/crates/redis-store/src/spend_counters/types.rs index 61c211ec7e4..0509046d8c8 100644 --- a/litellm-rust/crates/redis-store/src/spend_counters/types.rs +++ b/litellm-rust/crates/redis-store/src/spend_counters/types.rs @@ -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); + +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 { diff --git a/litellm-rust/crates/redis-store/tests/spend_counters.rs b/litellm-rust/crates/redis-store/tests/spend_counters.rs index 75b7af2c321..7c89fb5b41e 100644 --- a/litellm-rust/crates/redis-store/tests/spend_counters.rs +++ b/litellm-rust/crates/redis-store/tests/spend_counters.rs @@ -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) -> Arc { 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::(&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::(&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]) diff --git a/litellm/caching/redis_batch.py b/litellm/caching/redis_batch.py index d58ddb9c629..b3e6d2eb291 100644 --- a/litellm/caching/redis_batch.py +++ b/litellm/caching/redis_batch.py @@ -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()) diff --git a/litellm/chat_completions/dispatch.py b/litellm/chat_completions/dispatch.py index 06ba8e9ad76..05572ee8837 100644 --- a/litellm/chat_completions/dispatch.py +++ b/litellm/chat_completions/dispatch.py @@ -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, ) diff --git a/litellm/constants.py b/litellm/constants.py index b43f6d9f358..6cca43f8c8b 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -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) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 5f7d646e150..58947f9d0c2 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -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 diff --git a/litellm/embeddings/dispatch.py b/litellm/embeddings/dispatch.py index 772ac68737d..9547757ff8f 100644 --- a/litellm/embeddings/dispatch.py +++ b/litellm/embeddings/dispatch.py @@ -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, ) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index fc1a601f9da..7984c7d04f3 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -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, ) diff --git a/litellm/litellm_core_utils/audio_utils/utils.py b/litellm/litellm_core_utils/audio_utils/utils.py index 808a424815e..a89ab5b0c6c 100644 --- a/litellm/litellm_core_utils/audio_utils/utils.py +++ b/litellm/litellm_core_utils/audio_utils/utils.py @@ -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) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 2cc9446c203..04072cd52c0 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -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 diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 0572ca66fc9..2e6f865b575 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -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 diff --git a/litellm/litellm_core_utils/prompt_templates/tool_call_json_repair.py b/litellm/litellm_core_utils/prompt_templates/tool_call_json_repair.py index 317ffa7b3b1..ace6d25c9f6 100644 --- a/litellm/litellm_core_utils/prompt_templates/tool_call_json_repair.py +++ b/litellm/litellm_core_utils/prompt_templates/tool_call_json_repair.py @@ -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 "", context or "unknown context", len(arguments), - arguments, - "..." if len(arguments) > 200 else "", + _logged_arguments(arguments), ) return repaired diff --git a/litellm/litellm_core_utils/routing_common_utils.py b/litellm/litellm_core_utils/routing_common_utils.py index d13ff4059aa..89d6729766f 100644 --- a/litellm/litellm_core_utils/routing_common_utils.py +++ b/litellm/litellm_core_utils/routing_common_utils.py @@ -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 () diff --git a/litellm/llms/anthropic/messages/native/handler.py b/litellm/llms/anthropic/messages/native/handler.py index 6e9ddba3042..ef1d08f4290 100644 --- a/litellm/llms/anthropic/messages/native/handler.py +++ b/litellm/llms/anthropic/messages/native/handler.py @@ -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 = ( diff --git a/litellm/llms/anthropic/messages/native/utils.py b/litellm/llms/anthropic/messages/native/utils.py index ea3afdb9a8f..8f33d000c7c 100644 --- a/litellm/llms/anthropic/messages/native/utils.py +++ b/litellm/llms/anthropic/messages/native/utils.py @@ -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 diff --git a/litellm/llms/anthropic/prompt_cache_prediction.py b/litellm/llms/anthropic/prompt_cache_prediction.py index 2da071ebdd2..28896c99eaf 100644 --- a/litellm/llms/anthropic/prompt_cache_prediction.py +++ b/litellm/llms/anthropic/prompt_cache_prediction.py @@ -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, diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 8ab8a43ba3a..b0ba7470433 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -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") diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index 6f0a18b0088..3866240030b 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -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) diff --git a/litellm/llms/bedrock/count_tokens/handler.py b/litellm/llms/bedrock/count_tokens/handler.py index 84067923788..eb29459f8b1 100644 --- a/litellm/llms/bedrock/count_tokens/handler.py +++ b/litellm/llms/bedrock/count_tokens/handler.py @@ -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 diff --git a/litellm/llms/bedrock/count_tokens/mantle_handler.py b/litellm/llms/bedrock/count_tokens/mantle_handler.py index 5e5c5516b94..0b589eda896 100644 --- a/litellm/llms/bedrock/count_tokens/mantle_handler.py +++ b/litellm/llms/bedrock/count_tokens/mantle_handler.py @@ -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 ) diff --git a/litellm/llms/bedrock/count_tokens/transformation.py b/litellm/llms/bedrock/count_tokens/transformation.py index cd59e714ba0..aaeb57303f9 100644 --- a/litellm/llms/bedrock/count_tokens/transformation.py +++ b/litellm/llms/bedrock/count_tokens/transformation.py @@ -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) diff --git a/litellm/llms/bedrock/web_identity_session_policy.py b/litellm/llms/bedrock/web_identity_session_policy.py index d567038cd1b..31187953dbc 100644 --- a/litellm/llms/bedrock/web_identity_session_policy.py +++ b/litellm/llms/bedrock/web_identity_session_policy.py @@ -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")), ) diff --git a/litellm/llms/minimax/chat/transformation.py b/litellm/llms/minimax/chat/transformation.py index f397bf7fbb9..84ed8a56430 100644 --- a/litellm/llms/minimax/chat/transformation.py +++ b/litellm/llms/minimax/chat/transformation.py @@ -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] diff --git a/litellm/llms/ollama/chat/transformation.py b/litellm/llms/ollama/chat/transformation.py index c6f3a7a8280..1e3e6186dd1 100644 --- a/litellm/llms/ollama/chat/transformation.py +++ b/litellm/llms/ollama/chat/transformation.py @@ -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) diff --git a/litellm/llms/opensearch/vector_stores/index_passthrough.py b/litellm/llms/opensearch/vector_stores/index_passthrough.py new file mode 100644 index 00000000000..9432af4c1b3 --- /dev/null +++ b/litellm/llms/opensearch/vector_stores/index_passthrough.py @@ -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//, for example /opensearch//_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 + ) diff --git a/litellm/llms/opensearch/vector_stores/transformation.py b/litellm/llms/opensearch/vector_stores/transformation.py index e153cdf43d8..f7ba28e7f3f 100644 --- a/litellm/llms/opensearch/vector_stores/transformation.py +++ b/litellm/llms/opensearch/vector_stores/transformation.py @@ -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" diff --git a/litellm/messages/dispatch.py b/litellm/messages/dispatch.py index 4c94f26c6c1..faa604a25ce 100644 --- a/litellm/messages/dispatch.py +++ b/litellm/messages/dispatch.py @@ -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, ) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 78f5f8c29b7..68bbf612e31 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -1970,7 +1970,7 @@ "source": "https://aws.amazon.com/bedrock/pricing/" }, "bedrock/claude-fable-latest": { - "alias_of": "global.anthropic.claude-fable-5-1" + "alias_of": "global.anthropic.claude-fable-5-2" }, "us.anthropic.claude-fable-5": { "cache_creation_input_token_cost": 1.375e-05, @@ -5343,7 +5343,7 @@ "supports_tool_choice": true }, "azure/eu/gpt-4o-realtime-preview-2024-12-17": { - "cache_read_input_audio_token_cost": 2.5e-06, + "cache_read_input_audio_token_cost": 2.75e-06, "cache_read_input_token_cost": 2.75e-06, "input_cost_per_audio_token": 4.4e-05, "input_cost_per_token": 5.5e-06, @@ -5353,7 +5353,7 @@ "mode": "realtime", "output_cost_per_audio_token": 8.8e-05, "output_cost_per_token": 2.2e-05, - "source": "https://azure.microsoft.com/en-us/pricing/details/azure-openai/", + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_modalities": [ "text", "audio" @@ -6691,6 +6691,7 @@ "supports_tool_choice": true }, "azure/gpt-realtime": { + "cache_read_input_image_token_cost": 5e-07, "cache_creation_input_audio_token_cost": 4e-06, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, @@ -6726,9 +6727,11 @@ "supports_parallel_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure" + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/" }, "azure/gpt-realtime-2025-08-28": { + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "cache_read_input_image_token_cost": 5e-07, "cache_creation_input_audio_token_cost": 4e-06, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, @@ -6766,6 +6769,7 @@ "supports_tool_choice": true }, "azure/gpt-realtime-1.5": { + "cache_read_input_image_token_cost": 5e-07, "cache_creation_input_audio_token_cost": 4e-06, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, @@ -6802,9 +6806,11 @@ "supports_parallel_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure" + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/" }, "azure/gpt-realtime-1.5-2026-02-23": { + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "cache_read_input_image_token_cost": 5e-07, "cache_creation_input_audio_token_cost": 4e-06, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, @@ -6925,6 +6931,8 @@ "supports_tool_choice": true }, "azure/gpt-realtime-mini": { + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "cache_read_input_image_token_cost": 8e-08, "cache_creation_input_audio_token_cost": 3e-07, "cache_read_input_audio_token_cost": 3e-07, "cache_read_input_token_cost": 6e-08, @@ -7033,6 +7041,7 @@ ] }, "azure/gpt-4o-realtime-preview-2024-12-17": { + "cache_read_input_audio_token_cost": 2.5e-06, "cache_read_input_token_cost": 2.5e-06, "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 5e-06, @@ -7042,7 +7051,7 @@ "mode": "realtime", "output_cost_per_audio_token": 8e-05, "output_cost_per_token": 2e-05, - "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure", + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_modalities": [ "text", "audio" @@ -12430,7 +12439,7 @@ "supports_tool_choice": true }, "azure/us/gpt-4o-realtime-preview-2024-12-17": { - "cache_read_input_audio_token_cost": 2.5e-06, + "cache_read_input_audio_token_cost": 2.75e-06, "cache_read_input_token_cost": 2.75e-06, "input_cost_per_audio_token": 4.4e-05, "input_cost_per_token": 5.5e-06, @@ -12440,7 +12449,7 @@ "mode": "realtime", "output_cost_per_audio_token": 8.8e-05, "output_cost_per_token": 2.2e-05, - "source": "https://azure.microsoft.com/en-us/pricing/details/azure-openai/", + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_modalities": [ "text", "audio" @@ -17082,7 +17091,7 @@ "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-fable-latest": { - "alias_of": "claude-fable-5-1" + "alias_of": "claude-fable-5-2" }, "claude-opus-5-5": { "supports_web_search": true, @@ -29363,52 +29372,6 @@ "comment": "Open flagship GLM for long-horizon coding agents and million-token context work", "source": "https://api.friendli.ai/serverless/v1/models" }, - "friendliai/deepseek-ai/DeepSeek-V3.2": { - "litellm_provider": "friendliai", - "max_input_tokens": 163840, - "max_output_tokens": 163840, - "input_cost_per_token": 5e-07, - "output_cost_per_token": 1.5e-06, - "cache_read_input_token_cost": 2.5e-07, - "supports_prompt_caching": true, - "supports_reasoning": true, - "reasoning_effort_levels": [], - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_native_structured_output": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": false, - "supports_image_input": false, - "supports_video_input": false, - "mode": "chat", - "comment": "DeepSeek chat model for instruction following, coding, and analysis", - "source": "https://api.friendli.ai/serverless/v1/models" - }, - "friendliai/MiniMaxAI/MiniMax-M2.5": { - "litellm_provider": "friendliai", - "max_input_tokens": 196608, - "max_output_tokens": 196608, - "input_cost_per_token": 3e-07, - "output_cost_per_token": 1.2e-06, - "cache_read_input_token_cost": 6e-08, - "supports_prompt_caching": true, - "supports_reasoning": true, - "reasoning_effort_levels": [], - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_native_structured_output": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": false, - "supports_image_input": false, - "supports_video_input": false, - "mode": "chat", - "comment": "Prior MiniMax coding model for agent workflows, office edits, and automation", - "source": "https://api.friendli.ai/serverless/v1/models" - }, "friendliai/zai-org/GLM-5.1": { "litellm_provider": "friendliai", "max_input_tokens": 202752, @@ -30575,7 +30538,6 @@ "successor_model": "gemini-3.8-flash", "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "cache_read_input_token_cost_above_200k_tokens_priority": 4.5e-07, "cache_read_input_token_cost_flex": 1.25e-07, "cache_read_input_token_cost_priority": 2.25e-07, @@ -30634,7 +30596,6 @@ "output_cost_per_token_priority": 1.8e-05 }, "gemini-3.1-pro-preview": { - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, @@ -30698,7 +30659,6 @@ "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, "input_cost_per_token_batches": 1e-06, @@ -30750,7 +30710,6 @@ "vertex_ai/gemini-3-pro-preview": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, "input_cost_per_token_batches": 1e-06, @@ -31194,7 +31153,6 @@ "supports_web_search": true }, "vertex_ai/gemini-3.1-pro-preview": { - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, @@ -31258,7 +31216,6 @@ "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, "input_cost_per_token_batches": 1e-06, @@ -43150,6 +43107,29 @@ "max_input_tokens": 1000000, "max_output_tokens": 128000 }, + "minimax/MiniMax-M3.1-Flash": { + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "cache_read_input_token_cost": 1.5e-08, + "litellm_provider": "minimax", + "mode": "chat", + "source": "https://platform.minimax.io/docs/guides/pricing-paygo", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_none_reasoning_effort": false, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "reasoning_effort_levels": [ + "low", + "medium", + "high", + "xhigh", + "max" + ], + "supports_vision": true, + "max_input_tokens": 1000000 + }, "minimax/minimax-latest": { "alias_of": "minimax/MiniMax-M3" }, @@ -53162,8 +53142,49 @@ "output_cost_per_token_batches": 2.5e-05, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" }, + "vertex_ai/claude-fable-5-2": { + "regional_endpoint_uplift_multiplier": 1.1, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_creation_input_token_cost_batches": 5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "prompt_cache_min_tokens": 512, + "deprecation_date": "2027-11-09", + "input_cost_per_token_batches": 4e-06, + "output_cost_per_token_batches": 2e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" + }, "vertex_ai/claude-fable-latest": { - "alias_of": "vertex_ai/claude-fable-5-1" + "alias_of": "vertex_ai/claude-fable-5-2" }, "vertex_ai/claude-fable-5@default": { "deprecation_date": "2027-06-08", @@ -53245,6 +53266,47 @@ "output_cost_per_token_batches": 2.5e-05, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" }, + "vertex_ai/claude-fable-5-2@default": { + "regional_endpoint_uplift_multiplier": 1.1, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_creation_input_token_cost_batches": 5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "prompt_cache_min_tokens": 512, + "deprecation_date": "2027-11-09", + "input_cost_per_token_batches": 4e-06, + "output_cost_per_token_batches": 2e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" + }, "vertex_ai/claude-opus-5": { "deprecation_date": "2027-01-24", "regional_endpoint_uplift_multiplier": 1.1, @@ -65172,6 +65234,14 @@ "supports_reasoning": true } }, + { + "name": "minimax-reasoning-baseline", + "pattern": "^minimax/", + "description": "Any MiniMax model id under the minimax/ namespace, anchored so only that provider's ids match. MiniMax ships reasoning models faster than this registry names them, so an id the map has not described yet is treated as reasoning-capable and keeps the caller's reasoning_effort and thinking instead of dropping them. Rules lose to exact entries, so mapped speech, video and image models are unaffected. Carries no mode and no pricing, so cost stays on the standard unpriced behavior.", + "model_info": { + "supports_reasoning": true + } + }, { "name": "openai-reasoning-family-baseline", "pattern": "^(?!.*search-api)(?:[a-z0-9_.-]+/)*(?:ft:)?(?:o[1-9]\\d*(?![a-z0-9])|gpt-[5-9](?:\\.\\d+)?(?![0-9.])|(?:gpt-\\d+(?:\\.\\d+)?(?:-[a-z0-9]+)*-)?(?:codex|deep-research|chat-latest)(?![a-z0-9]))", @@ -67761,7 +67831,7 @@ "deepinfra/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B": { "max_output_tokens": 262144, "max_input_tokens": 262144, - "input_cost_per_token": 8.5e-08, + "input_cost_per_token": 7.5e-08, "output_cost_per_token": 4e-07, "litellm_provider": "deepinfra", "mode": "chat", @@ -69888,6 +69958,44 @@ "region": "us-gov-west-1" } }, + "bedrock/us-gov-west-1/anthropic.claude-fable-5-2": { + "cache_creation_input_token_cost": 1.2e-05, + "cache_creation_input_token_cost_above_1hr": 1.92e-05, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 9.6e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.8e-05, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh", + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512, + "token_count_model": { + "model": "anthropic.claude-fable-5-2", + "litellm_provider": "bedrock_mantle", + "region": "us-gov-west-1" + } + }, "bedrock/us-gov-east-1/nvidia.nemotron-nano-3-30b": { "input_cost_per_token": 7.2e-08, "input_cost_per_token_batches": 3.6e-08, @@ -70196,6 +70304,39 @@ "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512 }, + "bedrock/us-gov-east-1/anthropic.claude-fable-5-2": { + "cache_creation_input_token_cost": 1.2e-05, + "cache_creation_input_token_cost_above_1hr": 1.92e-05, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 9.6e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.8e-05, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh", + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512 + }, "bedrock_mantle/us-gov-west-1/openai.gpt-5.6-terra": { "litellm_provider": "bedrock_mantle", "max_input_tokens": 922000, @@ -72924,6 +73065,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/gpt-realtime-2": { + "cache_read_input_image_token_cost": 5e-07, "cache_creation_input_audio_token_cost": 4e-07, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, @@ -72942,7 +73084,7 @@ "provider_specific_entry": { "realtime_ga_only": 1 }, - "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_endpoints": [ "/v1/realtime" ], @@ -83307,13 +83449,16 @@ "cache_creation_input_token_cost": 2.5e-05, "cache_creation_input_token_cost_above_12hr": 5e-05, "cache_read_input_token_cost": 2e-06, + "cache_read_input_token_cost_above_272k_tokens": 4e-06, "input_cost_per_token": 2e-05, + "input_cost_per_token_above_272k_tokens": 4e-05, "litellm_provider": "openai", "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "output_cost_per_token": 0.0001, - "source": "https://developers.openai.com/api/docs/models/gpt-6-cyber", + "output_cost_per_token_above_272k_tokens": 0.00015, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -87201,13 +87346,13 @@ "supports_vision": true }, "wandb/deepseek-ai/DeepSeek-V4.1-Flash-1017": { - "cache_read_input_token_cost": 2e-09, - "input_cost_per_token": 1.3e-07, + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, "litellm_provider": "wandb", "max_input_tokens": 1049000, "max_output_tokens": 1048576, "mode": "chat", - "output_cost_per_token": 6.5e-07, + "output_cost_per_token": 1.2e-06, "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference", "supports_prompt_caching": true, "supports_reasoning": true, @@ -87352,5 +87497,574 @@ "supports_parallel_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true + }, + "anthropic.claude-fable-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "eu.anthropic.claude-fable-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.1e-05, + "cache_creation_input_token_cost_above_1hr": 1.76e-05, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 8.8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output_by_default": true, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "eu.anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.1e-05, + "cache_creation_input_token_cost_above_1hr": 1.76e-05, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 8.8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "global.anthropic.claude-fable-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output_by_default": true, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "global.anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "us-gov.anthropic.claude-fable-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.2e-05, + "cache_creation_input_token_cost_above_1hr": 1.92e-05, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 9.6e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.8e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "us-gov.anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.2e-05, + "cache_creation_input_token_cost_above_1hr": 1.92e-05, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 9.6e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.8e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "us.anthropic.claude-fable-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.1e-05, + "cache_creation_input_token_cost_above_1hr": 1.76e-05, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 8.8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output_by_default": true, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "us.anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.1e-05, + "cache_creation_input_token_cost_above_1hr": 1.76e-05, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 8.8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "claude-fable-5-2": { + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_creation_input_token_cost_batches": 5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, + "input_cost_per_token": 8e-06, + "input_cost_per_token_batches": 4e-06, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "output_cost_per_token_batches": 2e-05, + "prompt_cache_min_tokens": 512.0, + "provider_specific_entry": {"us":1.1,"fast":2}, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://platform.claude.com/docs/en/about-claude/pricing", + "supports_adaptive_thinking": true, + "supports_anthropic_compaction": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_fast_mode": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_speed": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "bedrock/claude_platform/claude-fable-5-2": { + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_creation_input_token_cost_batches": 5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, + "input_cost_per_token": 8e-06, + "input_cost_per_token_batches": 4e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "output_cost_per_token_batches": 2e-05, + "prompt_cache_min_tokens": 512.0, + "provider_specific_entry": {"us":1.1}, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing", + "supports_adaptive_thinking": true, + "supports_anthropic_compaction": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "azure_ai/claude-fable-5-2": { + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://platform.claude.com/docs/en/about-claude/pricing", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": false, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "azure_ai/claude-fable-5-2-fast": { + "cache_creation_input_token_cost": 2e-05, + "cache_creation_input_token_cost_above_1hr": 3.2e-05, + "cache_read_input_token_cost": 4e-07, + "input_cost_per_token": 1.6e-05, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 8e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://platform.claude.com/docs/en/about-claude/pricing", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": false, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "bedrock/us-gov-east-1/anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.2e-05, + "cache_creation_input_token_cost_above_1hr": 1.92e-05, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 9.6e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.8e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "bedrock/us-gov-west-1/anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.2e-05, + "cache_creation_input_token_cost_above_1hr": 1.92e-05, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 9.6e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.8e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "github_copilot/claude-fable-5.2": { + "cache_creation_input_token_cost": 1e-05, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "github_copilot", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/messages" + ], + "supports_adaptive_thinking": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "thinking_always_on": true } } diff --git a/litellm/ocr/dispatch.py b/litellm/ocr/dispatch.py index 529c3df2bb9..f90a42bb1c9 100644 --- a/litellm/ocr/dispatch.py +++ b/litellm/ocr/dispatch.py @@ -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, ) diff --git a/litellm/proxy/_lazy_features.py b/litellm/proxy/_lazy_features.py index 680b0e6ede4..85f71defe3b 100644 --- a/litellm/proxy/_lazy_features.py +++ b/litellm/proxy/_lazy_features.py @@ -243,6 +243,7 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = ( "/mistral/", "/nvidia_nim/", "/openai/", + "/opensearch/", "/openai_passthrough/", "/tinyfish/", "/transcribe", diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index fde5be049eb..746d6a1e8d7 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -35987,6 +35987,386 @@ ] } }, + "/opensearch/{endpoint}": { + "delete": { + "description": "Call an OpenSearch index through the proxy, e.g. `POST /opensearch//_search`.\n\n`` 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//_search`.\n\n`` 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//_search`.\n\n`` 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//_search`.\n\n`` 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//_search`.\n\n`` 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//_search`.\n\n`` 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)", diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index e53b5d4ad52..6b7aa4d15d2 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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.", diff --git a/litellm/proxy/auth/access_group_rules.py b/litellm/proxy/auth/access_group_rules.py new file mode 100644 index 00000000000..b8e3d23f2cc --- /dev/null +++ b/litellm/proxy/auth/access_group_rules.py @@ -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) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 4514882a45f..f0c9632837d 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -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, diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 230dc7f6c58..d3978a58c43 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -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"], diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 9e11179c686..4c4e506ce38 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -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( diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index a52ad14e293..9a3b18cf95c 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -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, diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 8f73d47d96b..eb356491014 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -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 diff --git a/litellm/proxy/common_utils/callback_config_validation.py b/litellm/proxy/common_utils/callback_config_validation.py index 302978d17a4..a486d6665fc 100644 --- a/litellm/proxy/common_utils/callback_config_validation.py +++ b/litellm/proxy/common_utils/callback_config_validation.py @@ -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 diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index c195779aa0a..0aa8d2291ef 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -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 diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index 08461aa3a44..80ba2f49745 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -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" diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 131f57568e8..9c06e3afc18 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -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 _: diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 4ed3ac30be7..ae45b43f15f 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -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: diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 68a6177cabc..dd7a24ae879 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -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: diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 8a23262acc1..0bb4a803609 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -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( diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index 0640bf0fa4b..672bbdb164a 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -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, diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 6c6118b2505..f88ac445a0b 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -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: diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index d54eda45160..038c1bdf5a6 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -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, diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 15d05633679..bc1963af0bb 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -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 diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 59fa8a18ac7..210aa1aaa6f 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -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, ...]: diff --git a/litellm/proxy/management_helpers/bulk_user_creation.py b/litellm/proxy/management_helpers/bulk_user_creation.py index cca37ebd7ed..77e37cb11cb 100644 --- a/litellm/proxy/management_helpers/bulk_user_creation.py +++ b/litellm/proxy/management_helpers/bulk_user_creation.py @@ -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: diff --git a/litellm/proxy/management_helpers/bulk_user_deletion.py b/litellm/proxy/management_helpers/bulk_user_deletion.py index 04dd0b3f57f..cc60f58efb0 100644 --- a/litellm/proxy/management_helpers/bulk_user_deletion.py +++ b/litellm/proxy/management_helpers/bulk_user_deletion.py @@ -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) diff --git a/litellm/proxy/middleware/billable_request_metrics_middleware.py b/litellm/proxy/middleware/billable_request_metrics_middleware.py index fd78b6b58da..95e5c9d1293 100644 --- a/litellm/proxy/middleware/billable_request_metrics_middleware.py +++ b/litellm/proxy/middleware/billable_request_metrics_middleware.py @@ -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 diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index ba8153ebb6d..5bdb2219194 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -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//_search`. + + `` 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) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 53317010af4..4b84fe19559 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -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): diff --git a/litellm/proxy/plugin_routes.py b/litellm/proxy/plugin_routes.py index cfbcfefe9e8..966136f6fa7 100644 --- a/litellm/proxy/plugin_routes.py +++ b/litellm/proxy/plugin_routes.py @@ -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() diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d21b6169133..37dc135d4d7 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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) diff --git a/litellm/proxy/spend_tracking/carried_budget_state.py b/litellm/proxy/spend_tracking/carried_budget_state.py index 0dfb38272b0..150ddd1efe9 100644 --- a/litellm/proxy/spend_tracking/carried_budget_state.py +++ b/litellm/proxy/spend_tracking/carried_budget_state.py @@ -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 diff --git a/litellm/proxy/spend_tracking/committed_spend.py b/litellm/proxy/spend_tracking/committed_spend.py new file mode 100644 index 00000000000..76066339cf1 --- /dev/null +++ b/litellm/proxy/spend_tracking/committed_spend.py @@ -0,0 +1,181 @@ +"""Committed spend counters: twins of the entity spend counters that only ever take charged cost, never budget holds. + +The remaining budget gauges read them, so a request's worst-case reservation never shows as spent. They are kept +only for entities with a budget, while a ``PrometheusLogger`` is registered, in the same ``SpendCounterStore`` as the +counters they twin. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import Iterable, Mapping, Sequence +from typing import TYPE_CHECKING, Final + +import litellm +from litellm._internal_context import with_service_target +from litellm._logging import verbose_proxy_logger +from litellm.caching.redis_batch import Batch +from litellm.integrations.prometheus import PrometheusLogger +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed +from litellm.proxy.spend_tracking.spend_counter_batch import ( + admission_counter_keys, + bind_spend_counter_keys, + read_batched_spend_counter, +) +from litellm.redis.stores.spend_counters import ( + COMMITTED_COUNTER_PREFIX, + END_USER_COUNTER_PREFIX, + SPEND_COUNTERS_TARGET, + SpendCounterStore, + committed_counter, +) +from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + +_BASE_COUNTER_PREFIXES: Final = ("spend:key:", "spend:team:", "spend:user:", "spend:org:", END_USER_COUNTER_PREFIX) + + +def committed_spend_key(counter_key: str) -> str: + return committed_counter(counter_key) + + +def has_committed_twin(counter_key: str) -> bool: + return counter_key.startswith(_BASE_COUNTER_PREFIXES) and ":window:" not in counter_key + + +def committed_spend_keys(counter_keys: Iterable[str]) -> tuple[str, ...]: + return tuple(committed_spend_key(counter_key) for counter_key in counter_keys if has_committed_twin(counter_key)) + + +def is_committed_spend_key(counter_key: str) -> bool: + return counter_key.startswith(COMMITTED_COUNTER_PREFIX) + + +def committed_spend_tracked() -> bool: + return litellm.logging_callback_manager.callback_is_active(callback_type=PrometheusLogger) + + +def bind_committed_admission_counters(token: UserAPIKeyAuth, end_user_id: str | None) -> None: + if committed_spend_tracked(): + bind_spend_counter_keys(frozenset(committed_spend_keys(admission_counter_keys(token, end_user_id)))) + + +async def _read_committed_values( + spend_counter_store: SpendCounterStore, twin_keys: tuple[str, ...] +) -> Mapping[str, float]: + batched: Final = await spend_counter_store.read_many(twin_keys, batch=Batch.NONE) + if batched is not None: + return {twin_key: value for twin_key, value in batched.items() if value is not None} + reads: Final = await asyncio.gather(*(spend_counter_store.read(twin_key) for twin_key in twin_keys)) + return {twin_key: value for twin_key, (value, _) in zip(twin_keys, reads, strict=True) if value is not None} + + +def _committed_budget_spend( + counter_keys: EntitySpendCounterKeys, committed: Mapping[str, float] +) -> CommittedBudgetSpend: + def spend_of(counter_key: str | None) -> float | None: + return committed.get(committed_spend_key(counter_key)) if counter_key is not None else None + + return CommittedBudgetSpend( + counter_keys=counter_keys, + key=spend_of(counter_keys.key), + team=spend_of(counter_keys.team), + user=spend_of(counter_keys.user), + org=spend_of(counter_keys.org), + end_user=spend_of(counter_keys.end_user), + ) + + +@with_service_target(SPEND_COUNTERS_TARGET) +async def read_committed_budget_spend( + spend_counter_store: SpendCounterStore, counter_keys: EntitySpendCounterKeys +) -> CommittedBudgetSpend: + twin_keys: Final = committed_spend_keys(counter_keys.tracked()) + if not twin_keys: + return CommittedBudgetSpend(counter_keys=counter_keys) + try: + committed: Final = await _read_committed_values(spend_counter_store, twin_keys) + except Exception as e: # noqa: BLE001 # a gauge read never fails the request; the gauges keep the stored spend + verbose_proxy_logger.debug("committed spend read failed for %s: %s", twin_keys, e) + return CommittedBudgetSpend(counter_keys=counter_keys) + return _committed_budget_spend(counter_keys=counter_keys, committed=committed) + + +async def warm_committed_counter_keys( + spend_counter_store: SpendCounterStore, counter_keys: EntitySpendCounterKeys +) -> EntitySpendCounterKeys: + """For a charge with no auth to name the budgeted entities: only a budgeted entity's twin exists.""" + committed: Final = await read_committed_budget_spend(spend_counter_store, counter_keys) + return EntitySpendCounterKeys( + key=counter_keys.key if committed.key is not None else None, + team=counter_keys.team if committed.team is not None else None, + user=counter_keys.user if committed.user is not None else None, + org=counter_keys.org if committed.org is not None else None, + end_user=counter_keys.end_user if committed.end_user is not None else None, + ) + + +async def read_admission_committed_spend( + spend_counter_store: SpendCounterStore, counter_keys: EntitySpendCounterKeys +) -> CommittedBudgetSpend: + """Rides the admission MGET: a counter the batch cannot answer stays unknown here and the emit-time read has it.""" + if not spend_counter_store.has_redis: + return await read_committed_budget_spend(spend_counter_store, counter_keys) + twin_keys: Final = committed_spend_keys(counter_keys.tracked()) + batched: Final = await asyncio.gather(*(read_batched_spend_counter(twin_key) for twin_key in twin_keys)) + return _committed_budget_spend( + counter_keys=counter_keys, + committed={ + twin_key: read[0] for twin_key, read in zip(twin_keys, batched) if read is not None and read[0] is not None + }, + ) + + +async def drop_committed_spend(spend_counter_store: SpendCounterStore, counter_keys: Sequence[str]) -> None: + """After a direct write to the stored spend the twin holds the old total, so the next charge seeds it again.""" + twin_keys: Final = committed_spend_keys(counter_keys) + if not twin_keys: + return + try: + await spend_counter_store.delete(twin_keys) + except Exception as e: # noqa: BLE001 # the write already committed; a kept twin lasts until charges stop for its TTL + verbose_proxy_logger.warning("committed spend drop failed: %s", e) + + +@with_service_target(SPEND_COUNTERS_TARGET) +async def seed_committed_spend( + spend_counter_store: SpendCounterStore, prisma_client: PrismaClient | None, counter_key: str +) -> bool: + """One DB read per worker for a cold twin: charges waiting on the lock find the twin the first one wrote.""" + twin_key: Final = committed_spend_key(counter_key) + async with await SpendCounterReseed.counter_lock(twin_key): + try: + stored, _ = await spend_counter_store.read(twin_key) + except Exception as e: # noqa: BLE001 # an unseeded twin only skips this charge; the next one seeds it again + verbose_proxy_logger.debug("committed spend seed read failed for %s: %s", twin_key, e) + return False + return stored is not None or await _seed_committed_spend_from_db( + spend_counter_store, prisma_client, counter_key + ) + + +async def _seed_committed_spend_from_db( + spend_counter_store: SpendCounterStore, prisma_client: PrismaClient | None, counter_key: str +) -> bool: + db_spend: Final = await ( + SpendCounterReseed.end_user_from_db(prisma_client=prisma_client, counter_key=counter_key) + if counter_key.startswith(END_USER_COUNTER_PREFIX) + else SpendCounterReseed.from_db(prisma_client=prisma_client, counter_key=counter_key) + ) + if db_spend is None: + return False + twin_key: Final = committed_spend_key(counter_key) + try: + await spend_counter_store.seed(twin_key, db_spend, increment_if_lost=False) + except Exception as e: # noqa: BLE001 # an unseeded twin only skips this charge; the next one seeds it again + verbose_proxy_logger.debug("committed spend seed failed for %s: %s", twin_key, e) + return False + return True diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 0348d8e1f82..f3dc672a072 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -2,7 +2,7 @@ import collections import json import os -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass from datetime import date, datetime, timedelta, timezone from functools import partial @@ -53,6 +53,7 @@ from litellm.proxy.auth.authorization_dependencies import ( from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.openai_error_payload import error_status_code, openai_error_param, openai_error_type from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup +from litellm.proxy.spend_tracking.committed_spend import committed_spend_tracked, drop_committed_spend from litellm.proxy.spend_tracking.session_page import ( AGENT_CALL_TYPE_SQL, MCP_CALL_TYPES_SQL, @@ -74,6 +75,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( get_spend_by_team_and_customer, ) from litellm.proxy.utils import handle_exception_on_proxy, query_raw_within_statement_timeout +from litellm.redis.stores.spend_counters import key_counter, team_counter from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import SpendLogsRepository from litellm.repositories.team_repository import TeamRepository @@ -3595,6 +3597,50 @@ async def view_spend_logs( ) +class _EntityIdRow(TypedDict): + id: ReadOnly[str] + + +_COMMITTED_DROP_PAGE_SIZE: Final = 1000 +_CommittedDropEntity: TypeAlias = Literal["key", "team"] +_COMMITTED_DROP_PAGE_QUERIES: Final[Mapping[_CommittedDropEntity, str]] = MappingProxyType( + { + "key": 'SELECT token AS id FROM "LiteLLM_VerificationToken" WHERE token > $1 ORDER BY token LIMIT $2', + "team": 'SELECT team_id AS id FROM "LiteLLM_TeamTable" WHERE team_id > $1 ORDER BY team_id LIMIT $2', + } +) +_COMMITTED_DROP_COUNTERS: Final[Mapping[_CommittedDropEntity, Callable[[str], str]]] = MappingProxyType( + {"key": key_counter, "team": team_counter} +) + + +async def _drop_committed_spend_page( + prisma_client: PrismaClient, entity: _CommittedDropEntity, after_id: str +) -> str | None: + from litellm.proxy.proxy_server import spend_counter_store + + rows: Final[Sequence[_EntityIdRow]] = await _query_raw( + prisma_client, _COMMITTED_DROP_PAGE_QUERIES[entity], after_id, _COMMITTED_DROP_PAGE_SIZE + ) + counter_of: Final = _COMMITTED_DROP_COUNTERS[entity] + await drop_committed_spend(spend_counter_store, tuple(counter_of(row["id"]) for row in rows)) + return rows[-1]["id"] if len(rows) == _COMMITTED_DROP_PAGE_SIZE else None + + +async def _drop_committed_spend_pages(prisma_client: PrismaClient, entity: _CommittedDropEntity) -> None: + after_id: str | None = "" # rebind-ok: the keyset cursor moves one page at a time + while after_id is not None: + after_id = await _drop_committed_spend_page(prisma_client, entity, after_id) + + +async def _drop_all_key_and_team_committed_spend(prisma_client: PrismaClient) -> None: + try: + await _drop_committed_spend_pages(prisma_client, "key") + await _drop_committed_spend_pages(prisma_client, "team") + except Exception as e: # noqa: BLE001 # the reset already committed; a kept twin lasts until charges stop for its TTL + verbose_proxy_logger.warning("committed spend drop after a global reset failed: %s", e) + + @router.post( "/global/spend/reset", tags=["Budget & Spend Tracking"], @@ -3626,6 +3672,8 @@ async def global_spend_reset(): _ = await transaction.litellm_teamtable.update_many(data={"spend": 0.0}, where={}) _ = await transaction.query_raw(_LOCK_ALL_KEYS_SQL) _ = await transaction.litellm_verificationtoken.update_many(data={"spend": 0.0}, where={}) + if committed_spend_tracked(): + await _drop_all_key_and_team_committed_spend(prisma_client) return { "message": "Spend for all API Keys and Teams reset successfully", diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8b9189f143f..97c327cbe4e 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -9263,13 +9263,14 @@ async def monitor_spend_logs_queue( proxy_logging_obj: Proxy logging object """ from litellm.constants import ( + SPEND_LOG_QUEUE_MAX_BACKOFF_SECONDS, SPEND_LOG_QUEUE_POLL_INTERVAL, SPEND_LOG_QUEUE_SIZE_THRESHOLD, ) threshold: Final = SPEND_LOG_QUEUE_SIZE_THRESHOLD base_interval: Final = SPEND_LOG_QUEUE_POLL_INTERVAL - max_backoff: Final = 30.0 # Maximum backoff interval in seconds + max_backoff: Final = SPEND_LOG_QUEUE_MAX_BACKOFF_SECONDS backoff_multiplier: Final = 1.5 # Exponential backoff multiplier current_interval = base_interval flush_requested: Final = asyncio.Event() diff --git a/litellm/proxy/vector_store_endpoints/utils.py b/litellm/proxy/vector_store_endpoints/utils.py index 8ce48fc3516..4dec4c7b38b 100644 --- a/litellm/proxy/vector_store_endpoints/utils.py +++ b/litellm/proxy/vector_store_endpoints/utils.py @@ -25,7 +25,7 @@ from litellm.utils import ProviderConfigManager _JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) -def _normalize_litellm_params( +def normalize_litellm_params( vector_store: LiteLLM_ManagedVectorStore, ) -> LiteLLM_ManagedVectorStore: litellm_params: Final = vector_store.get("litellm_params") @@ -46,6 +46,10 @@ def _is_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> bool: ) +def _key_metadata(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, object]: + return user_api_key_dict.metadata + + def assert_proxy_admin_for_vector_store_index_management( user_api_key_dict: UserAPIKeyAuth, *, @@ -278,7 +282,7 @@ async def get_litellm_managed_vector_store( vector_store_id=vector_store_id ) if vector_store is not None: - return _normalize_litellm_params(vector_store) + return normalize_litellm_params(vector_store) except Exception as e: verbose_proxy_logger.warning( "Failed to resolve vector store id=%s from registry: %s", @@ -310,7 +314,7 @@ async def get_litellm_managed_vector_store( ) if not rows: return None - return _normalize_litellm_params(LiteLLM_ManagedVectorStore(**rows[0].model_dump())) + return normalize_litellm_params(LiteLLM_ManagedVectorStore(**rows[0].model_dump())) except Exception as e: verbose_proxy_logger.warning( "Failed to resolve vector store id=%s from shared cache: %s", @@ -413,6 +417,34 @@ def check_vector_store_permission( return False +def vector_store_endpoint_forbidden(index_name: str) -> HTTPException: + return HTTPException( + status_code=403, + detail=( + f"User does not have permission to call vector store endpoint {index_name}. " + "Ask your administrator to add the necessary permissions to your API key/Team." + ), + ) + + +def has_vector_store_index_access( + user_api_key_dict: UserAPIKeyAuth, index_name: str, permission: Literal["read", "write"] | None +) -> bool: + return _is_proxy_admin(user_api_key_dict) or ( + permission is not None + and check_vector_store_permission( + index_name, permission, _key_metadata(user_api_key_dict), user_api_key_dict.team_metadata + ) + ) + + +def assert_vector_store_index_access( + user_api_key_dict: UserAPIKeyAuth, index_name: str, permission: Literal["read", "write"] | None +) -> None: + if not has_vector_store_index_access(user_api_key_dict, index_name, permission): + raise vector_store_endpoint_forbidden(index_name) + + def is_allowed_to_call_vector_store_endpoint( provider: LlmProviders, index_name: str, @@ -432,7 +464,7 @@ def is_allowed_to_call_vector_store_endpoint( ): return True # check what allowed permissions are for the key - key_metadata: Final = user_api_key_dict.metadata + key_metadata: Final = _key_metadata(user_api_key_dict) team_metadata: Final = user_api_key_dict.team_metadata provider_config: Final = ProviderConfigManager.get_provider_vector_stores_config(provider=provider) @@ -478,14 +510,7 @@ def is_allowed_to_call_vector_store_endpoint( break if permission_type is None: - raise HTTPException( - status_code=403, - detail=( - f"User does not have permission to call vector store endpoint " - f"{index_name}. Ask your administrator to add the necessary " - "permissions to your API key/Team." - ), - ) + raise vector_store_endpoint_forbidden(index_name) # Check if key has specific permission for allowed_vector_store_indexes has_permission: Final = check_vector_store_permission( @@ -496,10 +521,7 @@ def is_allowed_to_call_vector_store_endpoint( ) if not has_permission: - raise HTTPException( - status_code=403, - detail=f"User does not have permission to call vector store endpoint {index_name}. Ask your administrator to add the necessary permissions to your API key/Team.", - ) + raise vector_store_endpoint_forbidden(index_name) return has_permission @@ -516,7 +538,7 @@ def is_allowed_to_call_vector_store_files_endpoint( ): return True - key_metadata: Final = user_api_key_dict.metadata + key_metadata: Final = _key_metadata(user_api_key_dict) team_metadata: Final = user_api_key_dict.team_metadata provider_config: Final = ProviderConfigManager.get_provider_vector_store_files_config(provider=provider) diff --git a/litellm/redis/stores/spend_counters/__init__.py b/litellm/redis/stores/spend_counters/__init__.py index 3ad06589f53..6a6dc85586e 100644 --- a/litellm/redis/stores/spend_counters/__init__.py +++ b/litellm/redis/stores/spend_counters/__init__.py @@ -6,15 +6,17 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Final from litellm._internal_context import with_service_target -from litellm.caching.redis_batch import Batch +from litellm.caching.redis_batch import Batch, PostCallSettlement from litellm.redis.stores._route import post_call_request, ready, route from litellm.redis.stores.spend_counters.handle import SpendCountersHandle from litellm.redis.stores.spend_counters.types import ( + COMMITTED_COUNTER_PREFIX, END_USER_COUNTER_PREFIX, SPEND_COUNTER_PREFIX, SPEND_COUNTERS_TARGET, SpendIncrement, SpendReading, + committed_counter, end_user_counter, is_unavailable, key_counter, @@ -34,12 +36,14 @@ if TYPE_CHECKING: from litellm.redis.stores.connection import RedisStores __all__ = [ + "COMMITTED_COUNTER_PREFIX", "END_USER_COUNTER_PREFIX", "SPEND_COUNTERS_TARGET", "SPEND_COUNTER_PREFIX", "SpendCounterStore", "SpendIncrement", "SpendReading", + "committed_counter", "end_user_counter", "is_unavailable", "key_counter", @@ -81,6 +85,11 @@ class SpendCounterStore: def has_redis(self) -> bool: return self._handle.has_redis + def set_committed_ttl(self, ttl: float) -> None: + """The least TTL a Redis write gives a committed twin, which keeps it past the DB write that holds its + charges.""" + self._handle.set_committed_ttl(ttl) + @with_service_target(SPEND_COUNTERS_TARGET) async def read(self, counter: str) -> SpendReading: reading: Final = await self._handle.read(counter) @@ -104,12 +113,13 @@ class SpendCounterStore: return tuple(await self._handle.increment_many(increments)) @with_service_target(SPEND_COUNTERS_TARGET) - def increment_many_post_call(self, increments: Sequence[SpendIncrement]) -> bool: - """Declares the increments on the request's post-call batch, False when there is no Redis or request. - Each total lands in memory once the batch goes out; a counter whose increment failed is dropped everywhere.""" + def increment_many_post_call(self, increments: Sequence[SpendIncrement]) -> PostCallSettlement | None: + """Declares the increments on the request's post-call batch, None when there is no Redis or request. + Each total lands in memory once the batch goes out; a counter whose increment failed is dropped everywhere. + The settlement reports when every reply is back, whoever sent the batch.""" request: Final = post_call_request() if self.has_redis and increments else None if request is None: - return False + return None return self._handle.increment_many_post_call(request, increments) @with_service_target(SPEND_COUNTERS_TARGET) diff --git a/litellm/redis/stores/spend_counters/handle.py b/litellm/redis/stores/spend_counters/handle.py index e4bd789c341..5ee29b7ad8e 100644 --- a/litellm/redis/stores/spend_counters/handle.py +++ b/litellm/redis/stores/spend_counters/handle.py @@ -3,7 +3,7 @@ from __future__ import annotations from collections.abc import Awaitable, Mapping, Sequence from typing import Protocol -from litellm.caching.redis_batch import NativePending +from litellm.caching.redis_batch import NativePending, PostCallSettlement from litellm.redis.stores.spend_counters.types import SpendIncrement from litellm.rust_bridge._native import NativeRedisClient, NativeRequestRedis, NativeSend, SpendRead @@ -13,6 +13,7 @@ class SpendCountersHandle(Protocol): def has_redis(self) -> bool: ... @property def redis_ttl_seconds(self) -> float | None: ... + def set_committed_ttl(self, ttl_seconds: float, /) -> None: ... def attach(self, client: NativeRedisClient, /) -> bool: ... def read(self, counter: str) -> Awaitable[SpendRead]: ... def read_many( @@ -20,7 +21,9 @@ class SpendCountersHandle(Protocol): ) -> NativePending[Sequence[float | None] | None] | None: ... def increment(self, counter: str, amount: float) -> Awaitable[float]: ... def increment_many(self, increments: Sequence[SpendIncrement]) -> Awaitable[Sequence[float]]: ... - def increment_many_post_call(self, request: NativeRequestRedis, increments: Sequence[SpendIncrement]) -> bool: ... + def increment_many_post_call( + self, request: NativeRequestRedis, increments: Sequence[SpendIncrement] + ) -> PostCallSettlement | None: ... def set_max(self, counter: str, spend: float) -> Awaitable[float | None]: ... def seed(self, counter: str, spend: float, *, increment_if_lost: bool) -> Awaitable[float]: ... def set(self, counter: str, spend: float, ttl_seconds: float | None = None) -> Awaitable[None]: ... diff --git a/litellm/redis/stores/spend_counters/types.py b/litellm/redis/stores/spend_counters/types.py index 7dc846dabaa..9be029b6953 100644 --- a/litellm/redis/stores/spend_counters/types.py +++ b/litellm/redis/stores/spend_counters/types.py @@ -7,6 +7,7 @@ from litellm.rust_bridge._native import NativeRedisStoreError SPEND_COUNTERS_TARGET: Final = "spend_counters" SPEND_COUNTER_PREFIX: Final = "spend:" END_USER_COUNTER_PREFIX: Final = "spend:end_user:" +COMMITTED_COUNTER_PREFIX: Final = "spend:committed:" SpendIncrement: TypeAlias = tuple[str, float] SpendReading: TypeAlias = tuple[float | None, bool] @@ -61,6 +62,11 @@ def shadow_eval_counter(job_id: str) -> str: return f"spend:shadow_eval:{job_id}" +def committed_counter(counter_key: str) -> str: + """The charges-only twin of a key, team, user, org or end user lifetime counter.""" + return f"{COMMITTED_COUNTER_PREFIX}{counter_key.removeprefix(SPEND_COUNTER_PREFIX)}" + + def is_unavailable(error: BaseException) -> bool: """Redis timed out or its breaker is open, the failures a post-call increment drops.""" return isinstance(error, NativeRedisStoreError) and error.args[:1] in (("timeout",), ("breaker_open",)) diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index aa33da54572..fbf9e5d07f0 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -16,6 +16,7 @@ from pydantic import TypeAdapter from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.models.access_group import LiteLLM_AccessGroupTable from litellm.repositories.accounting_delivery_repository import LedgerDatabase, LedgerTransaction from litellm.repositories.prisma_protocols import RowT_co, TableActions @@ -300,6 +301,10 @@ class ModelTableRepository(PrismaTableRepository["prisma_models.LiteLLM_ModelTab class AccessGroupRepository(PrismaTableRepository["prisma_models.LiteLLM_AccessGroupTable"]): table_name = "litellm_accessgrouptable" + async def find_by_name(self, name: str) -> LiteLLM_AccessGroupTable | None: + row: Final = await self.table.find_unique(where={"access_group_name": name}) + return LiteLLM_AccessGroupTable.model_validate(row, from_attributes=True) if row is not None else None + class SSOConfigRepository(PrismaTableRepository["prisma_models.LiteLLM_SSOConfig"]): table_name = "litellm_ssoconfig" diff --git a/litellm/responses/dispatch.py b/litellm/responses/dispatch.py index 0dda3f93660..acf6610d06e 100644 --- a/litellm/responses/dispatch.py +++ b/litellm/responses/dispatch.py @@ -14,6 +14,7 @@ from litellm.rust_bridge.public_call import ( AsyncFlag, NativeCall, bind, + native_acall_hook, native_call, native_call_hook, optional_str, @@ -150,7 +151,7 @@ async def aresponses(*args: object, **kwargs: object) -> ResponsesResult: # kwa _FLAG.dropped(kwargs), python=python, binding=NATIVE_ARESPONSES, - native=native_call_hook, + native=native_acall_hook, ) diff --git a/litellm/router_rust/invoke/endpoints.py b/litellm/router_rust/invoke/endpoints.py index 7a5b52761e2..70c83bed728 100644 --- a/litellm/router_rust/invoke/endpoints.py +++ b/litellm/router_rust/invoke/endpoints.py @@ -2118,7 +2118,10 @@ class RustEndpoints: def text_completion(self, request: Mapping[str, object]) -> object: """The sync `text_completion`: one pick, no retries, the group name as the model.""" - from litellm.litellm_core_utils.routing_common_utils import deployment_model_id_pin + from litellm.litellm_core_utils.routing_common_utils import ( + deployment_model_id_pin, + metadata_defaults_under_request, + ) kwargs: Final = without(request, "is_retry", "is_fallback", "is_async") model: Final = kwargs.get("model") @@ -2139,9 +2142,9 @@ class RustEndpoints: if key not in kwargs: kwargs[key] = value elif key == "metadata": - bucket = kwargs[key] - if isinstance(bucket, dict): - bucket.update(as_dict(value)) # pyright: ignore[reportUnknownMemberType] # the request's metadata dict + metadata_bucket(kwargs, key).update( + metadata_defaults_under_request(defaults=as_dict(value), request_metadata=as_dict(kwargs[key])) + ) pinned: Final = deployment_model_id_pin(data) if pinned is not None: kwargs["model_id"] = pinned diff --git a/litellm/router_rust/invoke/shaping.py b/litellm/router_rust/invoke/shaping.py index 258f919c9f1..bf6410d8242 100644 --- a/litellm/router_rust/invoke/shaping.py +++ b/litellm/router_rust/invoke/shaping.py @@ -32,7 +32,7 @@ from litellm.litellm_core_utils.core_helpers import without from litellm.litellm_core_utils.get_litellm_params import SUPPLEMENTAL_KWARGS_KEYS from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider, inferred_provider from litellm.litellm_core_utils.request_timeout_resolver import get_service_tier_timeout -from litellm.litellm_core_utils.routing_common_utils import deployment_model_id_pin +from litellm.litellm_core_utils.routing_common_utils import deployment_model_id_pin, metadata_defaults_under_request from litellm.litellm_core_utils.tool_fields import TOOL_FIELDS, drop_policy from litellm.responses.retrieve_stream_query import requests_streamed_retrieval, without_retrieve_stream_defaults from litellm.router_api import deployment_id as generate_deployment_id @@ -359,14 +359,15 @@ def update_kwargs_with_default_litellm_params( metadata_variable_name: str, ) -> None: """``_update_kwargs_with_default_litellm_params``: unset non-None defaults fill the request; - the default metadata is merged over the request's.""" + the default metadata fills the keys the request leaves unset, the router's caching groups excepted.""" defaults: Final = dict(without_retrieve_stream_defaults(shaping.default_litellm_params, kwargs)) metadata_defaults: Final = defaults.pop("metadata", {}) or {} # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values for key, value in defaults.items(): if value is None: continue kwargs.setdefault(key, value) - metadata_bucket(kwargs, metadata_variable_name).update(as_dict(metadata_defaults)) # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values + bucket: Final = metadata_bucket(kwargs, metadata_variable_name) + bucket.update(metadata_defaults_under_request(defaults=as_dict(metadata_defaults), request_metadata=bucket)) # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values # --- clientside credentials ------------------------------------------------------------------- diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 600b9d7c455..4fcf555035d 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -801,6 +801,27 @@ def _newest_turn_is_human_ask( return _last_human_ask_index(messages, marker_pairs) == newest_activity +_CLAUDE_CODE_REQUEST_CLASS_HEADER: Final = "x-claude-code-request-class" +_CLAUDE_CODE_AUXILIARY_SENTINEL: Final = f"{_CLAUDE_CODE_REQUEST_CLASS_HEADER}: auxiliary" + + +def _claude_code_request_class(request_kwargs: Mapping[str, object]) -> str | None: + proxy_request: Final = request_kwargs.get("proxy_server_request") + if not isinstance(proxy_request, Mapping): + return None + headers: Final = cast(Mapping[str, object], proxy_request).get("headers") # cast-ok: proxy request is JSON-shaped + if not isinstance(headers, Mapping): + return None + return next( + ( + value.strip().lower() + for key, value in cast(Mapping[object, object], headers).items() # cast-ok: header map from the proxy + if isinstance(key, str) and key.lower() == _CLAUDE_CODE_REQUEST_CLASS_HEADER and isinstance(value, str) + ), + None, + ) + + _CODEX_TURN_METADATA_KEY: Final = "x-codex-turn-metadata" _CODEX_COMPACTION_SUMMARY_PREFIX: Final = ( "Another language model started to solve this problem and produced a summary of its thinking process." @@ -3475,8 +3496,11 @@ class ComplexityRouter(CustomLogger): self._reminder_markers_for_request(request_kwargs), ) - def _matched_housekeeping_sentinel(self, newest_ask: str | None) -> str | None: - """The client housekeeping sentinel on this request's newest ask, or None. + def _matched_housekeeping_sentinel( + self, newest_ask: str | None, request_kwargs: Mapping[str, object] + ) -> str | None: + """The client housekeeping sentinel on this request's newest ask, or None. Claude Code's + `auxiliary` request-class hint counts as one, so it keeps working when the prompts change. Read from the newest ask alone, never the whole history, for the reason `_newest_turn_ask` exists: a title request quoted into a later turn's context would otherwise keep matching and @@ -3490,6 +3514,8 @@ class ComplexityRouter(CustomLogger): """ if self.config.classifier_type == "custom" or not self.config.route_housekeeping_to_cheapest_tier: return None + if _claude_code_request_class(request_kwargs) == "auxiliary": + return _CLAUDE_CODE_AUXILIARY_SENTINEL if not newest_ask: return None return next( @@ -4970,7 +4996,7 @@ class ComplexityRouter(CustomLogger): ), ) - housekeeping_sentinel: Final = self._matched_housekeeping_sentinel(newest_ask) + housekeeping_sentinel: Final = self._matched_housekeeping_sentinel(newest_ask, request_kwargs) housekeeping_tier: Final = self._cheapest_configured_tier() if housekeeping_sentinel is not None else None outcome: Final = ( ClassificationOutcome(tier=housekeeping_tier, score=None, signals=("housekeeping",), cause="housekeeping") diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index d5870bb6510..e42fa6e426b 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -414,9 +414,14 @@ PLAN_MODE_SYSTEM_SENTINELS: Final[tuple[str, ...]] = ('You are currently running # never a paraphrase: a sentinel that matches nothing costs a substring scan per request and reads # as coverage the router does not have. These are client-owned strings that drift with client # releases, so operators extend coverage via housekeeping_patterns rather than editing these. +# The agent progress summary and the next-prompt suggestion were captured from Claude Code 2.1.316 +# on 2026-11-09 (#50444). Both ride the conversation they describe, so only the sentinel tells them +# apart when the client does not send the x-claude-code-request-class hint. HOUSEKEEPING_ASK_SENTINELS: Final[tuple[str, ...]] = ( "Write the title in the predominant language of the session", "You are coming up with a succinct title for a coding session", + "Describe your most recent action in 3-5 words using present tense", + "[SUGGESTION MODE: Suggest what the user might naturally type next into Claude Code.]", ) PLAN_MODE_TOOL_NAME: Final[str] = "exit_plan_mode" @@ -1496,8 +1501,9 @@ class ComplexityRouterConfig(LiteLLMBaseModel): default=None, description=( "Additional case-sensitive literal sentinels that mark a request as client " - "housekeeping, on top of the built-in conversation-title ones. For clients whose " - "wording the built-ins don't cover, or after a client release changes its strings." + "housekeeping, on top of the built-in conversation-title, agent-summary and " + "prompt-suggestion ones. For clients whose wording the built-ins don't cover, or " + "after a client release changes its strings." ), ) diff --git a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py index 9badee854fb..1564c0da634 100644 --- a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py @@ -38,7 +38,7 @@ Safe to enable globally: import hashlib from collections.abc import Callable, Iterator, Mapping, Sequence -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from functools import cache, lru_cache, partial from time import monotonic from typing import TYPE_CHECKING, Final, Optional, cast @@ -62,7 +62,7 @@ from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.router_deprecated import RenamedRouterAttribute, renamed_router_param from litellm.secret_managers.main import secret_manager_would_be_consulted from litellm.types.llms.openai import AllMessageValues -from litellm.types.router import Deployment, GenericLiteLLMParams +from litellm.types.router import Deployment, GenericLiteLLMParams, ModelInfo from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager, get_order_filtered_deployments, with_named_credential @@ -91,8 +91,11 @@ class _EncryptionBoundary: model: Callable[[], str | None] = field(compare=False) account_scope: frozenset[tuple[str, object]] explicit_pair: bool + declared: str | None = None def is_shared_with(self, other: "_EncryptionBoundary") -> bool: + if self.declared is not None and other.declared is not None: + return self.declared == other.declared if (self.api_base, self.api_key) != (other.api_base, other.api_key): return False if self.explicit_pair and other.explicit_pair: @@ -290,6 +293,37 @@ class EncryptedContentAffinityCheck(CustomLogger): ) ) + @staticmethod + def _deployment_encryption_boundary( + litellm_params: object, + model_info: object, + clock: Callable[[], float] = monotonic, + ) -> _EncryptionBoundary | None: + """ + ``_encryption_boundary_key``, overridden by an operator-declared + ``model_info.encrypted_content_boundary``. Two deployments that both declare one are peers exactly + when the names match, which lets endpoints of one provider organization with different ``api_base`` + values share encrypted reasoning. + """ + declared: Final = EncryptedContentAffinityCheck._declared_encryption_boundary(model_info) + resolved: Final = EncryptedContentAffinityCheck._encryption_boundary_key(litellm_params, clock) + if declared is None: + return resolved + if resolved is None: + return _EncryptionBoundary(None, None, lambda: None, frozenset(), explicit_pair=False, declared=declared) + return replace(resolved, declared=declared) + + @staticmethod + def _declared_encryption_boundary(model_info: object) -> str | None: + if isinstance(model_info, ModelInfo): + return model_info.encrypted_content_boundary + if not isinstance(model_info, Mapping): + return None + declared: Final = cast(Mapping[str, object], model_info).get( # cast-ok: narrowed by isinstance + "encrypted_content_boundary" + ) + return declared if isinstance(declared, str) and declared else None + @staticmethod @lru_cache(maxsize=1024) def _encryption_boundary_in_refresh_window(key: _ParamsInRefreshWindow) -> _EncryptionBoundary | None: @@ -408,9 +442,9 @@ class EncryptedContentAffinityCheck(CustomLogger): def _find_deployments_on_same_encryption_boundary( self, - healthy_deployments: list[dict], + healthy_deployments: list[dict[str, object]], model_id: str, - ) -> tuple[list[dict], Deployment | None]: + ) -> tuple[list[dict[str, object]], Deployment | None]: """ Deployments in ``healthy_deployments`` sharing the originating deployment's encryption boundary, alongside the originating @@ -421,15 +455,20 @@ class EncryptedContentAffinityCheck(CustomLogger): originating: Final = self.ctx.registry().deployment(model_id) if originating is None: return [], None - boundary: Final = self._encryption_boundary_key( - originating.litellm_params.model_dump(exclude_none=True), self._clock + boundary: Final = self._deployment_encryption_boundary( + originating.litellm_params.model_dump(exclude_none=True), originating.model_info, self._clock ) if boundary is None: return [], originating matches: Final = [ d for d in healthy_deployments - if (candidate := self._encryption_boundary_key(d.get("litellm_params", {}), self._clock)) is not None + if ( + candidate := self._deployment_encryption_boundary( + d.get("litellm_params", {}), d.get("model_info"), self._clock + ) + ) + is not None and boundary.is_shared_with(candidate) ] return matches, originating @@ -464,8 +503,8 @@ class EncryptedContentAffinityCheck(CustomLogger): boundary for target in target_deployments if ( - boundary := EncryptedContentAffinityCheck._encryption_boundary_key( - target.get("litellm_params"), clock + boundary := EncryptedContentAffinityCheck._deployment_encryption_boundary( + target.get("litellm_params"), target.get("model_info"), clock ) ) is not None @@ -482,8 +521,8 @@ class EncryptedContentAffinityCheck(CustomLogger): return False origin: Final = ctx.registry().deployment(origin_model_id) origin_boundary: Final = ( - EncryptedContentAffinityCheck._encryption_boundary_key( - origin.litellm_params.model_dump(exclude_none=True), clock + EncryptedContentAffinityCheck._deployment_encryption_boundary( + origin.litellm_params.model_dump(exclude_none=True), origin.model_info, clock ) if origin is not None else None diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 77af87f5eb3..114d2df026e 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -358,6 +358,11 @@ class SpendRead: @property def from_redis(self) -> bool: ... +@final +class PostCallSettlement: + @property + def settled(self) -> bool: ... + @final class NativeSpendCounters: def __new__(cls, client: NativeRedisClient | None) -> NativeSpendCounters: ... @@ -366,6 +371,7 @@ class NativeSpendCounters: def has_redis(self) -> bool: ... @property def redis_ttl_seconds(self) -> float | None: ... + def set_committed_ttl(self, ttl_seconds: float) -> None: ... def read(self, counter: str) -> Future[SpendRead]: ... def read_many( self, counters: Sequence[str], send: NativeSend @@ -374,7 +380,7 @@ class NativeSpendCounters: def increment_many(self, increments: Sequence[tuple[str, float]]) -> Future[list[float]]: ... def increment_many_post_call( self, request: NativeRequestRedis, increments: Sequence[tuple[str, float]] - ) -> bool: ... + ) -> PostCallSettlement | None: ... def set_max(self, counter: str, spend: float) -> Future[float | None]: ... def seed(self, counter: str, spend: float, *, increment_if_lost: bool) -> Future[float]: ... def set(self, counter: str, spend: float, ttl_seconds: float | None = None) -> Future[None]: ... @@ -596,6 +602,7 @@ __all__ = [ "NativeSend", "NativeSensitiveRoutePins", "NativeSpendCounters", + "PostCallSettlement", "ProcessReservedForForking", "ResponsesWebSocketConnection", "RouterCall", diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index 78040ba7225..c493617003d 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -39,7 +39,7 @@ class Feature(str, Enum): PROXY_MCP_TOOLS = "proxy_mcp_tools" """litellm_proxy MCP tools that Python expands and may auto-execute.""" CALLER_CLIENT = "caller_client" - """A caller-supplied HTTP client or shared session that Python sends through.""" + """A caller-supplied HTTP client that Python sends through.""" TIMEOUT_VARIANTS = "timeout_variants" """request_timeout, stream_timeout or httpx.Timeout phases beyond a plain timeout.""" DROP_PARAMS = "drop_params" @@ -172,7 +172,7 @@ CHAT_PYTHON_FEATURES: Final = COMMON_PYTHON_FEATURES | { Feature.MODEL_ALIAS_MAP, } -MESSAGES_PYTHON_FEATURES: Final = COMMON_PYTHON_FEATURES | { +MESSAGES_PYTHON_FEATURES: Final = (COMMON_PYTHON_FEATURES - {Feature.TIMEOUT_VARIANTS}) | { Feature.PROVIDER_AFFINITY, Feature.FOUNDRY_BETA_ROW, Feature.SAFEGUARDS_CLASSIFIER, diff --git a/litellm/rust_bridge/detectors.py b/litellm/rust_bridge/detectors.py index f2591bdbcd1..c38c3ea7135 100644 --- a/litellm/rust_bridge/detectors.py +++ b/litellm/rust_bridge/detectors.py @@ -44,8 +44,8 @@ def latest_alias(request: NativeCall) -> bool: def caller_client(request: NativeCall) -> bool: - """Python sends through the caller's HTTP client or shared session.""" - return present(request, "client") or present(request, "shared_session") + """Python sends through the caller's HTTP client; a shared session only reuses connections, which native pools.""" + return present(request, "client") def extra_body(request: NativeCall) -> bool: diff --git a/litellm/rust_bridge/host/cache.py b/litellm/rust_bridge/host/cache.py index cac4b75677f..bf2655f9f62 100644 --- a/litellm/rust_bridge/host/cache.py +++ b/litellm/rust_bridge/host/cache.py @@ -1,12 +1,54 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from collections.abc import Mapping +from typing import TYPE_CHECKING, Final + +from pydantic import TypeAdapter + +from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict if TYPE_CHECKING: from litellm.caching.caching import Cache +_RUST_RESPONSE_HEADER: Final = "x-litellm-rust" +_HEADERS: Final = TypeAdapter(dict[str, object]) + def configured_cache() -> Cache | None: import litellm return litellm.cache + + +def _python_entry(response: Mapping[str, object]) -> Mapping[str, object]: + hidden: Final = get_hidden_params_dict(response) + headers: Final = hidden.get("additional_headers") + if not isinstance(headers, Mapping) or _RUST_RESPONSE_HEADER not in headers: + return response + return { + **response, + "_hidden_params": { + **hidden, + "additional_headers": { + name: value + for name, value in _HEADERS.validate_python(headers).items() + if name != _RUST_RESPONSE_HEADER + }, + }, + } + + +def store_response( + cache: Cache, + response: Mapping[str, object], + /, + **kwargs: object, # kwargs-ok: the cache call's own arguments +) -> None: + from litellm.caching.caching_handler import create_cache_write_task + + create_cache_write_task( + lambda: cache.async_add_cache( # pyright: ignore[reportUnknownMemberType] # Cache.async_add_cache leaves its result untyped + _python_entry(response), + **kwargs, # pyright: ignore[reportArgumentType] # the cache call's own keyword arguments, as Python passes them + ) + ) diff --git a/litellm/rust_bridge/host/callbacks_legacy_python.py b/litellm/rust_bridge/host/callbacks_legacy_python.py index 92590531cdc..ba557f769eb 100644 --- a/litellm/rust_bridge/host/callbacks_legacy_python.py +++ b/litellm/rust_bridge/host/callbacks_legacy_python.py @@ -9,10 +9,13 @@ from __future__ import annotations import asyncio import contextvars import datetime +import inspect +import json +import sys import traceback import uuid import weakref -from collections.abc import Awaitable, Coroutine, Mapping, Sequence +from collections.abc import Awaitable, Callable, Coroutine, Mapping, Sequence from dataclasses import dataclass from typing import ( TYPE_CHECKING, @@ -111,6 +114,9 @@ def finalize( class LoggingSurface(Protocol): + deferred_success_result: tuple[object, datetime.datetime, datetime.datetime] | None + enqueue_deferred_logging: Callable[[], None] | None + @property def model_call_details(self) -> Mapping[str, object]: ... @@ -139,6 +145,10 @@ class LoggingSurface(Protocol): additional_args: dict[str, object] = ..., ) -> object: ... + def claim_async_success_log(self) -> bool: ... + + def mark_async_success_log_enqueued(self) -> None: ... + def handle_sync_success_callbacks_for_async_calls( self, result: object, start_time: datetime.datetime, end_time: datetime.datetime, cache_hit: bool | None = None ) -> None: ... @@ -234,10 +244,33 @@ def update_logging( ) -def pre_call(logger: LoggingSurface, input: str, api_key: str | None, additional_args: dict[str, object]) -> None: +def pre_call( + logger: LoggingSurface, + input: str | Sequence[Mapping[str, str]], + api_key: str | None, + additional_args: dict[str, object], +) -> None: logger.pre_call(input=input, api_key=api_key, additional_args=additional_args) +def messages_input(body: Mapping[str, object], kwargs: Mapping[str, object]) -> Sequence[Mapping[str, str]]: + from litellm.messages import anthropic_messages_handler + + # Python's Messages handler builds its body as messages, max_tokens, model, then its own parameters in + # signature order and the caller's other keywords in their order, and logs that body as serialized + order: Final = ( + "messages", + "max_tokens", + "model", + *inspect.signature(anthropic_messages_handler).parameters, + *kwargs, + *body, + ) + return [ + {"role": "user", "content": json.dumps({name: body[name] for name in dict.fromkeys(order) if name in body})} + ] + + def post_call( logger: LoggingSurface, original_response: str, api_key: str | None, additional_args: dict[str, object] ) -> None: @@ -269,8 +302,23 @@ def defers_async_logging(logger: LoggingSurface) -> bool: return bool(getattr(logger, "defer_async_logging", False)) -def defer_success(logger: LoggingSurface, pending: object) -> None: - setattr(logger, "_native_pending_logging", pending) +def claim_async_success_log(logger: LoggingSurface) -> bool: + return logger.claim_async_success_log() + + +def mark_async_success_enqueued(logger: LoggingSurface) -> None: + logger.mark_async_success_log_enqueued() + + +def defer_success( + logger: LoggingSurface, + pending: Callable[[], None], + response: object, + start: datetime.datetime, + end: datetime.datetime, +) -> None: + logger.deferred_success_result = (response, start, end) # rebind-ok: capture the native attempt before dispatch + logger.enqueue_deferred_logging = pending # rebind-ok: the proxy releases logging here def _cache_hit(logger: LoggingSurface) -> Literal[True] | None: @@ -410,6 +458,14 @@ def stream_opened(logger: Logging, head: Mapping[str, object]) -> None: ) +def _cancelling(task: asyncio.Task[object]) -> bool: + return task.cancelling() > 0 if sys.version_info >= (3, 11) else True + + +def _proxy_request(logger: LoggingSurface) -> object: + return logger.litellm_params.get("proxy_server_request") + + def stream_success( logger: Logging, url_route: str, @@ -419,6 +475,7 @@ def stream_success( start: datetime.datetime, end: datetime.datetime, first_chunk: datetime.datetime | None, + stream_teardown: bool, ) -> None: from litellm.llms.anthropic.messages.native.streaming_iterator import ( GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ, @@ -443,7 +500,7 @@ def stream_success( raw_bytes=chunks, end_time=end, ) - if getattr(logger, "_on_deferred_stream_complete", None) is not None: + if not stream_teardown and getattr(logger, "_on_deferred_stream_complete", None) is not None: logger._deferred_stream_complete_args = (coroutine,) # pyright: ignore[reportAttributeAccessIssue] # the proxy's deferred stream release reads this slot return try: @@ -453,6 +510,11 @@ def stream_success( executor.submit(contextvars.copy_context().run, asyncio.run, coroutine) return + cancelled: Final = asyncio.current_task() + if stream_teardown and cancelled is not None and _cancelling(cancelled) and _proxy_request(logger): + # the proxy's disconnect cleanup runs later in this task and stamps the 499 error_information this row carries + cancelled.add_done_callback(lambda _: enqueue_logging(coroutine)) + return enqueue_logging(coroutine) @@ -463,9 +525,15 @@ def stream_failure( chunks: list[bytes], error: Exception, ) -> Coroutine[object, object, None]: + from litellm._logging import verbose_proxy_logger from litellm.proxy.pass_through_endpoints.streaming_handler import PassThroughStreamingHandler from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType + verbose_proxy_logger.error("Error in chunk_processor: %s", error) + # Python's stream fails before its success handler records `cache_hit`; only a cache hit + # stamps it earlier. + if logger.model_call_details.get("cache_hit") is not True: + logger.model_call_details.pop("cache_hit", None) return PassThroughStreamingHandler.schedule_stream_failure_logging( litellm_logging_obj=logger, endpoint_type=EndpointType(endpoint_type), diff --git a/litellm/rust_bridge/host/routes/messages.py b/litellm/rust_bridge/host/routes/messages.py index c8b53097f46..c6568264edb 100644 --- a/litellm/rust_bridge/host/routes/messages.py +++ b/litellm/rust_bridge/host/routes/messages.py @@ -1,11 +1,16 @@ from __future__ import annotations import datetime +import errno +import ipaddress +import os +import socket from collections.abc import Callable, Mapping, Sequence from dataclasses import asdict, dataclass from typing import Final, TypeAlias, cast # noqa: TID251 # narrows native payloads to public TypedDicts import httpx +from aiohttp.http_exceptions import ContentLengthError, TransferEncodingError from pydantic import TypeAdapter, ValidationError import litellm @@ -14,7 +19,7 @@ from litellm.litellm_core_utils.core_helpers import normalize_drop_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.reasoning_effort_capability import deployment_reasoning_effort_levels from litellm.litellm_core_utils.request_timeout_resolver import get_configured_request_timeout -from litellm.llms.anthropic.common_utils import AnthropicError +from litellm.llms.anthropic.common_utils import AnthropicError, anthropic_pass_through_api_base from litellm.llms.anthropic.messages.native import handler from litellm.llms.anthropic.messages.native.handler import ( _native_messages_config as native_messages_config, # pyright: ignore[reportPrivateUsage] # the handler's own config lookup @@ -23,12 +28,13 @@ from litellm.llms.anthropic.messages.native.handler import ( _with_deployment_output_effort as _deployment_output_effort, # pyright: ignore[reportPrivateUsage] # the handler's own effort rounding ) from litellm.llms.anthropic.messages.utils import is_reasoning_auto_summary_enabled +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.custom_httpx.http_handler import get_ssl_configuration from litellm.rust_bridge import failures -from litellm.rust_bridge.host import callbacks_legacy_python +from litellm.rust_bridge.host import callbacks_legacy_python, http_transport from litellm.rust_bridge.public_call import known_params, optional_str from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse -_DROP_PATHS: Final = TypeAdapter(list[object]) _METADATA: Final = TypeAdapter(dict[str, object]) _BODY_FIELDS: Final = frozenset( { @@ -85,7 +91,7 @@ def inferred_provider(model: str, api_base: str | None) -> str | None: def body_params(request: Mapping[str, object]) -> Mapping[str, object]: - params: Final = known_params(_BODY_FIELDS, request) + params: Final = known_params(_BODY_FIELDS.difference(handler.additional_drop_paths(request)), request) if deployment_reasoning_effort_levels(request.get("model_info")) is None: return params return _with_deployment_effort(params, request) @@ -120,6 +126,7 @@ def stream_hidden_params(headers: Sequence[tuple[str, str]]) -> Mapping[str, obj _UPSTREAM_ARGS: Final = TypeAdapter(tuple[int, str]) _UPSTREAM_HEADERS: Final = TypeAdapter(list[tuple[str, str]]) +_BODY_CUT: Final = TypeAdapter(tuple[int, int | None, bool]) _ERROR_CLASS: TypeAlias = Callable[..., Exception] _validate_metadata: Final = cast( # cast-ok: the handler's untyped metadata validator Callable[[dict[str, object]], object], handler.validate_anthropic_api_metadata @@ -128,6 +135,17 @@ _exception_type: Final = cast( # cast-ok: the legacy public exception mapper is Callable[..., Exception], exception_mapping_utils.exception_type ) _BODY_READ_TIMEOUT: Final = "Timeout on reading data from socket" # aiohttp's SocketTimeoutError text +_UNREACHABLE: Final = "could not reach the provider: " +_SEND_FAILED: Final = "upstream network error: error sending request: client error (SendRequest): " +_DNS_FAILURE: Final = "failed to lookup address information: " +_HTTPX_DNS_ERRNO: Final = { + "Name or service not known": socket.EAI_NONAME, + "Temporary failure in name resolution": socket.EAI_AGAIN, +} +_TRANSPORT_TIMEOUTS: Final[Mapping[object, type[httpx.TimeoutException]]] = { + "Connection": httpx.ConnectTimeout, + "Read": httpx.ReadTimeout, +} _MISSING_KEY: Final = ( "Missing Anthropic API Key - A call is being made to anthropic but no key is set " "either in the environment variables or via params. Please set `ANTHROPIC_API_KEY` " @@ -174,6 +192,82 @@ def _timeout_text(phase: str, request: Mapping[str, object]) -> str: ) +def _body_cut(received: int, expected: int | None, reset: bool) -> httpx.TransportError: + if not http_transport.uses_aiohttp_transport(): + if reset: + return httpx.ReadError("") + detail: Final = ( + "incomplete chunked read" if expected is None else f"received {received} bytes, expected {expected}" + ) + return httpx.RemoteProtocolError(f"peer closed connection without sending complete message body ({detail})") + payload: Final = ( + TransferEncodingError("Not enough data to satisfy transfer length header.") + if expected is None + else ContentLengthError( + f"Not enough data to satisfy content length header (received {received} of {expected} bytes)." + ) + ) + text: Final = f"Response payload is not completed: {payload!r}" + if not reset: + return httpx.ReadError(text) + return httpx.ReadError(f"{text}. {ConnectionResetError(errno.ECONNRESET, os.strerror(errno.ECONNRESET))!r}") + + +def _connect_failure(native: str, request: Mapping[str, object], request_provider: str) -> str | None: + api_base: Final = request.get("api_base") + if not isinstance(api_base, str) and request_provider != "anthropic": + return None + refused: Final = "tcp connect error: Connection refused" in native + dns_reason: Final = native.rsplit(_DNS_FAILURE, 1)[1] if _DNS_FAILURE in native else None + tls_eof: Final = native.endswith("tls handshake eof") + if not http_transport.uses_aiohttp_transport(): + if refused: + return "All connection attempts failed" + dns_errno: Final = None if dns_reason is None else _HTTPX_DNS_ERRNO.get(dns_reason) + if dns_errno is not None: + return f"[Errno {dns_errno}] {dns_reason}" + return "" if tls_eof else None + url: Final = httpx.URL(api_base if isinstance(api_base, str) else anthropic_pass_through_api_base()) + port: Final = url.port or (443 if url.scheme == "https" else 80) + if refused: + try: + address: Final = ipaddress.ip_address(url.host) + except ValueError: + return None + peer: Final = (str(address), port) if address.version == 4 else (str(address), port, 0, 0) + return _aiohttp_connect_failure(url, port, request, f"Connect call failed {peer}") + if dns_reason is not None: + return _aiohttp_connect_failure(url, port, request, dns_reason) + if tls_eof: + return _aiohttp_connect_failure(url, port, request, None) + return None + + +def _aiohttp_connect_failure(url: httpx.URL, port: int, request: Mapping[str, object], reason: str | None) -> str: + ssl_verify: Final = request.get("ssl_verify") + ssl_config: Final = get_ssl_configuration(ssl_verify if isinstance(ssl_verify, bool | str) else None) + return f"Cannot connect to host {url.host}:{port} ssl:{'default' if ssl_config is True else ssl_config} [{reason}]" + + +def _transport_failure(error: Exception, request: Mapping[str, object], request_provider: str) -> str | None: + """The text Python's HTTP client raises for a call that failed before any response arrived.""" + native: Final = str(error) + if native.startswith(_SEND_FAILED): + reset: Final = native.endswith(f"Connection reset by peer (os error {errno.ECONNRESET})") + closed: Final = native.endswith("connection closed before message completed") + if http_transport.uses_aiohttp_transport(): + if reset: + return str(ConnectionResetError(errno.ECONNRESET, os.strerror(errno.ECONNRESET))) + return "Server disconnected" if closed else native + if reset: + return "" + return "Server disconnected without sending a response." if closed else native + if not native.startswith(_UNREACHABLE): + return None + connect: Final = _connect_failure(native, request, request_provider) + return native if connect is None else connect + + def _upstream_failure( error: Exception, request: Mapping[str, object] ) -> tuple[int, str, Sequence[tuple[str, str]]] | None: @@ -183,14 +277,27 @@ def _upstream_failure( return 408, _timeout_text(timeout_phase, request), () try: status, text = _UPSTREAM_ARGS.validate_python(error.args) - return status, text, _UPSTREAM_HEADERS.validate_python(getattr(error, "headers", None)) + headers: Final = _UPSTREAM_HEADERS.validate_python(getattr(error, "headers", None)) except ValidationError: return None + if not headers: + return status, text, headers + # Python's `MaskedHTTPStatusError` rebuilds the response from its decoded body without the + # wire's content-encoding and content-length, so httpx states a length only for a non-empty body. + rebuilt: Final = httpx.Response( + status, + content=text.encode(), + headers=[ + (name, value) for name, value in headers if name.lower() not in ("content-encoding", "content-length") + ], + ) + return status, text, tuple(rebuilt.headers.items()) def _provider_error(error: Exception, request: Mapping[str, object], request_provider: str) -> Exception | None: """The provider exception Python's `_handle_error` builds for the same failure.""" - failure: Final = _upstream_failure(error, request) + transport: Final = _transport_failure(error, request, request_provider) + failure: Final = _upstream_failure(error, request) if transport is None else (500, transport, ()) if failure is None: return None status, text, headers = failure @@ -201,13 +308,17 @@ def _provider_error(error: Exception, request: Mapping[str, object], request_pro request.get("model_info"), api_base if isinstance(api_base, str) else None, ) - return ( + provider_error: Final = ( cast(_ERROR_CLASS, config.get_error_class)( # cast-ok: the config's untyped headers parameter error_message=text, status_code=status, headers=dict(headers) ) if config is not None else AnthropicError(status_code=status, message=text, headers=httpx.Headers(headers)) ) + if transport is not None and isinstance(provider_error, BaseLLMException): + # Python's `_handle_error` gives a failure without a response the status 500. + provider_error.status_code_is_synthesized = True + return provider_error def _log_resolved_call(request: Mapping[str, object], request_provider: str) -> None: @@ -232,6 +343,17 @@ def map_failure(error: Exception, request: Mapping[str, object], request_provide # Python reads a stream's body outside its HTTP client's error mapping, so the transport's # own timeout reaches the caller unmapped. return httpx.ReadTimeout(_BODY_READ_TIMEOUT) + body_cut: Final = getattr(error, "messages_body_cut", None) + if body_cut is not None: + # Python reads a stream's body outside its HTTP client's error mapping, so a body the + # connection cut short reaches the caller as its transport raised it. + return _body_cut(*_BODY_CUT.validate_python(body_cut)) + transport_timeout: Final = _TRANSPORT_TIMEOUTS.get(getattr(error, "messages_timeout", None)) + if transport_timeout is not None: + # Python's `litellm.Timeout` is raised while handling httpx's own timeout, which the proxy's + # failure spend reads to bill a read timeout's prompt. + chained: Final = transport_timeout(str(error)) + error.__context__ = chained # rebind-ok: Python chains the transport timeout as the exception context if getattr(error, "messages_late_rejection", False): _log_resolved_call(request, request_provider) if getattr(error, "messages_request_error", False): @@ -263,20 +385,12 @@ def _drop_params(kwargs: Mapping[str, object]) -> bool: return bool(litellm.drop_params) or normalize_drop_params(kwargs.get("drop_params")) is True -def _additional_drop_params(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 settings(kwargs: Mapping[str, object]) -> dict[str, object]: return asdict( MessagesSettings( drop_params=_drop_params(kwargs), reasoning_auto_summary=is_reasoning_auto_summary_enabled(), - additional_drop_params=_additional_drop_params(kwargs), + additional_drop_params=handler.additional_drop_paths(kwargs), request_timeout=get_configured_request_timeout(), ) ) diff --git a/litellm/rust_bridge/public_call.py b/litellm/rust_bridge/public_call.py index 8625ca71bb2..1ab3d5e10c9 100644 --- a/litellm/rust_bridge/public_call.py +++ b/litellm/rust_bridge/public_call.py @@ -2,12 +2,17 @@ from __future__ import annotations +import asyncio import inspect -from collections.abc import Callable, Collection, Mapping, Sequence +from collections.abc import Awaitable, Callable, Collection, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType from typing import Final, TypeVar, cast # noqa: TID251 # narrows caller-owned containers without copying them +from litellm.utils import ( + _get_wrapper_timeout, # pyright: ignore[reportPrivateUsage] # the @client wrapper's own failure stamp +) + def signature(legacy: Callable[..., object]) -> inspect.Signature: return inspect.signature(legacy) @@ -115,3 +120,32 @@ def native_call_hook( _kwargs: Mapping[str, object], ) -> NativeResultT: return hook(call) + + +async def native_acall_hook( + hook: Callable[[NativeCall], Awaitable[NativeResultT]], + call: NativeCall, + args: tuple[object, ...], + kwargs: Mapping[str, object], +) -> NativeResultT: + # Python's async entries leave the event loop (run_in_executor) before they build the provider body, so + # tasks scheduled next to the call, such as the proxy's during_call guardrails, edit the request first + await asyncio.sleep(0) + return await native_acall_hook_on_loop(hook, call, args, kwargs) + + +async def native_acall_hook_on_loop( + hook: Callable[[NativeCall], Awaitable[NativeResultT]], + call: NativeCall, + _args: tuple[object, ...], + kwargs: Mapping[str, object], +) -> NativeResultT: + """``native_acall_hook`` for an async entry whose Python path built its request on the loop, as OCR's did.""" + try: + return await hook(call) + except Exception as error: + # utils.client stamps every async failure, which the proxy reads for x-litellm-timeout and the router for retries + if kwargs.get("num_retries") is not None: + setattr(error, "num_retries", kwargs.get("num_retries")) # noqa: B010 # Exception declares no such attribute, as in utils.client + setattr(error, "timeout", _get_wrapper_timeout(kwargs=dict(kwargs))) # noqa: B010 # Exception declares no such attribute, as in utils.client + raise diff --git a/litellm/types/proxy/auth/user_api_key_auth.py b/litellm/types/proxy/auth/user_api_key_auth.py index 1183c36eda7..8975a4057eb 100644 --- a/litellm/types/proxy/auth/user_api_key_auth.py +++ b/litellm/types/proxy/auth/user_api_key_auth.py @@ -16,6 +16,7 @@ from litellm.types.agents import AgentCaller, AgentResponse from litellm.types.proxy.agent_identity import ManagedAgentContext from litellm.types.proxy.auth.user_roles import LitellmUserRoles from litellm.types.proxy.carried_budget_state import ( + CommittedBudgetSpend, OrgBudgetSnapshot, TeamBudgetSnapshot, UserBudgetSnapshot, @@ -132,6 +133,17 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob """ api_key: str | None = None + _access_group_rule_grants: frozenset[str] = PrivateAttr(default=frozenset()) + + @property + def access_group_rule_grants(self) -> frozenset[str]: + return self._access_group_rule_grants + + def apply_access_group_rule_grants(self, grants: tuple[str, ...]) -> None: + if grants: + self._access_group_rule_grants = frozenset(grants) + self.access_group_ids = list(dict.fromkeys((*(self.access_group_ids or ()), *grants))) + user_role: LitellmUserRoles | None = None allowed_model_region: AllowedModelRegion | None = None parent_otel_span: Span | None = None @@ -206,6 +218,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob team_budget_snapshot: TeamBudgetSnapshot | None = Field(default=None, exclude=True) user_budget_snapshot: UserBudgetSnapshot | None = Field(default=None, exclude=True) org_budget_snapshot: OrgBudgetSnapshot | None = Field(default=None, exclude=True) + committed_budget_spend: CommittedBudgetSpend | None = Field(default=None, exclude=True, repr=False) matched_model_access_groups: list[str] | None = Field(default=None, exclude=True) budget_throttle_pct: float | None = Field(default=None, exclude=True) user: object | None = None # Expanded user object when expand=user is used @@ -329,6 +342,21 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob ) +COMMITTED_BUDGET_SPEND_PARAM: Final = "committed_budget_spend" + + +def committed_budget_spend_of_row( + metadata: Mapping[str, object], litellm_params: Mapping[str, object] +) -> CommittedBudgetSpend | None: + auth: Final = metadata.get("user_api_key_auth") + if isinstance(auth, UserAPIKeyAuth) and auth.committed_budget_spend is not None: + return auth.committed_budget_spend + if "user_api_key_auth" in metadata: + return CommittedBudgetSpend() + carried: Final = litellm_params.get(COMMITTED_BUDGET_SPEND_PARAM) + return carried if isinstance(carried, CommittedBudgetSpend) else None + + def user_api_key_has_admin_view(user_api_key_dict: UserAPIKeyAuth) -> bool: """Return True if the caller's role grants unscoped read access to all tenant resources (managed files, batches, vector stores, spend rows, etc). diff --git a/litellm/types/proxy/carried_budget_state.py b/litellm/types/proxy/carried_budget_state.py index 1a0c55a38ff..83ef6a21188 100644 --- a/litellm/types/proxy/carried_budget_state.py +++ b/litellm/types/proxy/carried_budget_state.py @@ -9,10 +9,18 @@ reads them back with ``from_metadata``. ``None`` means this request never carrie from collections.abc import Mapping from datetime import datetime from types import MappingProxyType +from typing import Final from pydantic import ConfigDict, Field, ValidationError from typing_extensions import Self +from litellm.redis.stores.spend_counters.types import ( + end_user_counter, + key_counter, + org_counter, + team_counter, + user_counter, +) from litellm.types.llms.base import LiteLLMBaseModel @@ -25,7 +33,7 @@ class _BudgetSnapshot(LiteLLMBaseModel): @classmethod def from_metadata(cls, metadata: Mapping[str, object]) -> Self | None: try: - return cls.model_validate(metadata) + return cls.model_validate(metadata, by_name=False) except ValidationError: return None @@ -64,3 +72,51 @@ class OrgBudgetSnapshot(_BudgetSnapshot): max_budget: float | None = Field( validation_alias="user_api_key_org_max_budget", serialization_alias="user_api_key_org_max_budget" ) + + +class EntitySpendCounterKeys(LiteLLMBaseModel): + """Spend counter of each entity whose remaining budget a gauge reports. ``None`` for an entity with no budget.""" + + model_config = ConfigDict(frozen=True) + + key: str | None = None + team: str | None = None + user: str | None = None + org: str | None = None + end_user: str | None = None + + @classmethod + def of( + cls, + token: str | None, + team_id: str | None, + user_id: str | None, + org_id: str | None, + end_user_id: str | None, + ) -> Self: + return cls( + key=key_counter(token) if token else None, + team=team_counter(team_id) if team_id else None, + user=user_counter(user_id) if user_id else None, + org=org_counter(org_id) if org_id else None, + end_user=end_user_counter(end_user_id) if end_user_id else None, + ) + + def tracked(self) -> tuple[str, ...]: + counter_keys: Final = (self.key, self.team, self.user, self.org, self.end_user) + return tuple(counter_key for counter_key in counter_keys if counter_key is not None) + + +class CommittedBudgetSpend(LiteLLMBaseModel): + """Spend each entity was charged, budget holds excluded, read from the committed spend counters for the + remaining budget gauges only. ``None`` when the entity has no committed counter. ``counter_keys`` names the + budgeted entities whose committed counters this request charges and reads.""" + + model_config = ConfigDict(frozen=True) + + counter_keys: EntitySpendCounterKeys = Field(default=EntitySpendCounterKeys(), repr=False) + key: float | None = Field(default=None, repr=False) + team: float | None = Field(default=None, repr=False) + user: float | None = Field(default=None, repr=False) + org: float | None = Field(default=None, repr=False) + end_user: float | None = Field(default=None, repr=False) diff --git a/litellm/types/router.py b/litellm/types/router.py index c0b182f0f20..6f5652db771 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -353,6 +353,10 @@ class ModelInfo(MirroredPricingParams): max_embedding_batch_size: Annotated[int, Field(gt=0)] | None = None + # deployments that set the same value share one provider encryption boundary, so encrypted reasoning + # minted on one can be replayed to another (e.g. regional endpoints of one OpenAI organization) + encrypted_content_boundary: Annotated[str, Field(min_length=1)] | None = None + def __init__(self, id: str | int | None = None, **params) -> None: if id is None: id = str(uuid.uuid4()) # Generate a UUID if id is None or not provided diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 78f5f8c29b7..68bbf612e31 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -1970,7 +1970,7 @@ "source": "https://aws.amazon.com/bedrock/pricing/" }, "bedrock/claude-fable-latest": { - "alias_of": "global.anthropic.claude-fable-5-1" + "alias_of": "global.anthropic.claude-fable-5-2" }, "us.anthropic.claude-fable-5": { "cache_creation_input_token_cost": 1.375e-05, @@ -5343,7 +5343,7 @@ "supports_tool_choice": true }, "azure/eu/gpt-4o-realtime-preview-2024-12-17": { - "cache_read_input_audio_token_cost": 2.5e-06, + "cache_read_input_audio_token_cost": 2.75e-06, "cache_read_input_token_cost": 2.75e-06, "input_cost_per_audio_token": 4.4e-05, "input_cost_per_token": 5.5e-06, @@ -5353,7 +5353,7 @@ "mode": "realtime", "output_cost_per_audio_token": 8.8e-05, "output_cost_per_token": 2.2e-05, - "source": "https://azure.microsoft.com/en-us/pricing/details/azure-openai/", + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_modalities": [ "text", "audio" @@ -6691,6 +6691,7 @@ "supports_tool_choice": true }, "azure/gpt-realtime": { + "cache_read_input_image_token_cost": 5e-07, "cache_creation_input_audio_token_cost": 4e-06, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, @@ -6726,9 +6727,11 @@ "supports_parallel_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure" + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/" }, "azure/gpt-realtime-2025-08-28": { + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "cache_read_input_image_token_cost": 5e-07, "cache_creation_input_audio_token_cost": 4e-06, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, @@ -6766,6 +6769,7 @@ "supports_tool_choice": true }, "azure/gpt-realtime-1.5": { + "cache_read_input_image_token_cost": 5e-07, "cache_creation_input_audio_token_cost": 4e-06, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, @@ -6802,9 +6806,11 @@ "supports_parallel_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure" + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/" }, "azure/gpt-realtime-1.5-2026-02-23": { + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "cache_read_input_image_token_cost": 5e-07, "cache_creation_input_audio_token_cost": 4e-06, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, @@ -6925,6 +6931,8 @@ "supports_tool_choice": true }, "azure/gpt-realtime-mini": { + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "cache_read_input_image_token_cost": 8e-08, "cache_creation_input_audio_token_cost": 3e-07, "cache_read_input_audio_token_cost": 3e-07, "cache_read_input_token_cost": 6e-08, @@ -7033,6 +7041,7 @@ ] }, "azure/gpt-4o-realtime-preview-2024-12-17": { + "cache_read_input_audio_token_cost": 2.5e-06, "cache_read_input_token_cost": 2.5e-06, "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 5e-06, @@ -7042,7 +7051,7 @@ "mode": "realtime", "output_cost_per_audio_token": 8e-05, "output_cost_per_token": 2e-05, - "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure", + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_modalities": [ "text", "audio" @@ -12430,7 +12439,7 @@ "supports_tool_choice": true }, "azure/us/gpt-4o-realtime-preview-2024-12-17": { - "cache_read_input_audio_token_cost": 2.5e-06, + "cache_read_input_audio_token_cost": 2.75e-06, "cache_read_input_token_cost": 2.75e-06, "input_cost_per_audio_token": 4.4e-05, "input_cost_per_token": 5.5e-06, @@ -12440,7 +12449,7 @@ "mode": "realtime", "output_cost_per_audio_token": 8.8e-05, "output_cost_per_token": 2.2e-05, - "source": "https://azure.microsoft.com/en-us/pricing/details/azure-openai/", + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_modalities": [ "text", "audio" @@ -17082,7 +17091,7 @@ "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, "claude-fable-latest": { - "alias_of": "claude-fable-5-1" + "alias_of": "claude-fable-5-2" }, "claude-opus-5-5": { "supports_web_search": true, @@ -29363,52 +29372,6 @@ "comment": "Open flagship GLM for long-horizon coding agents and million-token context work", "source": "https://api.friendli.ai/serverless/v1/models" }, - "friendliai/deepseek-ai/DeepSeek-V3.2": { - "litellm_provider": "friendliai", - "max_input_tokens": 163840, - "max_output_tokens": 163840, - "input_cost_per_token": 5e-07, - "output_cost_per_token": 1.5e-06, - "cache_read_input_token_cost": 2.5e-07, - "supports_prompt_caching": true, - "supports_reasoning": true, - "reasoning_effort_levels": [], - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_native_structured_output": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": false, - "supports_image_input": false, - "supports_video_input": false, - "mode": "chat", - "comment": "DeepSeek chat model for instruction following, coding, and analysis", - "source": "https://api.friendli.ai/serverless/v1/models" - }, - "friendliai/MiniMaxAI/MiniMax-M2.5": { - "litellm_provider": "friendliai", - "max_input_tokens": 196608, - "max_output_tokens": 196608, - "input_cost_per_token": 3e-07, - "output_cost_per_token": 1.2e-06, - "cache_read_input_token_cost": 6e-08, - "supports_prompt_caching": true, - "supports_reasoning": true, - "reasoning_effort_levels": [], - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_native_structured_output": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": false, - "supports_image_input": false, - "supports_video_input": false, - "mode": "chat", - "comment": "Prior MiniMax coding model for agent workflows, office edits, and automation", - "source": "https://api.friendli.ai/serverless/v1/models" - }, "friendliai/zai-org/GLM-5.1": { "litellm_provider": "friendliai", "max_input_tokens": 202752, @@ -30575,7 +30538,6 @@ "successor_model": "gemini-3.8-flash", "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "cache_read_input_token_cost_above_200k_tokens_priority": 4.5e-07, "cache_read_input_token_cost_flex": 1.25e-07, "cache_read_input_token_cost_priority": 2.25e-07, @@ -30634,7 +30596,6 @@ "output_cost_per_token_priority": 1.8e-05 }, "gemini-3.1-pro-preview": { - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, @@ -30698,7 +30659,6 @@ "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, "input_cost_per_token_batches": 1e-06, @@ -30750,7 +30710,6 @@ "vertex_ai/gemini-3-pro-preview": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, "input_cost_per_token_batches": 1e-06, @@ -31194,7 +31153,6 @@ "supports_web_search": true }, "vertex_ai/gemini-3.1-pro-preview": { - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, @@ -31258,7 +31216,6 @@ "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, "input_cost_per_token_batches": 1e-06, @@ -43150,6 +43107,29 @@ "max_input_tokens": 1000000, "max_output_tokens": 128000 }, + "minimax/MiniMax-M3.1-Flash": { + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "cache_read_input_token_cost": 1.5e-08, + "litellm_provider": "minimax", + "mode": "chat", + "source": "https://platform.minimax.io/docs/guides/pricing-paygo", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_none_reasoning_effort": false, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "reasoning_effort_levels": [ + "low", + "medium", + "high", + "xhigh", + "max" + ], + "supports_vision": true, + "max_input_tokens": 1000000 + }, "minimax/minimax-latest": { "alias_of": "minimax/MiniMax-M3" }, @@ -53162,8 +53142,49 @@ "output_cost_per_token_batches": 2.5e-05, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" }, + "vertex_ai/claude-fable-5-2": { + "regional_endpoint_uplift_multiplier": 1.1, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_creation_input_token_cost_batches": 5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "prompt_cache_min_tokens": 512, + "deprecation_date": "2027-11-09", + "input_cost_per_token_batches": 4e-06, + "output_cost_per_token_batches": 2e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" + }, "vertex_ai/claude-fable-latest": { - "alias_of": "vertex_ai/claude-fable-5-1" + "alias_of": "vertex_ai/claude-fable-5-2" }, "vertex_ai/claude-fable-5@default": { "deprecation_date": "2027-06-08", @@ -53245,6 +53266,47 @@ "output_cost_per_token_batches": 2.5e-05, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" }, + "vertex_ai/claude-fable-5-2@default": { + "regional_endpoint_uplift_multiplier": 1.1, + "supports_mid_conversation_system": true, + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_creation_input_token_cost_batches": 5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "vertex_ai-anthropic_models", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "prompt_cache_min_tokens": 512, + "deprecation_date": "2027-11-09", + "input_cost_per_token_batches": 4e-06, + "output_cost_per_token_batches": 2e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" + }, "vertex_ai/claude-opus-5": { "deprecation_date": "2027-01-24", "regional_endpoint_uplift_multiplier": 1.1, @@ -65172,6 +65234,14 @@ "supports_reasoning": true } }, + { + "name": "minimax-reasoning-baseline", + "pattern": "^minimax/", + "description": "Any MiniMax model id under the minimax/ namespace, anchored so only that provider's ids match. MiniMax ships reasoning models faster than this registry names them, so an id the map has not described yet is treated as reasoning-capable and keeps the caller's reasoning_effort and thinking instead of dropping them. Rules lose to exact entries, so mapped speech, video and image models are unaffected. Carries no mode and no pricing, so cost stays on the standard unpriced behavior.", + "model_info": { + "supports_reasoning": true + } + }, { "name": "openai-reasoning-family-baseline", "pattern": "^(?!.*search-api)(?:[a-z0-9_.-]+/)*(?:ft:)?(?:o[1-9]\\d*(?![a-z0-9])|gpt-[5-9](?:\\.\\d+)?(?![0-9.])|(?:gpt-\\d+(?:\\.\\d+)?(?:-[a-z0-9]+)*-)?(?:codex|deep-research|chat-latest)(?![a-z0-9]))", @@ -67761,7 +67831,7 @@ "deepinfra/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B": { "max_output_tokens": 262144, "max_input_tokens": 262144, - "input_cost_per_token": 8.5e-08, + "input_cost_per_token": 7.5e-08, "output_cost_per_token": 4e-07, "litellm_provider": "deepinfra", "mode": "chat", @@ -69888,6 +69958,44 @@ "region": "us-gov-west-1" } }, + "bedrock/us-gov-west-1/anthropic.claude-fable-5-2": { + "cache_creation_input_token_cost": 1.2e-05, + "cache_creation_input_token_cost_above_1hr": 1.92e-05, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 9.6e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.8e-05, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh", + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512, + "token_count_model": { + "model": "anthropic.claude-fable-5-2", + "litellm_provider": "bedrock_mantle", + "region": "us-gov-west-1" + } + }, "bedrock/us-gov-east-1/nvidia.nemotron-nano-3-30b": { "input_cost_per_token": 7.2e-08, "input_cost_per_token_batches": 3.6e-08, @@ -70196,6 +70304,39 @@ "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512 }, + "bedrock/us-gov-east-1/anthropic.claude-fable-5-2": { + "cache_creation_input_token_cost": 1.2e-05, + "cache_creation_input_token_cost_above_1hr": 1.92e-05, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 9.6e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.8e-05, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_native_structured_output": false, + "supports_bedrock_native_structured_output": false, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "bedrock_output_config_effort_ceiling": "xhigh", + "supports_parallel_tool_use_config": true, + "prompt_cache_min_tokens": 512 + }, "bedrock_mantle/us-gov-west-1/openai.gpt-5.6-terra": { "litellm_provider": "bedrock_mantle", "max_input_tokens": 922000, @@ -72924,6 +73065,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/gpt-realtime-2": { + "cache_read_input_image_token_cost": 5e-07, "cache_creation_input_audio_token_cost": 4e-07, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, @@ -72942,7 +73084,7 @@ "provider_specific_entry": { "realtime_ga_only": 1 }, - "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_endpoints": [ "/v1/realtime" ], @@ -83307,13 +83449,16 @@ "cache_creation_input_token_cost": 2.5e-05, "cache_creation_input_token_cost_above_12hr": 5e-05, "cache_read_input_token_cost": 2e-06, + "cache_read_input_token_cost_above_272k_tokens": 4e-06, "input_cost_per_token": 2e-05, + "input_cost_per_token_above_272k_tokens": 4e-05, "litellm_provider": "openai", "max_input_tokens": 922000, "max_output_tokens": 128000, "mode": "chat", "output_cost_per_token": 0.0001, - "source": "https://developers.openai.com/api/docs/models/gpt-6-cyber", + "output_cost_per_token_above_272k_tokens": 0.00015, + "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -87201,13 +87346,13 @@ "supports_vision": true }, "wandb/deepseek-ai/DeepSeek-V4.1-Flash-1017": { - "cache_read_input_token_cost": 2e-09, - "input_cost_per_token": 1.3e-07, + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, "litellm_provider": "wandb", "max_input_tokens": 1049000, "max_output_tokens": 1048576, "mode": "chat", - "output_cost_per_token": 6.5e-07, + "output_cost_per_token": 1.2e-06, "source": "https://www.coreweave.com/products/coreweave-forge/serverless-inference", "supports_prompt_caching": true, "supports_reasoning": true, @@ -87352,5 +87497,574 @@ "supports_parallel_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true + }, + "anthropic.claude-fable-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "eu.anthropic.claude-fable-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.1e-05, + "cache_creation_input_token_cost_above_1hr": 1.76e-05, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 8.8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output_by_default": true, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "eu.anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.1e-05, + "cache_creation_input_token_cost_above_1hr": 1.76e-05, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 8.8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "global.anthropic.claude-fable-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output_by_default": true, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "global.anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "us-gov.anthropic.claude-fable-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.2e-05, + "cache_creation_input_token_cost_above_1hr": 1.92e-05, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 9.6e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.8e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "us-gov.anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.2e-05, + "cache_creation_input_token_cost_above_1hr": 1.92e-05, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 9.6e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.8e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "us.anthropic.claude-fable-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.1e-05, + "cache_creation_input_token_cost_above_1hr": 1.76e-05, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 8.8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output_by_default": true, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "us.anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.1e-05, + "cache_creation_input_token_cost_above_1hr": 1.76e-05, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 8.8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "claude-fable-5-2": { + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_creation_input_token_cost_batches": 5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, + "input_cost_per_token": 8e-06, + "input_cost_per_token_batches": 4e-06, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "output_cost_per_token_batches": 2e-05, + "prompt_cache_min_tokens": 512.0, + "provider_specific_entry": {"us":1.1,"fast":2}, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://platform.claude.com/docs/en/about-claude/pricing", + "supports_adaptive_thinking": true, + "supports_anthropic_compaction": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_fast_mode": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_speed": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "bedrock/claude_platform/claude-fable-5-2": { + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_creation_input_token_cost_batches": 5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, + "input_cost_per_token": 8e-06, + "input_cost_per_token_batches": 4e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "output_cost_per_token_batches": 2e-05, + "prompt_cache_min_tokens": 512.0, + "provider_specific_entry": {"us":1.1}, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing", + "supports_adaptive_thinking": true, + "supports_anthropic_compaction": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "azure_ai/claude-fable-5-2": { + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://platform.claude.com/docs/en/about-claude/pricing", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": false, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "azure_ai/claude-fable-5-2-fast": { + "cache_creation_input_token_cost": 2e-05, + "cache_creation_input_token_cost_above_1hr": 3.2e-05, + "cache_read_input_token_cost": 4e-07, + "input_cost_per_token": 1.6e-05, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 8e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://platform.claude.com/docs/en/about-claude/pricing", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": false, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "bedrock/us-gov-east-1/anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.2e-05, + "cache_creation_input_token_cost_above_1hr": 1.92e-05, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 9.6e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.8e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "bedrock/us-gov-west-1/anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.2e-05, + "cache_creation_input_token_cost_above_1hr": 1.92e-05, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 9.6e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.8e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "github_copilot/claude-fable-5.2": { + "cache_creation_input_token_cost": 1e-05, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "github_copilot", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "source": "https://raw.githubusercontent.com/github/docs/main/data/tables/copilot/models-and-pricing.yml", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/messages" + ], + "supports_adaptive_thinking": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "thinking_always_on": true } } diff --git a/pyproject.toml b/pyproject.toml index 02834f604d2..de1ea342017 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,6 +56,7 @@ Documentation = "https://docs.litellm.ai" [project.optional-dependencies] switchyard = ["nemo-switchyard==0.4.0a1"] proxy = [ + "casbin==1.43.0", "gunicorn>=23.0.0,<24.0", "uvicorn>=0.33.0,<1.0", "granian>=2.7.4,<3.0; python_version < '3.15'", diff --git a/tests/e2e/llm_translation/test_together_ai_e2e.py b/tests/e2e/llm_translation/test_together_ai_e2e.py index 2d5d5f2d93a..7775c5c6ed2 100644 --- a/tests/e2e/llm_translation/test_together_ai_e2e.py +++ b/tests/e2e/llm_translation/test_together_ai_e2e.py @@ -7,9 +7,9 @@ backends are pinned because the registry has no flag for what they prove: ``enab and the ``{"reasoning": {"enabled": false}}`` toggle that ``reasoning_effort="none"`` maps to are Qwen hybrid-model contracts, and the structured-output case runs on that hybrid model with reasoning off, since a reasoning-only model can spend the whole token budget -thinking and return no content. MiniMax-M3 is the serverless model whose +thinking and return no content. MiniMax-M3.1-Flash is the serverless model whose template renders a replayed ``reasoning_content`` back into the prompt (Qwen and -DeepSeek silently drop it). MiniMax-M3 honors that replayed field on nearly every call, not +DeepSeek silently drop it). MiniMax-M3.1-Flash honors that replayed field on nearly every call, not every call (one miss in dozens of otherwise identical calls), so the replay case asks up to ``REPLAY_ATTEMPTS`` times and fails only when no answer carries the secret, which a proxy that strips the field guarantees. Requires TOGETHER_API_KEY on the proxy; no @@ -56,7 +56,7 @@ from pydantic import BaseModel pytestmark = pytest.mark.e2e HYBRID_REASONING_BACKEND = "together_ai/Qwen/Qwen3.5-9B" -REASONING_REPLAY_BACKEND = "together_ai/MiniMaxAI/MiniMax-M3" +REASONING_REPLAY_BACKEND = "together_ai/MiniMaxAI/MiniMax-M3.1-Flash" SECRET_PROMPT = "Remember this for later and reply with just OK." SECRET_REASONING = "The user told me their favorite color is chartreuse. I must remember it." diff --git a/tests/integration/_support/database_relay.py b/tests/integration/_support/database_relay.py index 0ac1e1d2521..0ac060dc867 100644 --- a/tests/integration/_support/database_relay.py +++ b/tests/integration/_support/database_relay.py @@ -91,7 +91,6 @@ class HeldStatementRelay: self._upstream_port: Final = upstream_port self._trigger: Final = trigger self._after: Final = after - self._marker_length: Final = max(len(trigger), len(after or b"")) self._server: Final = LoopbackServer(self.port, self._serve) self._loop: Final = self._server.loop self._released: Final = asyncio.Event() @@ -115,27 +114,25 @@ class HeldStatementRelay: self.release() self._server.stop() - def _holds(self, window: bytes) -> bool: + def _holds(self, markers: frozenset[bytes]) -> bool: if not self._armed.is_set() or self.held.is_set(): return False if self._after is not None and not self._after_seen.is_set(): - if self._after in window: + if self._after in markers: self._after_seen.set() return False - return self._trigger in window + return self._trigger in markers async def _serve(self, client_reader: asyncio.StreamReader, client_writer: asyncio.StreamWriter) -> None: server_reader, server_writer = await asyncio.open_connection(self._upstream_host, self._upstream_port) async def forward(reader: asyncio.StreamReader, writer: asyncio.StreamWriter, inspect: bool) -> None: - tail = b"" # rebind-ok: carries the previous read's end so a trigger split across reads still matches + scanner: Final = StatementMarkerScanner((self._trigger,) + ((self._after,) if self._after else ())) try: while chunk := await reader.read(65536): - window: Final = tail + chunk - if inspect and self._holds(window): + if inspect and any(self._holds(markers) for markers in scanner.feed(chunk)): self.held.set() await self._released.wait() - tail = window[-(self._marker_length - 1) :] writer.write(chunk) await writer.drain() except (ConnectionError, asyncio.IncompleteReadError): @@ -149,6 +146,47 @@ class HeldStatementRelay: ) +class StatementMarkerScanner: + def __init__(self, markers: tuple[bytes, ...]) -> None: + self._markers: Final = markers + self._statements: Final[dict[bytes, frozenset[bytes]]] = {} + self._buffer = b"" + self._startup = True + + def feed(self, chunk: bytes) -> Generator[frozenset[bytes]]: + self._buffer += chunk + while len(self._buffer) >= (4 if self._startup else 5): + offset: Final = 0 if self._startup else 1 + length: Final = int.from_bytes(self._buffer[offset : offset + 4], "big") + if length < 4: + raise ValueError("Invalid PostgreSQL message length") + if len(self._buffer) < length + offset: + return + message: Final = self._buffer[: length + offset] + self._buffer = self._buffer[length + offset :] + if self._startup: + self._startup = message[4:8] == b"\x04\xd2\x16/" + continue + payload: Final = message[5:] + if message[:1] == b"P": + name, _, query = payload.partition(b"\x00") + matches: Final = self._matches(query.split(b"\x00", 1)[0]) + if name: + self._statements[name] = matches + yield matches + elif message[:1] == b"B": + fields: Final = payload.split(b"\x00", 2) + prepared: Final = self._statements.get(fields[1], frozenset()) if len(fields) > 1 else frozenset() + yield prepared | self._matches(payload) + elif message[:1] == b"Q": + yield self._matches(payload) + elif message[:1] == b"C" and payload[:1] == b"S": + self._statements.pop(payload[1:].split(b"\x00", 1)[0], None) + + def _matches(self, query: bytes) -> frozenset[bytes]: + return frozenset(marker for marker in self._markers if marker in query) + + class TriggerScanner: def __init__(self, trigger: bytes) -> None: self._trigger: Final = trigger diff --git a/tests/integration/authorization/test_access_group_rules.py b/tests/integration/authorization/test_access_group_rules.py new file mode 100644 index 00000000000..68decdbf7aa --- /dev/null +++ b/tests/integration/authorization/test_access_group_rules.py @@ -0,0 +1,221 @@ +import json +import time +import uuid +from pathlib import Path +from typing import Final + +import jwt +import pytest +import yaml +from cryptography.hazmat.primitives.asymmetric import rsa + +from litellm.proxy.mcp_server.idp_token_exchange import REJECTED_SUBJECT_TOKEN +from tests.integration._support.client import Gateway, string_value +from tests.integration._support.database import read_rows +from tests.integration._support.process import owned_proxy +from tests.integration._support.wire import Reply, Request, wire_server + +pytestmark: Final = pytest.mark.shares_worker_state + + +@pytest.mark.parametrize( + "enabled,mapping,claimed", ((True, False, True), (False, False, True), (True, True, False), (True, True, True)) +) +def test_rules_add_models_to_jwt_and_scoped_virtual_key_without_opening_other_models( + gateway: Gateway, tmp_path: Path, enabled: bool, mapping: bool, claimed: bool +) -> None: + private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + public_key: Final = json.loads(jwt.algorithms.RSAAlgorithm.to_jwk(private_key.public_key())) + key_id: Final = "rule-signing-key" + + def respond(request: Request) -> Reply: + assert request.method == "GET", request + return Reply(body=json.dumps({"keys": [{**public_key, "kid": key_id}]}).encode()) + + with wire_server(respond) as jwks, gateway.scenario() as scenario: + explicit: Final = scenario.model() + group_model: Final = scenario.model() + outside: Final = scenario.model() + group_name: Final = f"integration-rule-{uuid.uuid4().hex}" + created: Final = gateway.request( + "POST", "/v1/access_group", {"access_group_name": group_name, "access_model_names": [group_model]} + ) + assert created.status_code == 201, created.text + group_id: Final = string_value(created.json()["access_group_id"]) + scenario.cleanups.callback(gateway.request, "DELETE", f"/v1/access_group/{group_id}") + virtual_user: Final = scenario.user() + virtual_key: Final = scenario.key(models=[explicit], user_id=virtual_user) + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["general_settings"] = { + **config["general_settings"], + "enable_jwt_auth": True, + "litellm_jwtauth": { + "user_id_jwt_field": "sub", + "team_ids_jwt_field": "groups", + "user_id_upsert": True, + "enforce_team_based_model_access": True, + **( + {"virtual_key_claim_field": "sub", "unregistered_jwt_client_behavior": "auto_register"} + if mapping + else {} + ), + }, + **( + { + "access_group_rules": [ + {"attribute": "jwt_claims.groups", "contains": "science", "access_group": group_name}, + {"attribute": "user_id", "equals": virtual_user, "access_group": group_name}, + ] + } + if enabled + else {} + ), + } + path: Final = tmp_path / "access_group_rules.yaml" + path.write_text(yaml.safe_dump(config)) + subject: Final = f"integration-rule-user-{uuid.uuid4().hex}" + gateway.post("/user/new", {"user_id": subject, "user_role": "internal_user"}) + scenario.cleanups.callback(scenario.delete_user, subject) + team_id: Final = ( + scenario.team(models=[explicit], members_with_roles=[{"role": "user", "user_id": subject}]) + if claimed + else None + ) + team_claims: Final = [team_id] if team_id is not None else [] + token: Final = jwt.encode( + {"sub": subject, "groups": ["science", *team_claims], "exp": int(time.time()) + 300}, + private_key, + algorithm="RS256", + headers={"kid": key_id}, + ) + unmatched: Final = jwt.encode( + {"sub": subject, "groups": ["other", *team_claims], "exp": int(time.time()) + 300}, + private_key, + algorithm="RS256", + headers={"kid": key_id}, + ) + with owned_proxy(gateway, tmp_path, {"JWT_PUBLIC_KEY_URL": jwks.url}, config=path, workers=2) as candidate: + fallback_response: Final = candidate.request( + "POST", + "/v1/chat/completions", + { + "model": group_model, + "fallbacks": [outside], + "messages": [{"role": "user", "content": "ungranted fallback control"}], + }, + key=token, + ) + assert fallback_response.status_code == 403, fallback_response.text + for key, model, allowed in ( + (token, group_model, enabled and claimed), + *(((token, group_model, enabled and claimed),) if mapping else ()), + (token, outside, False), + (unmatched, group_model, False), + (virtual_key, explicit, True), + (virtual_key, group_model, enabled), + (virtual_key, outside, False), + ): + response: Final = candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "rule authorization control"}]}, + key=key, + ) + assert response.status_code == (200 if allowed else 403), response.text + if allowed: + assert ( + response.json()["choices"][0]["message"]["content"] + == "Hello! This is a mock response from the fake OpenAI endpoint." + ), response.text + + if mapping and claimed: + rows: Final = read_rows( + 'SELECT models, access_group_ids FROM "LiteLLM_VerificationToken" ' + 'WHERE user_id=%s AND team_id=%s', + (subject, team_id), + ) + assert len(rows) == 1, rows + assert group_model not in rows[0]["models"], rows + assert not rows[0]["access_group_ids"], rows + + +@pytest.mark.parametrize("consumer", ("auto_register", "token_exchange")) +@pytest.mark.parametrize("matching", (True, False)) +def test_teamless_jwt_cannot_provision_credentials_without_native_team_access( + gateway: Gateway, tmp_path: Path, consumer: str, matching: bool +) -> None: + private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + public_key: Final = json.loads(jwt.algorithms.RSAAlgorithm.to_jwk(private_key.public_key())) + key_id: Final = "credential-rule-signing-key" + + def respond(request: Request) -> Reply: + assert request.method == "GET", request + return Reply(body=json.dumps({"keys": [{**public_key, "kid": key_id}]}).encode()) + + with wire_server(respond) as jwks, gateway.scenario() as scenario: + model: Final = scenario.model() + subject: Final = f"integration-credential-user-{uuid.uuid4().hex}" + gateway.post("/user/new", {"user_id": subject, "user_role": "internal_user"}) + scenario.cleanups.callback(scenario.delete_user, subject) + group_name: Final = f"credential-rule-{uuid.uuid4().hex}" + created_group: Final = gateway.request( + "POST", "/v1/access_group", {"access_group_name": group_name, "access_model_names": [model]} + ) + assert created_group.status_code == 201, created_group.text + group_id: Final = string_value(created_group.json()["access_group_id"]) + scenario.cleanups.callback(gateway.request, "DELETE", f"/v1/access_group/{group_id}") + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["general_settings"] = { + **config["general_settings"], + "enable_jwt_auth": True, + "litellm_jwtauth": { + "user_id_jwt_field": "sub", + "enforce_team_based_model_access": True, + **( + {"virtual_key_claim_field": "sub", "unregistered_jwt_client_behavior": "auto_register"} + if consumer == "auto_register" + else {} + ), + }, + "access_group_rules": [ + {"attribute": "jwt_claims.groups", "contains": "science", "access_group": group_name} + ], + } + path: Final = tmp_path / "credential_access_group_rules.yaml" + path.write_text(yaml.safe_dump(config)) + token: Final = jwt.encode( + {"sub": subject, "groups": ["science" if matching else "other"], "exp": int(time.time()) + 300}, + private_key, + algorithm="RS256", + headers={"kid": key_id}, + ) + with owned_proxy(gateway, tmp_path, {"JWT_PUBLIC_KEY_URL": jwks.url}, config=path, workers=2) as candidate: + if consumer == "auto_register": + for _ in range(2): + response: Final = candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "credential rejection control"}]}, + key=token, + ) + assert response.status_code == 403, response.text + assert "No teams found in token" in response.text, response.text + else: + registered: Final = candidate.client.post( + "/register", json={"redirect_uris": ["http://localhost/callback"], "client_name": "rule-control"} + ) + assert registered.status_code in (200, 201), registered.text + exchanged: Final = candidate.client.post( + "/token", + data={ + "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange", + "client_id": registered.json()["client_id"], + "subject_token": token, + "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", + "requested_token_type": "urn:ietf:params:oauth:token-type:access_token", + }, + ) + assert exchanged.status_code == 400, exchanged.text + assert exchanged.json() == {"error": "invalid_request", "error_description": REJECTED_SUBJECT_TOKEN}, ( + exchanged.text + ) diff --git a/tests/integration/authorization/test_responses_compact_id_ownership.py b/tests/integration/authorization/test_responses_compact_id_ownership.py index b5a6215dd15..c6486a61c85 100644 --- a/tests/integration/authorization/test_responses_compact_id_ownership.py +++ b/tests/integration/authorization/test_responses_compact_id_ownership.py @@ -442,7 +442,7 @@ def test_provider_errors_on_an_owned_compact_reach_the_caller(gateway: Gateway, assert_compact_rows(actors.owner, frozenset({after})) -def test_response_minted_by_a_key_without_owner_is_addressable_by_any_key(gateway: Gateway) -> None: +def test_response_minted_by_a_key_without_user_or_team_is_addressable_only_by_that_key(gateway: Gateway) -> None: upstream: Final = CompactUpstream() with wire_server(upstream.respond) as wire, gateway.scenario() as scenario: actors: Final = tenants(scenario, wire) @@ -451,7 +451,18 @@ def test_response_minted_by_a_key_without_owner_is_addressable_by_any_key(gatewa minted: Final = mint(gateway, ownerless, actors.model, created) upstreamed(wire) text: Final = marker("compact_ownerless") - assert_compacted(sdk_compact(gateway, actors.other_user, "/v1", actors.model, text, minted), actors.model, text) + assert_refused( + compact_post( + gateway, + actors.other_user, + "/v1", + compact_body(actors.model, text, {"previous_response_id": minted}), + text, + ), + OTHER_USER_DETAIL, + ) + assert compacts(upstreamed(wire)) == () + assert_compacted(sdk_compact(gateway, ownerless, "/v1", actors.model, text, minted), actors.model, text) assert compacts(upstreamed(wire)) == ( upstream_compact_body(text, {"previous_response_id": minted_id(created)}), ) diff --git a/tests/integration/authorization/test_team_scoped_models.py b/tests/integration/authorization/test_team_scoped_models.py index 8f8e12ae498..433a9e34e04 100644 --- a/tests/integration/authorization/test_team_scoped_models.py +++ b/tests/integration/authorization/test_team_scoped_models.py @@ -137,12 +137,13 @@ def _v2_team_public_names(gateway: Gateway, key: str, model: str) -> list[JsonVa return [entry["model_info"].get("team_public_model_name") for entry in response.json()["data"]] -def test_v2_model_info_reports_a_team_model_to_team_and_non_team_keys(gateway: Gateway) -> None: +def test_v2_model_info_reports_a_team_model_only_to_team_and_admin_keys(gateway: Gateway) -> None: with gateway.scenario() as scenario: team: Final = scenario.team(models=[]) model: Final = scenario.model(model_info={"team_id": team}) assert _v2_team_public_names(gateway, scenario.key(team_id=team), model) == [model] - assert _v2_team_public_names(gateway, scenario.key(), model) == [model] + assert _v2_team_public_names(gateway, scenario.key(), model) == [] + assert _v2_team_public_names(gateway, gateway.key, model) == [model] @pytest.mark.parametrize("set_on", ["team_new", "team_update"]) diff --git a/tests/integration/compatibility/test_audio_transcriptions_wire.py b/tests/integration/compatibility/test_audio_transcriptions_wire.py index faebfda715b..f943400bb15 100644 --- a/tests/integration/compatibility/test_audio_transcriptions_wire.py +++ b/tests/integration/compatibility/test_audio_transcriptions_wire.py @@ -832,12 +832,13 @@ def test_gpt_transcribe_unmeasurable_upload_bills_the_seconds_openai_reports( assert _uploaded_files(requests[0]) == (("file", "a.m4a", upload),) -def test_gpt_transcribe_measured_upload_bills_the_measured_length_over_reported_seconds(gateway: Gateway) -> None: +def test_gpt_transcribe_measured_upload_bills_reported_seconds_over_the_measured_length(gateway: Gateway) -> None: measured_seconds: Final = 1.0 + reported_seconds: Final = 7.0 upload: Final = _silent_wav(measured_seconds) with ( - wire_server(lambda request: _openai_gpt_transcribe_reply(request, 7)) as wire, + wire_server(lambda request: _openai_gpt_transcribe_reply(request, reported_seconds)) as wire, gateway.scenario() as scenario, ): model: Final = scenario.model( @@ -851,9 +852,13 @@ def test_gpt_transcribe_measured_upload_bills_the_measured_length_over_reported_ {"model": model, "response_format": "json"}, {"file": ("a.wav", upload, _WAV_CONTENT_TYPE)}, ) - expected_cost: Final = measured_seconds * _DEPLOYMENT_COST_PER_SECOND + expected_cost: Final = reported_seconds * _DEPLOYMENT_COST_PER_SECOND assert response.status_code == 200, response.text + assert JSON_OBJECT.validate_json(response.content) == { + "text": "hello world", + "usage": {"type": "duration", "seconds": reported_seconds}, + }, response.text assert math.isclose(_response_cost(response), expected_cost), (_response_cost(response), expected_cost) assert math.isclose(_logged_spend(response.headers["x-litellm-call-id"]), expected_cost) requests: Final = wire.drain() diff --git a/tests/integration/cost_calculation/cases/anthropic.py b/tests/integration/cost_calculation/cases/anthropic.py index 776c5b51dd4..565d521aa95 100644 --- a/tests/integration/cost_calculation/cases/anthropic.py +++ b/tests/integration/cost_calculation/cases/anthropic.py @@ -4,8 +4,8 @@ PARITY pairs the streamed case with its plain twin where both bill the same row. from typing import Final +from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParityTestCase, bridge_pairs from integration.cost_calculation.cost_tracking_case import ( - with_usage, CostTrackingTestCase, ExactExpected, FailureDetails, @@ -14,11 +14,10 @@ from integration.cost_calculation.cost_tracking_case import ( RecountExpected, RecountRates, SseResponse, + with_usage, ) -from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParityTestCase, bridge_pairs from integration.cost_calculation.stream_parity.case import StreamParityTestCase, sse_frames - # claude-haiku-4-5 CLAUDE_HAIKU_4_5_INPUT_TEXT: Final = CostTrackingTestCase( name="claude-haiku-4-5-input_text", @@ -5344,6 +5343,66 @@ CLAUDE_FABLE_5_1_INPUT_TEXT: Final = CostTrackingTestCase( ), ) +CLAUDE_FABLE_5_2_INPUT_TEXT: Final = CostTrackingTestCase( + name="claude-fable-5-2-input_text", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-fable-5-2", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "7d8c0a2d1ecc summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-fable-5-2", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted answer 7d8c0a2d1ecc"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 412, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + "speed": "standard", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.03432, input_cost=0.016192, output_cost=0.018128, prompt_tokens=1840, completion_tokens=412 + ), +) + CLAUDE_FABLE_5_1_STREAM: Final = CostTrackingTestCase( name="claude-fable-5-1-stream", covers="quota_management.spend_tracking.scripted_wire.logs_cost", @@ -5429,6 +5488,91 @@ CLAUDE_FABLE_5_1_STREAM: Final = CostTrackingTestCase( ), ) +CLAUDE_FABLE_5_2_STREAM: Final = CostTrackingTestCase( + name="claude-fable-5-2-stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="claude-fable-5-2", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "459736de6b4a summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": True, + "stream_options": {"include_usage": True}, + "allowed_openai_params": [], + }, + response=SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "type": "message_start", + "message": { + "model": "claude-fable-5-2", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [], + "container": None, + "stop_reason": None, + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 1, + "service_tier": "standard", + "inference_geo": "global", + "speed": "standard", + }, + "diagnostics": None, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "ping"}, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "scripted answer 459736de6b4a"}, + }, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None, "stop_details": None, "container": None}, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "output_tokens": 412, + "output_tokens_details": {"thinking_tokens": 0}, + }, + }, + {"type": "message_stop"}, + ), + ), + expected=ExactExpected( + spend=0.03432, input_cost=0.016192, output_cost=0.018128, prompt_tokens=1840, completion_tokens=412 + ), +) + CLAUDE_FABLE_5_1_CACHE_READ: Final = CostTrackingTestCase( name="claude-fable-5-1-cache_read", covers="quota_management.spend_tracking.cost_matrix.logs_cost", @@ -5490,6 +5634,67 @@ CLAUDE_FABLE_5_1_CACHE_READ: Final = CostTrackingTestCase( ), ) +CLAUDE_FABLE_5_2_CACHE_READ: Final = CostTrackingTestCase( + name="claude-fable-5-2-cache_read", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-fable-5-2", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "c5a9841fc31b summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-fable-5-2", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted answer c5a9841fc31b"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 640, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 12288, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 380, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + "speed": "standard", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.02505536, input_cost=0.00833536, output_cost=0.01672, prompt_tokens=12928, completion_tokens=380 + ), +) + CLAUDE_FABLE_5_1_CACHE_WRITE_5M: Final = CostTrackingTestCase( name="claude-fable-5-1-cache_write_5m", covers="quota_management.spend_tracking.cost_matrix.logs_cost", @@ -5550,6 +5755,66 @@ CLAUDE_FABLE_5_1_CACHE_WRITE_5M: Final = CostTrackingTestCase( ), ) +CLAUDE_FABLE_5_2_CACHE_WRITE_5M: Final = CostTrackingTestCase( + name="claude-fable-5-2-cache_write_5m", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-fable-5-2", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "95e59856831e summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-fable-5-2", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted answer 95e59856831e"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 512, + "cache_creation_input_tokens": 9216, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 9216, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 350, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.1212816, input_cost=0.1058816, output_cost=0.0154, prompt_tokens=9728, completion_tokens=350 + ), +) + CLAUDE_FABLE_5_1_CACHE_WRITE_1H: Final = CostTrackingTestCase( name="claude-fable-5-1-cache_write_1h", covers="quota_management.spend_tracking.cost_matrix.logs_cost", @@ -5610,6 +5875,66 @@ CLAUDE_FABLE_5_1_CACHE_WRITE_1H: Final = CostTrackingTestCase( ), ) +CLAUDE_FABLE_5_2_CACHE_WRITE_1H: Final = CostTrackingTestCase( + name="claude-fable-5-2-cache_write_1h", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-fable-5-2", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "0057eda38d9c summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-fable-5-2", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted answer 0057eda38d9c"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 512, + "cache_creation_input_tokens": 9216, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 9216}, + "output_tokens": 350, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.1821072, input_cost=0.1667072, output_cost=0.0154, prompt_tokens=9728, completion_tokens=350 + ), +) + CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT: Final = CostTrackingTestCase( name="claude-fable-5-1-messages_input_text", covers="quota_management.spend_tracking.cost_matrix.logs_cost", @@ -5651,6 +5976,47 @@ CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT: Final = CostTrackingTestCase( ), ) +CLAUDE_FABLE_5_2_MESSAGES_INPUT_TEXT: Final = CostTrackingTestCase( + name="claude-fable-5-2-messages_input_text", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-fable-5-2", + endpoint="/v1/messages", + request={ + "model": "$MODEL", + "max_tokens": 412, + "messages": [{"role": "user", "content": "summarize the attached material in one line"}], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-fable-5-2", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted response"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 412, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + "speed": "standard", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.03432, input_cost=0.016192, output_cost=0.018128, prompt_tokens=1840, completion_tokens=412 + ), +) + CLAUDE_FABLE_5_1_MESSAGES_STREAM: Final = CostTrackingTestCase( name="claude-fable-5-1-messages_stream", @@ -5680,6 +6046,35 @@ CLAUDE_FABLE_5_1_MESSAGES_STREAM: Final = CostTrackingTestCase( spend=0.0624, input_cost=0.02944, output_cost=0.03296, prompt_tokens=1840, completion_tokens=412 ), ) + +CLAUDE_FABLE_5_2_MESSAGES_STREAM: Final = CostTrackingTestCase( + name="claude-fable-5-2-messages_stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="claude-fable-5-2", + endpoint="/v1/messages", + request={ + "model": "$MODEL", + "max_tokens": 412, + "messages": [{"role": "user", "content": "summarize the attached material in one line"}], + "stream": True, + }, + response=SseResponse( + content_type="text/event-stream", + frames=( + 'event: message_start\ndata: {"type":"message_start","message":{"model":"claude-fable-5-2","id":"msg_$REQUEST_ID","type":"message","role":"assistant","content":[],"container":null,"stop_reason":null,"stop_sequence":null,"stop_details":null,"usage":{"input_tokens":1840,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"cache_creation":{"ephemeral_5m_input_tokens":0,"ephemeral_1h_input_tokens":0},"output_tokens":1,"service_tier":"standard","inference_geo":"global","speed":"standard"},"diagnostics":null}}', + 'event: content_block_start\ndata: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}', + 'event: ping\ndata: {"type":"ping"}', + 'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"scripted "}}', + 'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"response"}}', + 'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}', + 'event: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null,"stop_details":null,"container":null},"usage":{"input_tokens":1840,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"output_tokens":412,"output_tokens_details":{"thinking_tokens":0}}}', + 'event: message_stop\ndata: {"type":"message_stop"}', + ), + ), + expected=ExactExpected( + spend=0.03432, input_cost=0.016192, output_cost=0.018128, prompt_tokens=1840, completion_tokens=412 + ), +) CLAUDE_OPUS_5_5_STREAM: Final = CostTrackingTestCase( name="claude-opus-5-5-stream", covers="quota_management.spend_tracking.scripted_wire.logs_cost", @@ -6560,6 +6955,92 @@ CLAUDE_FABLE_5_1_TOOL_CALL: Final = CostTrackingTestCase( ), ) +CLAUDE_FABLE_5_2_TOOL_CALL: Final = CostTrackingTestCase( + name="claude-fable-5-2-tool_call", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-fable-5-2", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "6c08522a6809 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather and a short forecast for a city.", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "City name"}, + "days": {"type": "integer", "description": "Forecast horizon in days"}, + "units": {"type": "string", "enum": ["metric", "imperial"]}, + }, + "required": ["city"], + }, + }, + } + ], + "tool_choice": "auto", + "allowed_openai_params": ["tool_choice"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-fable-5-2", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "toolu_$REQUEST_ID", + "name": "get_weather", + "input": {"city": "Berlin", "days": 7, "units": "metric"}, + "caller": {"type": "direct"}, + } + ], + "container": None, + "stop_reason": "tool_use", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 412, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.03432, input_cost=0.016192, output_cost=0.018128, prompt_tokens=1840, completion_tokens=412 + ), +) + CLAUDE_FABLE_5_1_THINKING: Final = CostTrackingTestCase( name="claude-fable-5-1-thinking", covers="quota_management.spend_tracking.cost_matrix.logs_cost", @@ -6621,6 +7102,67 @@ CLAUDE_FABLE_5_1_THINKING: Final = CostTrackingTestCase( ), ) +CLAUDE_FABLE_5_2_THINKING: Final = CostTrackingTestCase( + name="claude-fable-5-2-thinking", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="claude-fable-5-2", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "697b4f57a5df summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "reasoning_effort": "medium", + "allowed_openai_params": ["reasoning_effort"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "model": "claude-fable-5-2", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "scripted answer 697b4f57a5df"}], + "container": None, + "stop_reason": "end_turn", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1240, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 412, + "output_tokens_details": {"thinking_tokens": 0}, + "service_tier": "standard", + "inference_geo": "global", + "speed": "standard", + }, + "diagnostics": None, + }, + ), + expected=ExactExpected( + spend=0.02904, input_cost=0.010912, output_cost=0.018128, prompt_tokens=1240, completion_tokens=412 + ), +) + CLAUDE_OPUS_5_5_TOOL_CALL: Final = CostTrackingTestCase( name="claude-opus-5-5-tool_call", covers="quota_management.spend_tracking.cost_matrix.logs_cost", @@ -6895,6 +7437,14 @@ CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON: Final = CLAUDE_FABLE_5_1_MESSAGES_INPUT } ) +CLAUDE_FABLE_5_2_RESPONSES_BRIDGED_JSON: Final = CLAUDE_FABLE_5_2_MESSAGES_INPUT_TEXT.model_copy( + update={ + "name": "claude-fable-5-2-responses_bridged_json", + "endpoint": "/v1/responses", + "request": {"model": "$MODEL", "input": "summarize the attached material in one line"}, + } +) + CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_JSON: Final = CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT.model_copy( update={ "name": "claude-opus-5-5-responses_bridged_json", @@ -6903,6 +7453,7 @@ CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_JSON: Final = CLAUDE_OPUS_5_5_MESSAGES_INPUT_T } ) + def _responses_bridged_stream_case( sibling: CostTrackingTestCase, stream_sibling: CostTrackingTestCase, name: str ) -> CostTrackingTestCase: @@ -6916,19 +7467,29 @@ def _responses_bridged_stream_case( CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_STREAM: Final = _responses_bridged_stream_case( - CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON, CLAUDE_FABLE_5_1_MESSAGES_STREAM, + CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON, + CLAUDE_FABLE_5_1_MESSAGES_STREAM, "claude-fable-5-1-responses_bridged_stream", ) + +CLAUDE_FABLE_5_2_RESPONSES_BRIDGED_STREAM: Final = _responses_bridged_stream_case( + CLAUDE_FABLE_5_2_RESPONSES_BRIDGED_JSON, + CLAUDE_FABLE_5_2_MESSAGES_STREAM, + "claude-fable-5-2-responses_bridged_stream", +) CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_STREAM: Final = _responses_bridged_stream_case( - CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_JSON, CLAUDE_HAIKU_5_5_MESSAGES_STREAM, + CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_JSON, + CLAUDE_HAIKU_5_5_MESSAGES_STREAM, "claude-haiku-5-5-responses_bridged_stream", ) CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_STREAM: Final = _responses_bridged_stream_case( - CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_JSON, CLAUDE_OPUS_5_5_MESSAGES_STREAM, + CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_JSON, + CLAUDE_OPUS_5_5_MESSAGES_STREAM, "claude-opus-5-5-responses_bridged_stream", ) CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_STREAM: Final = _responses_bridged_stream_case( - CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_JSON, CLAUDE_SONNET_5_5_MESSAGES_STREAM, + CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_JSON, + CLAUDE_SONNET_5_5_MESSAGES_STREAM, "claude-sonnet-5-5-responses_bridged_stream", ) @@ -7028,11 +7589,13 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_JSON, CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_JSON, CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_STREAM, + CLAUDE_FABLE_5_2_RESPONSES_BRIDGED_STREAM, CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_STREAM, CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_STREAM, CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_STREAM, CLAUDE_FABLE_5_RESPONSES_BRIDGED_JSON, CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON, + CLAUDE_FABLE_5_2_RESPONSES_BRIDGED_JSON, CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_JSON, CLAUDE_HAIKU_4_5_INPUT_TEXT, CLAUDE_HAIKU_4_5_CACHE_READ, @@ -7129,12 +7692,19 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( CLAUDE_FABLE_5_CACHE_WRITE_1H, CLAUDE_FABLE_5_MESSAGES_INPUT_TEXT, CLAUDE_FABLE_5_1_INPUT_TEXT, + CLAUDE_FABLE_5_2_INPUT_TEXT, CLAUDE_FABLE_5_1_STREAM, + CLAUDE_FABLE_5_2_STREAM, CLAUDE_FABLE_5_1_CACHE_READ, + CLAUDE_FABLE_5_2_CACHE_READ, CLAUDE_FABLE_5_1_CACHE_WRITE_5M, + CLAUDE_FABLE_5_2_CACHE_WRITE_5M, CLAUDE_FABLE_5_1_CACHE_WRITE_1H, + CLAUDE_FABLE_5_2_CACHE_WRITE_1H, CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, + CLAUDE_FABLE_5_2_MESSAGES_INPUT_TEXT, CLAUDE_FABLE_5_1_MESSAGES_STREAM, + CLAUDE_FABLE_5_2_MESSAGES_STREAM, CLAUDE_OPUS_5_5_STREAM, CLAUDE_OPUS_5_5_CACHE_READ, CLAUDE_OPUS_5_5_CACHE_WRITE_5M, @@ -7148,7 +7718,9 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( CLAUDE_FABLE_5_TOOL_CALL, CLAUDE_FABLE_5_THINKING, CLAUDE_FABLE_5_1_TOOL_CALL, + CLAUDE_FABLE_5_2_TOOL_CALL, CLAUDE_FABLE_5_1_THINKING, + CLAUDE_FABLE_5_2_THINKING, CLAUDE_OPUS_5_5_TOOL_CALL, CLAUDE_OPUS_5_5_THINKING, CLAUDE_OPUS_5_5_UPSTREAM_404_ZERO_SPEND, @@ -7159,6 +7731,9 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( StreamParityTestCase( plain=CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON, streamed=CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_STREAM ), + StreamParityTestCase( + plain=CLAUDE_FABLE_5_2_RESPONSES_BRIDGED_JSON, streamed=CLAUDE_FABLE_5_2_RESPONSES_BRIDGED_STREAM + ), StreamParityTestCase( plain=CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_JSON, streamed=CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_STREAM ), @@ -7190,11 +7765,13 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( StreamParityTestCase(plain=CLAUDE_HAIKU_5_5_INPUT_TEXT, streamed=CLAUDE_HAIKU_5_5_STREAM), StreamParityTestCase(plain=CLAUDE_FABLE_5_INPUT_TEXT, streamed=CLAUDE_FABLE_5_STREAM), StreamParityTestCase(plain=CLAUDE_FABLE_5_1_INPUT_TEXT, streamed=CLAUDE_FABLE_5_1_STREAM), + StreamParityTestCase(plain=CLAUDE_FABLE_5_2_INPUT_TEXT, streamed=CLAUDE_FABLE_5_2_STREAM), StreamParityTestCase(plain=CLAUDE_OPUS_5_5_INPUT_TEXT, streamed=CLAUDE_OPUS_5_5_STREAM), StreamParityTestCase(plain=CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, streamed=CLAUDE_OPUS_5_5_MESSAGES_STREAM), StreamParityTestCase(plain=CLAUDE_SONNET_5_5_MESSAGES_INPUT_TEXT, streamed=CLAUDE_SONNET_5_5_MESSAGES_STREAM), StreamParityTestCase(plain=CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT, streamed=CLAUDE_HAIKU_5_5_MESSAGES_STREAM), StreamParityTestCase(plain=CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, streamed=CLAUDE_FABLE_5_1_MESSAGES_STREAM), + StreamParityTestCase(plain=CLAUDE_FABLE_5_2_MESSAGES_INPUT_TEXT, streamed=CLAUDE_FABLE_5_2_MESSAGES_STREAM), ) BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( @@ -7220,6 +7797,9 @@ BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( BridgeCostParityTestCase( native=CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, bridged=CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON ), + BridgeCostParityTestCase( + native=CLAUDE_FABLE_5_2_MESSAGES_INPUT_TEXT, bridged=CLAUDE_FABLE_5_2_RESPONSES_BRIDGED_JSON + ), BridgeCostParityTestCase( native=CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, bridged=CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_JSON ), diff --git a/tests/integration/cost_calculation/cases/batch_realtime.py b/tests/integration/cost_calculation/cases/batch_realtime.py index 0284036f31f..52e2021dc7d 100644 --- a/tests/integration/cost_calculation/cases/batch_realtime.py +++ b/tests/integration/cost_calculation/cases/batch_realtime.py @@ -46,8 +46,8 @@ GPT_5_6_BATCH_CACHED_INPUT_HALVED: Final = BatchCostCase( ), ) -GPT_5_4_BATCH_EXPLICIT_BATCH_RATES_BILL_CACHED_AT_BATCH_INPUT_RATE: Final = BatchCostCase( - name="gpt-5.4-batch-explicit_batch_rates_bill_cached_at_batch_input_rate", +GPT_5_4_BATCH_NO_BATCH_CACHE_RATE_BILLS_CACHED_AT_STANDARD_CACHE_RATE: Final = BatchCostCase( + name="gpt-5.4-batch-no_batch_cache_rate_bills_cached_at_standard_cache_rate", covers="quota_management.spend_tracking.batch_costs.explicit_rates", model="gpt-5.4", litellm_model="openai/gpt-5.4", @@ -56,8 +56,8 @@ GPT_5_4_BATCH_EXPLICIT_BATCH_RATES_BILL_CACHED_AT_BATCH_INPUT_RATE: Final = Batc BatchOutputLine(status_code=200, prompt_tokens=120, completion_tokens=30), ), expected=ExactExpected( - spend=0.000875, - input_cost=0.000275, + spend=0.000835, + input_cost=0.000235, output_cost=0.0006, prompt_tokens=220, completion_tokens=80, @@ -188,7 +188,7 @@ GPT_REALTIME_MINI_2025_12_15_REALTIME_SESSION_WITHOUT_TURNS_ZERO_SPEND: Final = BATCH_CASES: Final[tuple[BatchCostCase, ...]] = ( GPT_5_6_BATCH_HALVED_RATES_WHEN_MAP_HAS_NO_BATCH_KEYS, GPT_5_6_BATCH_CACHED_INPUT_HALVED, - GPT_5_4_BATCH_EXPLICIT_BATCH_RATES_BILL_CACHED_AT_BATCH_INPUT_RATE, + GPT_5_4_BATCH_NO_BATCH_CACHE_RATE_BILLS_CACHED_AT_STANDARD_CACHE_RATE, GPT_5_6_BATCH_ALL_REQUESTS_FAILED_ZERO_SPEND, ) diff --git a/tests/integration/cost_calculation/cases/bedrock_converse.py b/tests/integration/cost_calculation/cases/bedrock_converse.py index 858db06cafc..96033a8b49a 100644 --- a/tests/integration/cost_calculation/cases/bedrock_converse.py +++ b/tests/integration/cost_calculation/cases/bedrock_converse.py @@ -4672,6 +4672,57 @@ US_ANTHROPIC_CLAUDE_FABLE_5_1_INPUT_TEXT: Final = CostTrackingTestCase( spend=0.0702, input_cost=0.03312, output_cost=0.03708, prompt_tokens=1840, completion_tokens=412 ), ) +US_ANTHROPIC_CLAUDE_FABLE_5_2_INPUT_TEXT: Final = CostTrackingTestCase( + name="us.anthropic.claude-fable-5-2-input_text", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-fable-5-2", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "a5b0ca93d2e1 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "metrics": {"latencyMs": 878}, + "output": {"message": {"content": [{"text": "scripted answer a5b0ca93d2e1"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": { + "cacheReadInputTokenCount": 0, + "cacheReadInputTokens": 0, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, + "inputTokens": 1840, + "outputTokens": 412, + "serverToolUsage": {}, + "totalTokens": 2252, + }, + }, + ), + expected=ExactExpected( + spend=0.05616, input_cost=0.026496, output_cost=0.029664, prompt_tokens=1840, completion_tokens=412 + ), +) US_ANTHROPIC_CLAUDE_FABLE_5_1_STREAM: Final = CostTrackingTestCase( name="us.anthropic.claude-fable-5-1-stream", @@ -4726,6 +4777,59 @@ US_ANTHROPIC_CLAUDE_FABLE_5_1_STREAM: Final = CostTrackingTestCase( spend=0.0702, input_cost=0.03312, output_cost=0.03708, prompt_tokens=1840, completion_tokens=412 ), ) +US_ANTHROPIC_CLAUDE_FABLE_5_2_STREAM: Final = CostTrackingTestCase( + name="us.anthropic.claude-fable-5-2-stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="us.anthropic.claude-fable-5-2", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "d62b383256ff summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": True, + "stream_options": {"include_usage": True}, + "allowed_openai_params": [], + }, + response=EventStreamResponse( + content_type="application/vnd.amazon.eventstream", + events=( + EventStreamEvent(event_type="messageStart", payload={"role": "assistant"}), + EventStreamEvent( + event_type="contentBlockDelta", + payload={"contentBlockIndex": 0, "delta": {"text": "scripted answer d62b383256ff"}}, + ), + EventStreamEvent(event_type="contentBlockStop", payload={"contentBlockIndex": 0}), + EventStreamEvent(event_type="messageStop", payload={"stopReason": "end_turn"}), + EventStreamEvent( + event_type="metadata", + payload={ + "metrics": {"latencyMs": 883}, + "usage": {"inputTokens": 1840, "outputTokens": 412, "serverToolUsage": {}, "totalTokens": 2252}, + }, + ), + ), + ), + expected=ExactExpected( + spend=0.05616, input_cost=0.026496, output_cost=0.029664, prompt_tokens=1840, completion_tokens=412 + ), +) US_ANTHROPIC_CLAUDE_FABLE_5_1_CACHE_READ: Final = CostTrackingTestCase( name="us.anthropic.claude-fable-5-1-cache_read", @@ -4779,6 +4883,58 @@ US_ANTHROPIC_CLAUDE_FABLE_5_1_CACHE_READ: Final = CostTrackingTestCase( spend=0.0678384, input_cost=0.0336384, output_cost=0.0342, prompt_tokens=12928, completion_tokens=380 ), ) +US_ANTHROPIC_CLAUDE_FABLE_5_2_CACHE_READ: Final = CostTrackingTestCase( + name="us.anthropic.claude-fable-5-2-cache_read", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-fable-5-2", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "1a5a397f1134 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "metrics": {"latencyMs": 878}, + "output": {"message": {"content": [{"text": "scripted answer 1a5a397f1134"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": { + "cacheReadInputTokenCount": 12288, + "cacheReadInputTokens": 12288, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, + "inputTokens": 640, + "outputTokens": 380, + "serverToolUsage": {}, + "totalTokens": 13308, + }, + }, + ), + expected=ExactExpected( + spend=0.05427072, input_cost=0.02691072, output_cost=0.02736, prompt_tokens=12928, completion_tokens=380 + ), +) US_ANTHROPIC_CLAUDE_FABLE_5_1_CACHE_WRITE_5M: Final = CostTrackingTestCase( name="us.anthropic.claude-fable-5-1-cache_write_5m", @@ -4830,6 +4986,56 @@ US_ANTHROPIC_CLAUDE_FABLE_5_1_CACHE_WRITE_5M: Final = CostTrackingTestCase( spend=0.248076, input_cost=0.216576, output_cost=0.0315, prompt_tokens=9728, completion_tokens=350 ), ) +US_ANTHROPIC_CLAUDE_FABLE_5_2_CACHE_WRITE_5M: Final = CostTrackingTestCase( + name="us.anthropic.claude-fable-5-2-cache_write_5m", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-fable-5-2", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "3481bca67e1a summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "metrics": {"latencyMs": 878}, + "output": {"message": {"content": [{"text": "scripted answer 3481bca67e1a"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": { + "inputTokens": 512, + "outputTokens": 350, + "totalTokens": 10078, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 9216, + "cacheDetails": [{"ttl": "5m", "inputTokens": 9216}], + }, + }, + ), + expected=ExactExpected( + spend=0.1984608, input_cost=0.1732608, output_cost=0.0252, prompt_tokens=9728, completion_tokens=350 + ), +) US_ANTHROPIC_CLAUDE_FABLE_5_1_CACHE_WRITE_1H: Final = CostTrackingTestCase( name="us.anthropic.claude-fable-5-1-cache_write_1h", @@ -4881,6 +5087,56 @@ US_ANTHROPIC_CLAUDE_FABLE_5_1_CACHE_WRITE_1H: Final = CostTrackingTestCase( spend=0.372492, input_cost=0.340992, output_cost=0.0315, prompt_tokens=9728, completion_tokens=350 ), ) +US_ANTHROPIC_CLAUDE_FABLE_5_2_CACHE_WRITE_1H: Final = CostTrackingTestCase( + name="us.anthropic.claude-fable-5-2-cache_write_1h", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-fable-5-2", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "3f037d5317b0 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "allowed_openai_params": [], + }, + response=JsonResponse( + content_type="application/json", + body={ + "metrics": {"latencyMs": 878}, + "output": {"message": {"content": [{"text": "scripted answer 3f037d5317b0"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": { + "inputTokens": 512, + "outputTokens": 350, + "totalTokens": 10078, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 9216, + "cacheDetails": [{"ttl": "1h", "inputTokens": 9216}], + }, + }, + ), + expected=ExactExpected( + spend=0.2979936, input_cost=0.2727936, output_cost=0.0252, prompt_tokens=9728, completion_tokens=350 + ), +) US_ANTHROPIC_CLAUDE_SONNET_5_TOOL_CALL: Final = CostTrackingTestCase( name="us.anthropic.claude-sonnet-5-tool_call", @@ -5832,6 +6088,87 @@ US_ANTHROPIC_CLAUDE_FABLE_5_1_TOOL_CALL: Final = CostTrackingTestCase( spend=0.0702, input_cost=0.03312, output_cost=0.03708, prompt_tokens=1840, completion_tokens=412 ), ) +US_ANTHROPIC_CLAUDE_FABLE_5_2_TOOL_CALL: Final = CostTrackingTestCase( + name="us.anthropic.claude-fable-5-2-tool_call", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-fable-5-2", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "7f97ab74e808 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather and a short forecast for a city.", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "City name"}, + "days": {"type": "integer", "description": "Forecast horizon in days"}, + "units": {"type": "string", "enum": ["metric", "imperial"]}, + }, + "required": ["city"], + }, + }, + } + ], + "tool_choice": "auto", + "allowed_openai_params": ["tool_choice"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "output": { + "message": { + "role": "assistant", + "content": [ + { + "toolUse": { + "toolUseId": "tooluse_$REQUEST_ID", + "name": "get_weather", + "input": {"city": "Berlin", "days": 7, "units": "metric"}, + "type": "tool_use", + } + } + ], + } + }, + "stopReason": "tool_use", + "usage": { + "inputTokens": 1840, + "outputTokens": 412, + "totalTokens": 2252, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 0, + }, + "metrics": {"latencyMs": 2071}, + }, + ), + expected=ExactExpected( + spend=0.05616, input_cost=0.026496, output_cost=0.029664, prompt_tokens=1840, completion_tokens=412 + ), +) US_ANTHROPIC_CLAUDE_FABLE_5_1_THINKING: Final = CostTrackingTestCase( name="us.anthropic.claude-fable-5-1-thinking", @@ -5885,6 +6222,58 @@ US_ANTHROPIC_CLAUDE_FABLE_5_1_THINKING: Final = CostTrackingTestCase( spend=0.0594, input_cost=0.02232, output_cost=0.03708, prompt_tokens=1240, completion_tokens=412 ), ) +US_ANTHROPIC_CLAUDE_FABLE_5_2_THINKING: Final = CostTrackingTestCase( + name="us.anthropic.claude-fable-5-2-thinking", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="us.anthropic.claude-fable-5-2", + request={ + "model": "$MODEL", + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are a deterministic pricing-harness assistant. Keep answers to a single short line.", + } + ], + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "bca700fd52a2 summarize the attached material in one line and name the city weather", + } + ], + }, + ], + "stream": False, + "reasoning_effort": "medium", + "allowed_openai_params": ["reasoning_effort"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "output": {"message": {"role": "assistant", "content": [{"text": "scripted answer bca700fd52a2"}]}}, + "stopReason": "end_turn", + "usage": { + "cacheReadInputTokenCount": 0, + "cacheReadInputTokens": 0, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, + "inputTokens": 1240, + "outputTokens": 412, + "serverToolUsage": {}, + "totalTokens": 1652, + }, + "metrics": {"latencyMs": 2071}, + }, + ), + expected=ExactExpected( + spend=0.04752, input_cost=0.017856, output_cost=0.029664, prompt_tokens=1240, completion_tokens=412 + ), +) US_ANTHROPIC_CLAUDE_OPUS_5_5_UPSTREAM_400_ZERO_SPEND: Final = CostTrackingTestCase( name="us.anthropic.claude-opus-5-5-upstream_400_zero_spend", @@ -6014,6 +6403,13 @@ US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON: Final = US_ANTHROPIC_CLAUD "request": {"model": "$MODEL", "input": "summarize the attached material in one line"}, } ) +US_ANTHROPIC_CLAUDE_FABLE_5_2_RESPONSES_BRIDGED_JSON: Final = US_ANTHROPIC_CLAUDE_FABLE_5_2_INPUT_TEXT.model_copy( + update={ + "name": "us.anthropic.claude-fable-5-2-responses_bridged_json", + "endpoint": "/v1/responses", + "request": {"model": "$MODEL", "input": "summarize the attached material in one line"}, + } +) US_ANTHROPIC_CLAUDE_HAIKU_5_5_TIERED_INPUT_ABOVE_100K: Final = with_usage( US_ANTHROPIC_CLAUDE_HAIKU_5_5_INPUT_TEXT, @@ -6088,6 +6484,7 @@ US_ANTHROPIC_CLAUDE_HAIKU_5_5_TIERED_CACHE_WRITE_1H_ABOVE_100K: Final = with_usa ), ) + def _us_anthropic_messages_bridge_pairs( input_case: CostTrackingTestCase, stream_case: CostTrackingTestCase ) -> tuple[BridgeCostParityTestCase, ...]: @@ -6125,6 +6522,11 @@ US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_STREAM: Final = _us_anthropic_re US_ANTHROPIC_CLAUDE_FABLE_5_1_STREAM, "us.anthropic.claude-fable-5-1-responses_bridged_stream", ) +US_ANTHROPIC_CLAUDE_FABLE_5_2_RESPONSES_BRIDGED_STREAM: Final = _us_anthropic_responses_stream_case( + US_ANTHROPIC_CLAUDE_FABLE_5_2_RESPONSES_BRIDGED_JSON, + US_ANTHROPIC_CLAUDE_FABLE_5_2_STREAM, + "us.anthropic.claude-fable-5-2-responses_bridged_stream", +) US_ANTHROPIC_CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_STREAM: Final = _us_anthropic_responses_stream_case( US_ANTHROPIC_CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_JSON, US_ANTHROPIC_CLAUDE_HAIKU_5_5_STREAM, @@ -6143,6 +6545,9 @@ US_ANTHROPIC_CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_STREAM: Final = _us_anthropic_r _US_ANTHROPIC_FABLE_5_1_MESSAGES_PAIRS: Final = _us_anthropic_messages_bridge_pairs( US_ANTHROPIC_CLAUDE_FABLE_5_1_INPUT_TEXT, US_ANTHROPIC_CLAUDE_FABLE_5_1_STREAM ) +_US_ANTHROPIC_FABLE_5_2_MESSAGES_PAIRS: Final = _us_anthropic_messages_bridge_pairs( + US_ANTHROPIC_CLAUDE_FABLE_5_2_INPUT_TEXT, US_ANTHROPIC_CLAUDE_FABLE_5_2_STREAM +) _US_ANTHROPIC_HAIKU_5_5_MESSAGES_PAIRS: Final = _us_anthropic_messages_bridge_pairs( US_ANTHROPIC_CLAUDE_HAIKU_5_5_INPUT_TEXT, US_ANTHROPIC_CLAUDE_HAIKU_5_5_STREAM ) @@ -6161,6 +6566,7 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( US_ANTHROPIC_CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_JSON, US_ANTHROPIC_CLAUDE_FABLE_5_RESPONSES_BRIDGED_JSON, US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON, + US_ANTHROPIC_CLAUDE_FABLE_5_2_RESPONSES_BRIDGED_JSON, US_ANTHROPIC_CLAUDE_HAIKU_5_5_TIERED_INPUT_ABOVE_100K, US_ANTHROPIC_CLAUDE_HAIKU_5_5_TIERED_AT_100K_BOUNDARY, US_ANTHROPIC_CLAUDE_HAIKU_5_5_TIERED_CACHE_READ_ABOVE_100K, @@ -6255,10 +6661,15 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( US_ANTHROPIC_CLAUDE_FABLE_5_CACHE_WRITE_5M, US_ANTHROPIC_CLAUDE_FABLE_5_CACHE_WRITE_1H, US_ANTHROPIC_CLAUDE_FABLE_5_1_INPUT_TEXT, + US_ANTHROPIC_CLAUDE_FABLE_5_2_INPUT_TEXT, US_ANTHROPIC_CLAUDE_FABLE_5_1_STREAM, + US_ANTHROPIC_CLAUDE_FABLE_5_2_STREAM, US_ANTHROPIC_CLAUDE_FABLE_5_1_CACHE_READ, + US_ANTHROPIC_CLAUDE_FABLE_5_2_CACHE_READ, US_ANTHROPIC_CLAUDE_FABLE_5_1_CACHE_WRITE_5M, + US_ANTHROPIC_CLAUDE_FABLE_5_2_CACHE_WRITE_5M, US_ANTHROPIC_CLAUDE_FABLE_5_1_CACHE_WRITE_1H, + US_ANTHROPIC_CLAUDE_FABLE_5_2_CACHE_WRITE_1H, US_ANTHROPIC_CLAUDE_SONNET_5_TOOL_CALL, US_ANTHROPIC_CLAUDE_SONNET_5_THINKING, US_ANTHROPIC_CLAUDE_SONNET_5_5_TOOL_CALL, @@ -6272,14 +6683,18 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( US_ANTHROPIC_CLAUDE_FABLE_5_TOOL_CALL, US_ANTHROPIC_CLAUDE_FABLE_5_THINKING, US_ANTHROPIC_CLAUDE_FABLE_5_1_TOOL_CALL, + US_ANTHROPIC_CLAUDE_FABLE_5_2_TOOL_CALL, US_ANTHROPIC_CLAUDE_FABLE_5_1_THINKING, + US_ANTHROPIC_CLAUDE_FABLE_5_2_THINKING, US_ANTHROPIC_CLAUDE_OPUS_5_5_UPSTREAM_400_ZERO_SPEND, US_ANTHROPIC_CLAUDE_OPUS_5_5_INPUT_TEXT_ROLLUPS, US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_STREAM, + US_ANTHROPIC_CLAUDE_FABLE_5_2_RESPONSES_BRIDGED_STREAM, US_ANTHROPIC_CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_STREAM, US_ANTHROPIC_CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_STREAM, US_ANTHROPIC_CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_STREAM, *tuple(pair.bridged for pair in _US_ANTHROPIC_FABLE_5_1_MESSAGES_PAIRS), + *tuple(pair.bridged for pair in _US_ANTHROPIC_FABLE_5_2_MESSAGES_PAIRS), *tuple(pair.bridged for pair in _US_ANTHROPIC_HAIKU_5_5_MESSAGES_PAIRS), *tuple(pair.bridged for pair in _US_ANTHROPIC_OPUS_5_5_MESSAGES_PAIRS), *tuple(pair.bridged for pair in _US_ANTHROPIC_SONNET_5_5_MESSAGES_PAIRS), @@ -6308,10 +6723,15 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( StreamParityTestCase(plain=US_ANTHROPIC_CLAUDE_HAIKU_5_5_INPUT_TEXT, streamed=US_ANTHROPIC_CLAUDE_HAIKU_5_5_STREAM), StreamParityTestCase(plain=US_ANTHROPIC_CLAUDE_FABLE_5_INPUT_TEXT, streamed=US_ANTHROPIC_CLAUDE_FABLE_5_STREAM), StreamParityTestCase(plain=US_ANTHROPIC_CLAUDE_FABLE_5_1_INPUT_TEXT, streamed=US_ANTHROPIC_CLAUDE_FABLE_5_1_STREAM), + StreamParityTestCase(plain=US_ANTHROPIC_CLAUDE_FABLE_5_2_INPUT_TEXT, streamed=US_ANTHROPIC_CLAUDE_FABLE_5_2_STREAM), StreamParityTestCase( plain=US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON, streamed=US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_STREAM, ), + StreamParityTestCase( + plain=US_ANTHROPIC_CLAUDE_FABLE_5_2_RESPONSES_BRIDGED_JSON, + streamed=US_ANTHROPIC_CLAUDE_FABLE_5_2_RESPONSES_BRIDGED_STREAM, + ), StreamParityTestCase( plain=US_ANTHROPIC_CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_JSON, streamed=US_ANTHROPIC_CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_STREAM, @@ -6328,6 +6748,10 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( plain=_US_ANTHROPIC_FABLE_5_1_MESSAGES_PAIRS[0].bridged, streamed=_US_ANTHROPIC_FABLE_5_1_MESSAGES_PAIRS[1].bridged, ), + StreamParityTestCase( + plain=_US_ANTHROPIC_FABLE_5_2_MESSAGES_PAIRS[0].bridged, + streamed=_US_ANTHROPIC_FABLE_5_2_MESSAGES_PAIRS[1].bridged, + ), StreamParityTestCase( plain=_US_ANTHROPIC_HAIKU_5_5_MESSAGES_PAIRS[0].bridged, streamed=_US_ANTHROPIC_HAIKU_5_5_MESSAGES_PAIRS[1].bridged, @@ -6372,10 +6796,17 @@ BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( BridgeCostParityTestCase( native=US_ANTHROPIC_CLAUDE_FABLE_5_1_INPUT_TEXT, bridged=US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON ), + BridgeCostParityTestCase( + native=US_ANTHROPIC_CLAUDE_FABLE_5_2_INPUT_TEXT, bridged=US_ANTHROPIC_CLAUDE_FABLE_5_2_RESPONSES_BRIDGED_JSON + ), BridgeCostParityTestCase( native=US_ANTHROPIC_CLAUDE_FABLE_5_1_STREAM, bridged=US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_STREAM, ), + BridgeCostParityTestCase( + native=US_ANTHROPIC_CLAUDE_FABLE_5_2_STREAM, + bridged=US_ANTHROPIC_CLAUDE_FABLE_5_2_RESPONSES_BRIDGED_STREAM, + ), BridgeCostParityTestCase( native=US_ANTHROPIC_CLAUDE_HAIKU_5_5_STREAM, bridged=US_ANTHROPIC_CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_STREAM, @@ -6389,6 +6820,7 @@ BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( bridged=US_ANTHROPIC_CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_STREAM, ), *_US_ANTHROPIC_FABLE_5_1_MESSAGES_PAIRS, + *_US_ANTHROPIC_FABLE_5_2_MESSAGES_PAIRS, *_US_ANTHROPIC_HAIKU_5_5_MESSAGES_PAIRS, *_US_ANTHROPIC_OPUS_5_5_MESSAGES_PAIRS, *_US_ANTHROPIC_SONNET_5_5_MESSAGES_PAIRS, diff --git a/tests/integration/cost_calculation/cases/vertex_ai.py b/tests/integration/cost_calculation/cases/vertex_ai.py index 8590a19646b..47a484e972b 100644 --- a/tests/integration/cost_calculation/cases/vertex_ai.py +++ b/tests/integration/cost_calculation/cases/vertex_ai.py @@ -5,8 +5,7 @@ PARITY pairs the streamed case with its plain twin where both bill the same row. import json from typing import Final, Literal, cast -from pydantic import JsonValue - +from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParityTestCase, bridge_pairs from integration.cost_calculation.cases.anthropic import ( CLAUDE_FABLE_5_1_CACHE_READ, CLAUDE_FABLE_5_1_CACHE_WRITE_1H, @@ -17,6 +16,15 @@ from integration.cost_calculation.cases.anthropic import ( CLAUDE_FABLE_5_1_STREAM, CLAUDE_FABLE_5_1_THINKING, CLAUDE_FABLE_5_1_TOOL_CALL, + CLAUDE_FABLE_5_2_CACHE_READ, + CLAUDE_FABLE_5_2_CACHE_WRITE_1H, + CLAUDE_FABLE_5_2_CACHE_WRITE_5M, + CLAUDE_FABLE_5_2_INPUT_TEXT, + CLAUDE_FABLE_5_2_MESSAGES_INPUT_TEXT, + CLAUDE_FABLE_5_2_MESSAGES_STREAM, + CLAUDE_FABLE_5_2_STREAM, + CLAUDE_FABLE_5_2_THINKING, + CLAUDE_FABLE_5_2_TOOL_CALL, CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT, CLAUDE_HAIKU_5_5_MESSAGES_STREAM, CLAUDE_OPUS_5_5_CACHE_READ, @@ -38,7 +46,6 @@ from integration.cost_calculation.cases.anthropic import ( CLAUDE_SONNET_5_5_THINKING, CLAUDE_SONNET_5_5_TOOL_CALL, ) -from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParityTestCase, bridge_pairs from integration.cost_calculation.cost_tracking_case import ( CostTrackingTestCase, Deployment, @@ -49,6 +56,7 @@ from integration.cost_calculation.cost_tracking_case import ( SseResponse, ) from integration.cost_calculation.stream_parity.case import StreamParityTestCase, sse_frames +from pydantic import JsonValue def _gemini_case_response( @@ -3009,6 +3017,37 @@ VERTEX_AI_CLAUDE_FABLE_5_1_INPUT_TEXT: Final = CLAUDE_FABLE_5_1_INPUT_TEXT.model ), } ) +VERTEX_AI_CLAUDE_FABLE_5_2_INPUT_TEXT: Final = CLAUDE_FABLE_5_2_INPUT_TEXT.model_copy( + update={ + "name": "vertex_ai/claude-fable-5-2-input_text", + "model": "vertex_ai/claude-fable-5-2", + "response": JsonResponse( + content_type="application/json", + body={ + "model": "claude-fable-5-2", + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "content": [{"type": "thinking", "thinking": "scripted thought", "signature": "scripted signature"}], + "container": None, + "stop_reason": "max_tokens", + "stop_sequence": None, + "stop_details": None, + "usage": { + "input_tokens": 1840, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 412, + "output_tokens_details": {"thinking_tokens": 0}, + }, + }, + ), + "expected": ExactExpected( + spend=0.05928, input_cost=0.027968, output_cost=0.031312, prompt_tokens=1840, completion_tokens=412 + ), + } +) VERTEX_AI_CLAUDE_FABLE_5_1_STREAM: Final = CLAUDE_FABLE_5_1_STREAM.model_copy( update={ "name": "vertex_ai/claude-fable-5-1-stream", @@ -3017,6 +3056,14 @@ VERTEX_AI_CLAUDE_FABLE_5_1_STREAM: Final = CLAUDE_FABLE_5_1_STREAM.model_copy( "expected": VERTEX_AI_CLAUDE_FABLE_5_1_INPUT_TEXT.expected, } ) +VERTEX_AI_CLAUDE_FABLE_5_2_STREAM: Final = CLAUDE_FABLE_5_2_STREAM.model_copy( + update={ + "name": "vertex_ai/claude-fable-5-2-stream", + "model": "vertex_ai/claude-fable-5-2", + "response": _vertex_claude_stream_response("claude-fable-5-2", "max_tokens", ()), + "expected": VERTEX_AI_CLAUDE_FABLE_5_2_INPUT_TEXT.expected, + } +) VERTEX_AI_CLAUDE_SONNET_5_5_CACHE_READ: Final = _vertex_claude_shape_case( CLAUDE_SONNET_5_5_CACHE_READ, @@ -3135,6 +3182,15 @@ VERTEX_AI_CLAUDE_FABLE_5_1_CACHE_READ: Final = _vertex_claude_shape_case( spend=0.013567425, input_cost=0.013092425, output_cost=0.000475, prompt_tokens=27017, completion_tokens=5 ), ) +VERTEX_AI_CLAUDE_FABLE_5_2_CACHE_READ: Final = _vertex_claude_shape_case( + CLAUDE_FABLE_5_2_CACHE_READ, + "vertex_ai/claude-fable-5-2-cache_read", + "vertex_ai/claude-fable-5-2", + _vertex_claude_usage(input_tokens=14, output_tokens=5, cache_read_input_tokens=27003), + ExactExpected( + spend=0.01085394, input_cost=0.01047394, output_cost=0.00038, prompt_tokens=27017, completion_tokens=5 + ), +) VERTEX_AI_CLAUDE_FABLE_5_1_CACHE_WRITE_5M: Final = _vertex_claude_shape_case( CLAUDE_FABLE_5_1_CACHE_WRITE_5M, "vertex_ai/claude-fable-5-1-cache_write_5m", @@ -3146,6 +3202,17 @@ VERTEX_AI_CLAUDE_FABLE_5_1_CACHE_WRITE_5M: Final = _vertex_claude_shape_case( spend=0.64225225, input_cost=0.64177725, output_cost=0.000475, prompt_tokens=27025, completion_tokens=5 ), ) +VERTEX_AI_CLAUDE_FABLE_5_2_CACHE_WRITE_5M: Final = _vertex_claude_shape_case( + CLAUDE_FABLE_5_2_CACHE_WRITE_5M, + "vertex_ai/claude-fable-5-2-cache_write_5m", + "vertex_ai/claude-fable-5-2", + _vertex_claude_usage( + input_tokens=14, output_tokens=5, cache_creation_input_tokens=27011, ephemeral_5m_input_tokens=27011 + ), + ExactExpected( + spend=0.5138018, input_cost=0.5134218, output_cost=0.00038, prompt_tokens=27025, completion_tokens=5 + ), +) VERTEX_AI_CLAUDE_FABLE_5_1_CACHE_WRITE_1H: Final = _vertex_claude_shape_case( CLAUDE_FABLE_5_1_CACHE_WRITE_1H, "vertex_ai/claude-fable-5-1-cache_write_1h", @@ -3157,6 +3224,17 @@ VERTEX_AI_CLAUDE_FABLE_5_1_CACHE_WRITE_1H: Final = _vertex_claude_shape_case( spend=1.027159, input_cost=1.026684, output_cost=0.000475, prompt_tokens=27025, completion_tokens=5 ), ) +VERTEX_AI_CLAUDE_FABLE_5_2_CACHE_WRITE_1H: Final = _vertex_claude_shape_case( + CLAUDE_FABLE_5_2_CACHE_WRITE_1H, + "vertex_ai/claude-fable-5-2-cache_write_1h", + "vertex_ai/claude-fable-5-2", + _vertex_claude_usage( + input_tokens=14, output_tokens=5, cache_creation_input_tokens=27011, ephemeral_1h_input_tokens=27011 + ), + ExactExpected( + spend=0.8217272, input_cost=0.8213472, output_cost=0.00038, prompt_tokens=27025, completion_tokens=5 + ), +) VERTEX_AI_CLAUDE_FABLE_5_1_THINKING: Final = _vertex_claude_shape_case( CLAUDE_FABLE_5_1_THINKING, "vertex_ai/claude-fable-5-1-thinking", @@ -3171,6 +3249,20 @@ VERTEX_AI_CLAUDE_FABLE_5_1_THINKING: Final = _vertex_claude_shape_case( ], adaptive_thinking=True, ) +VERTEX_AI_CLAUDE_FABLE_5_2_THINKING: Final = _vertex_claude_shape_case( + CLAUDE_FABLE_5_2_THINKING, + "vertex_ai/claude-fable-5-2-thinking", + "vertex_ai/claude-fable-5-2", + _vertex_claude_usage(input_tokens=18, output_tokens=98, thinking_tokens=33), + ExactExpected( + spend=0.0077216, input_cost=0.0002736, output_cost=0.007448, prompt_tokens=18, completion_tokens=98 + ), + [ + {"type": "thinking", "thinking": "scripted thought", "signature": "scripted signature"}, + {"type": "text", "text": "scripted answer for cost fixtures"}, + ], + adaptive_thinking=True, +) VERTEX_AI_CLAUDE_FABLE_5_1_TOOL_CALL: Final = _vertex_claude_shape_case( CLAUDE_FABLE_5_1_TOOL_CALL, "vertex_ai/claude-fable-5-1-tool_call", @@ -3180,6 +3272,15 @@ VERTEX_AI_CLAUDE_FABLE_5_1_TOOL_CALL: Final = _vertex_claude_shape_case( spend=0.012065, input_cost=0.007315, output_cost=0.00475, prompt_tokens=385, completion_tokens=50 ), ) +VERTEX_AI_CLAUDE_FABLE_5_2_TOOL_CALL: Final = _vertex_claude_shape_case( + CLAUDE_FABLE_5_2_TOOL_CALL, + "vertex_ai/claude-fable-5-2-tool_call", + "vertex_ai/claude-fable-5-2", + _vertex_claude_usage(input_tokens=385, output_tokens=50), + ExactExpected( + spend=0.009652, input_cost=0.005852, output_cost=0.0038, prompt_tokens=385, completion_tokens=50 + ), +) VERTEX_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT: Final = _vertex_claude_shape_case( VERTEX_AI_CLAUDE_SONNET_5_5_INPUT_TEXT, @@ -3325,6 +3426,15 @@ VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT: Final = CLAUDE_FABLE_5_1_MESSAGE "expected": VERTEX_AI_CLAUDE_FABLE_5_1_INPUT_TEXT.expected, } ) +VERTEX_AI_CLAUDE_FABLE_5_2_MESSAGES_INPUT_TEXT: Final = CLAUDE_FABLE_5_2_MESSAGES_INPUT_TEXT.model_copy( + update={ + "name": "vertex_ai/claude-fable-5-2-messages_input_text", + "model": "vertex_ai/claude-fable-5-2", + "endpoint": "/v1/messages", + "response": VERTEX_AI_CLAUDE_FABLE_5_2_INPUT_TEXT.response, + "expected": VERTEX_AI_CLAUDE_FABLE_5_2_INPUT_TEXT.expected, + } +) VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_STREAM: Final = CLAUDE_FABLE_5_1_MESSAGES_STREAM.model_copy( update={ "name": "vertex_ai/claude-fable-5-1-messages_stream", @@ -3334,10 +3444,23 @@ VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_STREAM: Final = CLAUDE_FABLE_5_1_MESSAGES_ST "expected": VERTEX_AI_CLAUDE_FABLE_5_1_INPUT_TEXT.expected, } ) +VERTEX_AI_CLAUDE_FABLE_5_2_MESSAGES_STREAM: Final = CLAUDE_FABLE_5_2_MESSAGES_STREAM.model_copy( + update={ + "name": "vertex_ai/claude-fable-5-2-messages_stream", + "model": "vertex_ai/claude-fable-5-2", + "endpoint": "/v1/messages", + "response": VERTEX_AI_CLAUDE_FABLE_5_2_STREAM.response, + "expected": VERTEX_AI_CLAUDE_FABLE_5_2_INPUT_TEXT.expected, + } +) VERTEX_AI_CLAUDE_FABLE_5_1_BRIDGE_PAIRS: Final = bridge_pairs( (VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_STREAM), ("/v1/responses",), ) +VERTEX_AI_CLAUDE_FABLE_5_2_BRIDGE_PAIRS: Final = bridge_pairs( + (VERTEX_AI_CLAUDE_FABLE_5_2_MESSAGES_INPUT_TEXT, VERTEX_AI_CLAUDE_FABLE_5_2_MESSAGES_STREAM), + ("/v1/responses",), +) VERTEX_AI_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT: Final = CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT.model_copy( update={ "name": "vertex_ai/claude-opus-5-5-messages_input_text", @@ -3915,7 +4038,9 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( VERTEX_AI_CLAUDE_OPUS_5_5_INPUT_TEXT, VERTEX_AI_CLAUDE_OPUS_5_5_STREAM, VERTEX_AI_CLAUDE_FABLE_5_1_INPUT_TEXT, + VERTEX_AI_CLAUDE_FABLE_5_2_INPUT_TEXT, VERTEX_AI_CLAUDE_FABLE_5_1_STREAM, + VERTEX_AI_CLAUDE_FABLE_5_2_STREAM, VERTEX_AI_CLAUDE_SONNET_5_5_CACHE_READ, VERTEX_AI_CLAUDE_SONNET_5_5_CACHE_WRITE_5M, VERTEX_AI_CLAUDE_SONNET_5_5_CACHE_WRITE_1H, @@ -3927,10 +4052,15 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( VERTEX_AI_CLAUDE_OPUS_5_5_THINKING, VERTEX_AI_CLAUDE_OPUS_5_5_TOOL_CALL, VERTEX_AI_CLAUDE_FABLE_5_1_CACHE_READ, + VERTEX_AI_CLAUDE_FABLE_5_2_CACHE_READ, VERTEX_AI_CLAUDE_FABLE_5_1_CACHE_WRITE_5M, + VERTEX_AI_CLAUDE_FABLE_5_2_CACHE_WRITE_5M, VERTEX_AI_CLAUDE_FABLE_5_1_CACHE_WRITE_1H, + VERTEX_AI_CLAUDE_FABLE_5_2_CACHE_WRITE_1H, VERTEX_AI_CLAUDE_FABLE_5_1_THINKING, + VERTEX_AI_CLAUDE_FABLE_5_2_THINKING, VERTEX_AI_CLAUDE_FABLE_5_1_TOOL_CALL, + VERTEX_AI_CLAUDE_FABLE_5_2_TOOL_CALL, VERTEX_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT, VERTEX_AI_CLAUDE_HAIKU_5_5_STREAM, VERTEX_AI_CLAUDE_HAIKU_5_5_CACHE_READ, @@ -3944,8 +4074,11 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( VERTEX_AI_CLAUDE_HAIKU_5_5_MESSAGES_STREAM, *(pair.bridged for pair in VERTEX_AI_CLAUDE_HAIKU_5_5_BRIDGE_PAIRS), VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, + VERTEX_AI_CLAUDE_FABLE_5_2_MESSAGES_INPUT_TEXT, VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_STREAM, + VERTEX_AI_CLAUDE_FABLE_5_2_MESSAGES_STREAM, *(pair.bridged for pair in VERTEX_AI_CLAUDE_FABLE_5_1_BRIDGE_PAIRS), + *(pair.bridged for pair in VERTEX_AI_CLAUDE_FABLE_5_2_BRIDGE_PAIRS), VERTEX_AI_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, VERTEX_AI_CLAUDE_OPUS_5_5_MESSAGES_STREAM, *(pair.bridged for pair in VERTEX_AI_CLAUDE_OPUS_5_5_BRIDGE_PAIRS), @@ -3995,6 +4128,7 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( ), StreamParityTestCase(plain=VERTEX_AI_CLAUDE_OPUS_5_5_INPUT_TEXT, streamed=VERTEX_AI_CLAUDE_OPUS_5_5_STREAM), StreamParityTestCase(plain=VERTEX_AI_CLAUDE_FABLE_5_1_INPUT_TEXT, streamed=VERTEX_AI_CLAUDE_FABLE_5_1_STREAM), + StreamParityTestCase(plain=VERTEX_AI_CLAUDE_FABLE_5_2_INPUT_TEXT, streamed=VERTEX_AI_CLAUDE_FABLE_5_2_STREAM), StreamParityTestCase( plain=VERTEX_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT, streamed=VERTEX_AI_CLAUDE_HAIKU_5_5_STREAM ), @@ -4010,10 +4144,18 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( plain=VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, streamed=VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_STREAM, ), + StreamParityTestCase( + plain=VERTEX_AI_CLAUDE_FABLE_5_2_MESSAGES_INPUT_TEXT, + streamed=VERTEX_AI_CLAUDE_FABLE_5_2_MESSAGES_STREAM, + ), StreamParityTestCase( plain=VERTEX_AI_CLAUDE_FABLE_5_1_BRIDGE_PAIRS[0].bridged, streamed=VERTEX_AI_CLAUDE_FABLE_5_1_BRIDGE_PAIRS[1].bridged, ), + StreamParityTestCase( + plain=VERTEX_AI_CLAUDE_FABLE_5_2_BRIDGE_PAIRS[0].bridged, + streamed=VERTEX_AI_CLAUDE_FABLE_5_2_BRIDGE_PAIRS[1].bridged, + ), StreamParityTestCase( plain=VERTEX_AI_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, streamed=VERTEX_AI_CLAUDE_OPUS_5_5_MESSAGES_STREAM, @@ -4068,6 +4210,7 @@ BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( *GEMINI_4_ARGON_ROUTE_BRIDGES, *VERTEX_AI_CLAUDE_HAIKU_5_5_BRIDGE_PAIRS, *VERTEX_AI_CLAUDE_FABLE_5_1_BRIDGE_PAIRS, + *VERTEX_AI_CLAUDE_FABLE_5_2_BRIDGE_PAIRS, *VERTEX_AI_CLAUDE_OPUS_5_5_BRIDGE_PAIRS, *VERTEX_AI_CLAUDE_SONNET_5_5_BRIDGE_PAIRS, ) diff --git a/tests/integration/cost_calculation/cost_map.py b/tests/integration/cost_calculation/cost_map.py index ffdb6752bf2..4f87caf7854 100644 --- a/tests/integration/cost_calculation/cost_map.py +++ b/tests/integration/cost_calculation/cost_map.py @@ -906,6 +906,28 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "max_output_tokens": 128000, "max_tokens": 128000, }, + "vertex_ai/claude-fable-5-2": { + "litellm_provider": "vertex_ai-anthropic_models", + "mode": "chat", + "input_cost_per_token": 1.52e-05, + "input_cost_per_token_batches": 7.6e-06, + "output_cost_per_token": 7.6e-05, + "output_cost_per_token_batches": 3.8e-05, + "cache_creation_input_token_cost": 1.9e-05, + "cache_creation_input_token_cost_above_1hr": 3.04e-05, + "cache_read_input_token_cost": 3.8e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.019, + "search_context_size_low": 0.019, + "search_context_size_medium": 0.019, + }, + "supports_function_calling": True, + "supports_prompt_caching": True, + "supports_reasoning": True, + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + }, "vertex_ai/claude-haiku-5-5": { "litellm_provider": "vertex_ai-anthropic_models", "mode": "chat", @@ -1398,6 +1420,19 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "output_cost_per_token": 8e-05, "supports_function_calling": True, }, + "claude-fable-5-2": { + "cache_creation_input_token_cost": 1.1e-05, + "cache_creation_input_token_cost_above_1hr": 1.76e-05, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 8.8e-06, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-05, + "supports_function_calling": True, + }, "us.anthropic.claude-sonnet-5": { "cache_creation_input_token_cost": 4.5e-06, "cache_creation_input_token_cost_above_1hr": 7.2e-06, @@ -1481,6 +1516,19 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "output_cost_per_token": 9e-05, "supports_function_calling": True, }, + "us.anthropic.claude-fable-5-2": { + "cache_creation_input_token_cost": 1.8e-05, + "cache_creation_input_token_cost_above_1hr": 2.88e-05, + "cache_read_input_token_cost": 1.44e-06, + "input_cost_per_token": 1.44e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 7.2e-05, + "supports_function_calling": True, + }, "gpt-6-sol": { "cache_read_input_token_cost": 2.2e-07, "input_cost_per_token": 2.2e-06, diff --git a/tests/integration/database/test_harness_lifecycle.py b/tests/integration/database/test_harness_lifecycle.py index d600a92c5a2..4cb01dfeef6 100644 --- a/tests/integration/database/test_harness_lifecycle.py +++ b/tests/integration/database/test_harness_lifecycle.py @@ -219,6 +219,49 @@ async def _unused_peer(reader: asyncio.StreamReader, writer: asyncio.StreamWrite await reader.read() +async def _echo_protocol(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + while chunk := await reader.read(65536): + writer.write(chunk) + await writer.drain() + + +def _protocol_message(kind: bytes, payload: bytes) -> bytes: + return kind + (len(payload) + 4).to_bytes(4, "big") + payload + + +@pytest.mark.parametrize("prepared_trigger", (False, True)) +def test_relay_holds_cached_statements_parsed_before_arming(prepared_trigger: bool) -> None: + upstream_port: Final = allocate_port() + upstream: Final = LoopbackServer(upstream_port, _echo_protocol) + relay: Final = HeldStatementRelay("127.0.0.1", upstream_port, b"trigger", after=b"snapshot", armed=False) + startup: Final = b"\x00\x00\x00\x08\x00\x03\x00\x00" + snapshot: Final = _protocol_message(b"P", b"s1\x00SELECT snapshot\x00\x00\x00") + trigger: Final = _protocol_message(b"P", b"s2\x00SELECT trigger\x00\x00\x00") + bind_snapshot: Final = _protocol_message(b"B", b"\x00s1\x00\x00\x00\x00\x00\x00\x00") + bind_trigger: Final = _protocol_message(b"B", b"\x00s2\x00\x00\x00\x00\x00\x00\x00") + upstream.start() + try: + relay.start() + with ( + socket.create_connection(("127.0.0.1", relay.port), timeout=5) as client, + client.makefile("rb") as replies, + ): + before: Final = startup + snapshot + bind_snapshot + (trigger if prepared_trigger else b"") + client.sendall(before) + assert replies.read(len(before)) == before + relay.arm() + client.sendall(bind_snapshot) + assert replies.read(len(bind_snapshot)) == bind_snapshot + held: Final = bind_trigger if prepared_trigger else _protocol_message(b"Q", b"SELECT trigger\x00") + client.sendall(held) + assert relay.held.wait(timeout=5), "cached snapshot did not open the relay barrier" + relay.release() + assert replies.read(len(held)) == held + finally: + relay.stop() + upstream.stop() + + def test_relay_bind_failure_closes_its_loop_and_allows_repeated_cleanup() -> None: port: Final = allocate_port() with socket.socket() as occupied: diff --git a/tests/integration/messages_endpoint/test_messages_additional_drop_params_wire.py b/tests/integration/messages_endpoint/test_messages_additional_drop_params_wire.py new file mode 100644 index 00000000000..f19976bda74 --- /dev/null +++ b/tests/integration/messages_endpoint/test_messages_additional_drop_params_wire.py @@ -0,0 +1,119 @@ +import json +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final + +import pytest +from integration._support.client import Gateway, list_value, object_value +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +_JSON: Final = TypeAdapter(dict[str, JsonValue]) +_ANSWER: Final = "sampling params dropped" +_SAMPLING_SPELLINGS: Final = ("top_k", "top_p", "topK", "topP") +_AWS: Final[dict[str, JsonValue]] = { + "aws_access_key_id": "AKIASCRIPTEDPROVIDER", + "aws_secret_access_key": "scripted-secret", + "aws_region_name": "us-east-1", +} + + +@dataclass(frozen=True, slots=True) +class Backend: + model: str + path_suffix: str + credentials: Mapping[str, JsonValue] + + +_BEDROCK_INVOKE: Final = Backend("bedrock/invoke/us.anthropic.claude-opus-4-8", "/invoke", _AWS) +_ANTHROPIC: Final = Backend("anthropic/claude-opus-4-8", "/v1/messages", {}) +_BEDROCK_CONVERSE: Final = Backend("bedrock/converse/us.anthropic.claude-opus-4-8", "/converse", _AWS) +_NATIVE_BACKENDS: Final = ( + pytest.param(_BEDROCK_INVOKE, id="bedrock_invoke"), + pytest.param(_ANTHROPIC, id="anthropic"), +) + + +def _answer(request: Request) -> Reply: + if request.target.endswith("/converse"): + return Reply( + body=json.dumps( + { + "output": {"message": {"role": "assistant", "content": [{"text": _ANSWER}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 5, "outputTokens": 3, "totalTokens": 8}, + "metrics": {"latencyMs": 1}, + } + ).encode() + ) + return Reply( + body=json.dumps( + { + "id": "msg_dropped_sampling_params", + "type": "message", + "role": "assistant", + "model": "claude-opus-4-8", + "content": [{"type": "text", "text": _ANSWER}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 3}, + } + ).encode() + ) + + +def _messages(gateway: Gateway, model: str, **sampling: JsonValue) -> tuple[int, str]: + response: Final = gateway.request( + "POST", + "/v1/messages", + {"model": model, "max_tokens": 32, "messages": [{"role": "user", "content": "hi"}], **sampling}, + ) + return response.status_code, response.text + + +def _answered(status: int, text: str) -> None: + assert status == 200, text + content: Final = list_value(_JSON.validate_json(text)["content"]) + assert object_value(content[0])["text"] == _ANSWER, text + + +def _sent_without_sampling_params(wire: Wire, backend: Backend) -> None: + (request,) = wire.drain() + assert request.target.endswith(backend.path_suffix), request.target + sent: Final = request.body.decode() + assert not any(spelling in sent for spelling in _SAMPLING_SPELLINGS), sent + + +@pytest.mark.parametrize( + "backend", (*_NATIVE_BACKENDS, pytest.param(_BEDROCK_CONVERSE, id="bedrock_converse_chat_bridge")) +) +def test_messages_drops_the_sampling_params_a_deployment_lists_on_a_model_without_them( + gateway: Gateway, backend: Backend +) -> None: + with wire_server(_answer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=backend.model, + api_base=wire.url, + additional_drop_params=["top_k", "top_p"], + **backend.credentials, + ) + _answered(*_messages(gateway, model, top_k=5, top_p=0.9)) + _sent_without_sampling_params(wire, backend) + + +@pytest.mark.parametrize("backend", _NATIVE_BACKENDS) +def test_messages_rejects_a_sampling_param_the_deployment_does_not_list_before_calling_the_provider( + gateway: Gateway, backend: Backend +) -> None: + with wire_server(_answer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=backend.model, + api_base=wire.url, + additional_drop_params=["top_k"], + **backend.credentials, + ) + rejected_status, rejected_text = _messages(gateway, model, top_k=5, top_p=0.9) + assert rejected_status == 400, rejected_text + assert "does not support top_p=0.9" in rejected_text, rejected_text + _answered(*_messages(gateway, model, top_k=5)) + _sent_without_sampling_params(wire, backend) diff --git a/tests/integration/observability/_openinference_support.py b/tests/integration/observability/_openinference_support.py index 58785c89033..3fc4b2e7628 100644 --- a/tests/integration/observability/_openinference_support.py +++ b/tests/integration/observability/_openinference_support.py @@ -672,10 +672,12 @@ def _write_config( callback_settings: Mapping[str, JsonValue] | None = None, litellm_settings: Mapping[str, JsonValue] | None = None, general_settings: Mapping[str, JsonValue] | None = None, + extra_config: Mapping[str, JsonValue] | None = None, ) -> Path: loaded: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) config: Final = { **loaded, + **(extra_config or {}), "litellm_settings": { **loaded["litellm_settings"], "callbacks": list(callbacks), @@ -752,6 +754,7 @@ def _rig( callback_settings: Mapping[str, JsonValue] | None = None, litellm_settings: Mapping[str, JsonValue] | None = None, general_settings: Mapping[str, JsonValue] | None = None, + extra_config: Mapping[str, JsonValue] | None = None, environment: Mapping[str, str] | None = None, remove_environment: tuple[str, ...] = (), disabled_environment: tuple[str, ...] = (), @@ -782,6 +785,7 @@ def _rig( callback_settings=resolved_settings, litellm_settings=litellm_settings, general_settings=general_settings, + extra_config=extra_config, ) preset_environment: Final = { **({"PHOENIX_COLLECTOR_ENDPOINT": destination.url + "/v1/traces"} if "arize_phoenix" in callbacks else {}), diff --git a/tests/integration/observability/test_arize_key_credential_reference.py b/tests/integration/observability/test_arize_key_credential_reference.py new file mode 100644 index 00000000000..b1087c4f4c5 --- /dev/null +++ b/tests/integration/observability/test_arize_key_credential_reference.py @@ -0,0 +1,428 @@ +from __future__ import annotations + +import operator +import time +import uuid +from collections.abc import Callable, Iterator +from itertools import accumulate +from pathlib import Path +from typing import Final + +import httpx +import pytest +from _openinference_support import ( + _anthropic_response, + _chat_request_marker, + _chat_response, + _chat_stream_response, + _chat_tool_call, + _matching_marker_spans, + _responses_response, + _rig, +) +from integration._support.client import Gateway, eventually +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +CREDENTIAL_NAME: Final = "arize-space-a" +SECRET_API_KEY: Final = "space-a-arize-key-" + uuid.uuid4().hex +SECRET_SPACE_ID: Final = "space-a-id-" + uuid.uuid4().hex +AMBIENT_ARIZE_ENVIRONMENT: Final = ("ARIZE_API_KEY", "ARIZE_SPACE_ID") +CREDENTIAL_CONFIG: Final[dict[str, JsonValue]] = { + "credential_list": [ + { + "credential_name": CREDENTIAL_NAME, + "credential_values": { + "arize_api_key": "os.environ/SPACE_A_ARIZE_API_KEY", + "arize_space_id": "os.environ/SPACE_A_ARIZE_SPACE_ID", + }, + "credential_info": {}, + }, + { + "credential_name": "model-only", + "credential_values": {"api_key": "sk-not-a-logging-value"}, + "credential_info": {}, + }, + ] +} +SECRET_ENVIRONMENT: Final = {"SPACE_A_ARIZE_API_KEY": SECRET_API_KEY, "SPACE_A_ARIZE_SPACE_ID": SECRET_SPACE_ID} + + +def _reference_entry(credential_name: str = CREDENTIAL_NAME) -> dict[str, JsonValue]: + return { + "callback_name": "arize", + "callback_type": "success", + "litellm_credential_name": credential_name, + "callback_vars": {}, + } + + +def _echo_marker(request: Request) -> Reply: + return _chat_response(_chat_request_marker(request)) + + +def _send_chat(proxy: Gateway, model: str, key: str, marker: str) -> httpx.Response: + return proxy.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": marker}], + "metadata": {"trace_marker": marker}, + "cache": {"no-cache": True}, + }, + key=key, + ) + + +def _chat(proxy: Gateway, model: str, key: str, marker: str) -> None: + response: Final = _send_chat(proxy, model, key, marker) + assert response.status_code == 200, response.text + + +def _drains(destination: Wire) -> Iterator[tuple[Request, ...]]: + while True: + yield destination.drain() + time.sleep(0.1) + + +def _carries(request: Request, marker: str) -> bool: + return any(True for _ in _matching_marker_spans((request,), marker)) + + +def _received_until(destination: Wire, marker: str, seconds: float = 30) -> tuple[Request, ...]: + deadline: Final = time.monotonic() + seconds + received: Final = next( + batch + for batch in accumulate(_drains(destination), operator.add) + if any(_carries(request, marker) for request in batch) or time.monotonic() >= deadline + ) + assert any(_carries(request, marker) for request in received), f"no span for {marker} reached the destination" + return received + + +def test_key_referencing_a_credential_exports_with_the_secret_manager_values(gateway: Gateway, tmp_path: Path) -> None: + unreferenced_marker: Final = "plain-" + uuid.uuid4().hex + referenced_marker: Final = "ref-" + uuid.uuid4().hex + with ( + _rig( + gateway, + tmp_path, + _echo_marker, + extra_config=CREDENTIAL_CONFIG, + environment=SECRET_ENVIRONMENT, + remove_environment=AMBIENT_ARIZE_ENVIRONMENT, + disabled_environment=AMBIENT_ARIZE_ENVIRONMENT, + ) as rig, + rig.proxy.scenario() as scenario, + ): + team: Final = scenario.team() + referencing_key: Final = scenario.key(team_id=team, metadata={"logging": [_reference_entry()]}) + unreferenced_key: Final = scenario.key(team_id=team) + + _chat(rig.proxy, rig.model, unreferenced_key, unreferenced_marker) + _chat(rig.proxy, rig.model, referencing_key, referenced_marker) + received: Final = _received_until(rig.destination, referenced_marker) + + exports: Final = tuple(request for request in received if _carries(request, referenced_marker)) + for export in exports: + headers = {name.lower(): value for name, value in export.headers.items()} + assert headers.get("api_key") == SECRET_API_KEY, headers + assert headers.get("arize-space-id") == SECRET_SPACE_ID, headers + assert not any(_carries(request, unreferenced_marker) for request in received) + + info: Final = rig.proxy.request("GET", "/key/info", params={"key": referencing_key}) + assert info.status_code == 200, info.text + assert CREDENTIAL_NAME in info.text + assert SECRET_API_KEY not in info.text + assert SECRET_SPACE_ID not in info.text + + +def test_only_proxy_admins_change_logging_on_a_key_that_references_a_credential( + gateway: Gateway, tmp_path: Path +) -> None: + with ( + _rig(gateway, tmp_path, _echo_marker, extra_config=CREDENTIAL_CONFIG) as rig, + rig.proxy.scenario() as scenario, + ): + team: Final = scenario.team() + team_admin: Final = scenario.member(team, role="admin") + team_admin_key: Final = scenario.key(user_id=team_admin, team_id=team) + langfuse_entry: Final[dict[str, JsonValue]] = { + "callback_name": "langfuse", + "callback_type": "success", + "callback_vars": {"langfuse_public_key": "pk-lf-integration", "langfuse_secret_key": "sk-lf-integration"}, + } + referencing_key: Final = scenario.key(team_id=team, metadata={"logging": [_reference_entry(), langfuse_entry]}) + + def as_team_admin(path: str, body: dict[str, JsonValue]) -> httpx.Response: + return rig.proxy.request("POST", path, body, key=team_admin_key) + + def assert_gated(response: httpx.Response) -> None: + assert response.status_code == 403, response.text + assert "litellm_credential_name" in response.text, response.text + + assert_gated(as_team_admin("/key/generate", {"team_id": team, "metadata": {"logging": [_reference_entry()]}})) + + sibling: Final = as_team_admin( + "/key/update", + { + "key": referencing_key, + "metadata": { + "logging": [ + _reference_entry(), + { + "callback_name": "arize", + "callback_type": "success", + "callback_vars": {"arize_endpoint": "http://127.0.0.1:9/v1/traces"}, + }, + ] + }, + }, + ) + assert_gated(sibling) + assert_gated(as_team_admin("/key/update", {"key": referencing_key, "metadata": {"team": "unrelated"}})) + assert_gated(as_team_admin("/key/regenerate", {"key": referencing_key, "metadata": {"team": "unrelated"}})) + + bulk: Final = as_team_admin( + "/team/key/bulk_update", + {"team_id": team, "key_ids": [referencing_key], "update_fields": {"metadata": {"team": "unrelated"}}}, + ) + assert bulk.status_code == 200, bulk.text + assert bulk.json()["successful_updates"] == [], bulk.text + assert "litellm_credential_name" in bulk.json()["failed_updates"][0]["failed_reason"], bulk.text + + renamed: Final = as_team_admin( + "/key/update", {"key": referencing_key, "key_alias": "alias-" + uuid.uuid4().hex} + ) + assert renamed.status_code == 200, renamed.text + + read_back: Final = rig.proxy.request("GET", "/key/info", params={"key": referencing_key}, key=team_admin_key) + assert read_back.status_code == 200, read_back.text + resaved: Final = as_team_admin( + "/key/update", + {"key": referencing_key, "metadata": {**read_back.json()["info"]["metadata"], "note": "unrelated edit"}}, + ) + assert resaved.status_code == 200, resaved.text + + stored: Final = rig.proxy.get("/key/info", params={"key": referencing_key}) + logging_entries: Final = stored["info"]["metadata"]["logging"] + assert [entry["callback_name"] for entry in logging_entries] == ["arize", "langfuse"], stored + assert logging_entries[0] == _reference_entry(), stored + + undefined: Final = rig.proxy.request( + "POST", "/key/generate", {"metadata": {"logging": [_reference_entry("no-such-credential")]}} + ) + assert undefined.status_code == 400, undefined.text + assert "no-such-credential" in undefined.text, undefined.text + + +def test_only_proxy_admins_create_a_user_whose_key_references_a_credential(gateway: Gateway, tmp_path: Path) -> None: + with ( + _rig(gateway, tmp_path, _echo_marker, extra_config=CREDENTIAL_CONFIG) as rig, + rig.proxy.scenario() as scenario, + ): + organization: Final = scenario.organization() + org_admin_key: Final = scenario.key(user_id=scenario.org_member(organization, "org_admin")) + referencing_metadata: Final = {"logging": [_reference_entry()]} + + single: Final = rig.proxy.request( + "POST", + "/user/new", + {"user_id": "user-" + uuid.uuid4().hex, "organizations": [organization], "metadata": referencing_metadata}, + key=org_admin_key, + ) + assert single.status_code == 403, single.text + assert "litellm_credential_name" in single.text, single.text + + delegate_key: Final = scenario.key( + user_id=scenario.user(user_role="internal_user"), allowed_routes=["/management/v1/users/bulk"] + ) + bulk: Final = rig.proxy.request( + "POST", + "/management/v1/users/bulk", + { + "users": [ + {"user_id": "user-" + uuid.uuid4().hex, "auto_create_key": True, "metadata": referencing_metadata} + ] + }, + key=delegate_key, + ) + assert bulk.status_code == 200, bulk.text + assert [row["success"] for row in bulk.json()["data"]] == [False], bulk.text + assert "litellm_credential_name" in bulk.json()["data"][0]["error"], bulk.text + + +def _carries_credential(request: Request) -> bool: + headers: Final = {name.lower(): value for name, value in request.headers.items()} + return headers.get("api_key") == SECRET_API_KEY and headers.get("arize-space-id") == SECRET_SPACE_ID + + +def _credential_exports(destination: Wire, seconds: float = 30) -> tuple[Request, ...]: + return eventually( + lambda: tuple(request for request in destination.drain() if _carries_credential(request)), + bool, + seconds=seconds, + ) + + +def _surface_reply(marker: str) -> Callable[[Request], Reply]: + def reply(request: Request) -> Reply: + if request.target.endswith("/responses"): + return _responses_response(marker) + if request.target.endswith("/messages"): + return _anthropic_response(marker) + return _chat_stream_response(marker, (_chat_tool_call(marker),)) + + return reply + + +_SURFACES: Final[dict[str, tuple[str, str, dict[str, JsonValue]]]] = { + "responses": ("/v1/responses", "openai/gpt-4o-mini", {"input": "weather in Paris?"}), + "messages": ( + "/v1/messages", + "anthropic/claude-opus-5-5", + {"max_tokens": 64, "messages": [{"role": "user", "content": "weather in Paris?"}]}, + ), + "chat-stream": ( + "/v1/chat/completions", + "openai/gpt-4o-mini", + {"stream": True, "messages": [{"role": "user", "content": "weather in Paris?"}]}, + ), +} + + +@pytest.mark.parametrize("surface", sorted(_SURFACES)) +def test_referencing_key_exports_on_each_surface(surface: str, gateway: Gateway, tmp_path: Path) -> None: + marker: Final = f"{surface}-" + uuid.uuid4().hex + path, model_name, body = _SURFACES[surface] + with ( + _rig( + gateway, + tmp_path, + _surface_reply(marker), + extra_config=CREDENTIAL_CONFIG, + environment=SECRET_ENVIRONMENT, + remove_environment=AMBIENT_ARIZE_ENVIRONMENT, + disabled_environment=AMBIENT_ARIZE_ENVIRONMENT, + model_name=model_name, + api_base_suffix="" if surface == "messages" else "/v1", + ) as rig, + rig.proxy.scenario() as scenario, + ): + key: Final = scenario.key(metadata={"logging": [_reference_entry()]}) + response: Final = rig.proxy.request( + "POST", + path, + {"model": rig.model, **body, "metadata": {"trace_marker": marker}, "cache": {"no-cache": True}}, + key=key, + ) + assert response.status_code == 200, response.text + received: Final = _received_until(rig.destination, marker) + exports: Final = tuple(request for request in received if _carries(request, marker)) + assert exports and all(_carries_credential(request) for request in exports), [ + dict(request.headers) for request in exports + ] + + +def test_referencing_key_exports_through_anthropic_pass_through(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "passthrough-" + uuid.uuid4().hex + with wire_server(lambda _request: _anthropic_response(marker)) as anthropic: + with ( + _rig( + gateway, + tmp_path, + _echo_marker, + extra_config=CREDENTIAL_CONFIG, + environment={ + **SECRET_ENVIRONMENT, + "ANTHROPIC_API_BASE": anthropic.url, + "ANTHROPIC_API_KEY": "sk-ant-synthetic", + }, + remove_environment=AMBIENT_ARIZE_ENVIRONMENT, + disabled_environment=AMBIENT_ARIZE_ENVIRONMENT, + ) as rig, + rig.proxy.scenario() as scenario, + ): + key: Final = scenario.key(metadata={"logging": [_reference_entry()]}) + response: Final = rig.proxy.request( + "POST", + "/anthropic/v1/messages", + {"model": "claude-opus-5-5", "max_tokens": 16, "messages": [{"role": "user", "content": marker}]}, + key=key, + headers={"anthropic-version": "2023-06-01"}, + ) + assert response.status_code == 200, response.text + assert response.json()["id"] == marker, response.text + assert anthropic.drain(), "the pass-through never reached the scripted Anthropic upstream" + assert _credential_exports(rig.destination) + + +def test_referencing_key_exports_on_otel_v1(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "otel-v1-" + uuid.uuid4().hex + with ( + _rig( + gateway, + tmp_path, + _echo_marker, + extra_config=CREDENTIAL_CONFIG, + environment={**SECRET_ENVIRONMENT, "LITELLM_OTEL_V2": "false"}, + remove_environment=AMBIENT_ARIZE_ENVIRONMENT, + disabled_environment=AMBIENT_ARIZE_ENVIRONMENT, + ) as rig, + rig.proxy.scenario() as scenario, + ): + key: Final = scenario.key(metadata={"logging": [_reference_entry()]}) + _chat(rig.proxy, rig.model, key, marker) + assert _credential_exports(rig.destination) + + +def test_reference_to_a_credential_without_logging_values_exports_nothing_and_is_unhealthy( + gateway: Gateway, tmp_path: Path +) -> None: + dropped_marker: Final = "model-only-" + uuid.uuid4().hex + sentinel_marker: Final = "sentinel-" + uuid.uuid4().hex + with ( + _rig( + gateway, + tmp_path, + _echo_marker, + extra_config=CREDENTIAL_CONFIG, + environment=SECRET_ENVIRONMENT, + remove_environment=AMBIENT_ARIZE_ENVIRONMENT, + disabled_environment=AMBIENT_ARIZE_ENVIRONMENT, + ) as rig, + rig.proxy.scenario() as scenario, + ): + dropped_key: Final = scenario.key(metadata={"logging": [_reference_entry("model-only")]}) + sentinel_key: Final = scenario.key(metadata={"logging": [_reference_entry()]}) + + _chat(rig.proxy, rig.model, dropped_key, dropped_marker) + _chat(rig.proxy, rig.model, sentinel_key, sentinel_marker) + received: Final = _received_until(rig.destination, sentinel_marker) + assert not any(_carries(request, dropped_marker) for request in received) + + health: Final = rig.proxy.request("POST", "/key/health", {}, key=dropped_key) + assert health.status_code == 200, health.text + assert health.json()["key"] == "unhealthy", health.text + + +def test_malformed_key_logging_fails_requests_like_main(gateway: Gateway, tmp_path: Path) -> None: + stored_logging: Final[dict[str, JsonValue]] = { + "string": "abc", + "list-with-string": ["x"], + "object": {"callback_name": "arize"}, + "empty-list": [], + } + with ( + _rig(gateway, tmp_path, _echo_marker, extra_config=CREDENTIAL_CONFIG) as rig, + rig.proxy.scenario() as scenario, + ): + statuses: Final = { + label: _send_chat( + rig.proxy, rig.model, scenario.key(metadata={"logging": value}), "malformed-" + uuid.uuid4().hex + ).status_code + for label, value in stored_logging.items() + } + + assert statuses == {"string": 500, "list-with-string": 500, "object": 500, "empty-list": 200}, statuses diff --git a/tests/integration/observability/test_arize_otel_v2_openinference_family.py b/tests/integration/observability/test_arize_otel_v2_openinference_family.py index f638a7360be..881a65c687e 100644 --- a/tests/integration/observability/test_arize_otel_v2_openinference_family.py +++ b/tests/integration/observability/test_arize_otel_v2_openinference_family.py @@ -23,7 +23,7 @@ from integration._support.client import Gateway from integration._support.wire import Reply, Request from pydantic import JsonValue -_GENAI_B3_KEYS_WITHOUT_BAGGAGE: Final = frozenset( +_GENAI_B3_KEYS: Final = frozenset( { "gen_ai.input.messages", "gen_ai.operation.name", @@ -56,6 +56,7 @@ _GENAI_B3_KEYS_WITHOUT_BAGGAGE: Final = frozenset( "litellm.cost.output", "litellm.cost.tool_usage", "litellm.cost.total", + "litellm.metadata.trace_marker", "litellm.provider.model", "litellm.request.route", "litellm.request.tools.declared", @@ -66,7 +67,6 @@ _GENAI_B3_KEYS_WITHOUT_BAGGAGE: Final = frozenset( "server.port", } ) -_GENAI_B3_KEYS_WITH_BAGGAGE: Final = _GENAI_B3_KEYS_WITHOUT_BAGGAGE | frozenset({"litellm.metadata.trace_marker"}) _LANGFUSE_B3_KEYS: Final = frozenset( { "gen_ai.input.messages", @@ -109,6 +109,7 @@ _LANGFUSE_B3_KEYS: Final = frozenset( "litellm.cost.output", "litellm.cost.tool_usage", "litellm.cost.total", + "litellm.metadata.trace_marker", "litellm.provider.model", "litellm.request.route", "litellm.request.tools.declared", @@ -122,15 +123,14 @@ _LANGFUSE_B3_KEYS: Final = frozenset( _B3_ATTRIBUTE_KEYS: Final[Mapping[str, frozenset[str]]] = MappingProxyType( { "langfuse_otel": _LANGFUSE_B3_KEYS, - "langtrace": _GENAI_B3_KEYS_WITHOUT_BAGGAGE, - "signoz": _GENAI_B3_KEYS_WITH_BAGGAGE, - "newrelic": _GENAI_B3_KEYS_WITH_BAGGAGE, - "levo": _GENAI_B3_KEYS_WITH_BAGGAGE, - "agentops": _GENAI_B3_KEYS_WITH_BAGGAGE, - "otel": _GENAI_B3_KEYS_WITH_BAGGAGE, + "langtrace": _GENAI_B3_KEYS, + "signoz": _GENAI_B3_KEYS, + "newrelic": _GENAI_B3_KEYS, + "levo": _GENAI_B3_KEYS, + "agentops": _GENAI_B3_KEYS, + "otel": _GENAI_B3_KEYS, } ) -_B3_BAGGAGE_CALLBACKS: Final = frozenset({"signoz", "newrelic", "levo", "agentops", "otel"}) _B4_ATTRIBUTE_KEYS: Final = frozenset( { "input.value", @@ -305,14 +305,10 @@ def test_arize_otel_v2_b3_non_openinference_callback_family(callback: str, gatew baggage: Final = tuple( sorted((key, value) for key, value in attributes.items() if key.startswith("litellm.metadata.")) ) - expected_baggage: Final = ( - (("litellm.metadata.trace_marker", marker),) if callback in _B3_BAGGAGE_CALLBACKS else () - ) - assert baggage == expected_baggage, attributes + assert baggage == (("litellm.metadata.trace_marker", marker),), attributes def test_arize_otel_v2_b3_langfuse_carries_metadata_baggage(gateway: Gateway, tmp_path: Path) -> None: - pytest.skip("BUG: LIT-9128 Langfuse OTel v2 preset omits request-metadata baggage") marker: Final = "b3-langfuse-" + uuid.uuid4().hex def upstream(request: Request) -> Reply: diff --git a/tests/integration/providers/test_anthropic_fable_capability_flags_wire.py b/tests/integration/providers/test_anthropic_fable_capability_flags_wire.py new file mode 100644 index 00000000000..0b140e4626f --- /dev/null +++ b/tests/integration/providers/test_anthropic_fable_capability_flags_wire.py @@ -0,0 +1,208 @@ +"""Claude Fable 5.2 over Anthropic direct. Its cost-map row says `supports_sampling_params: false`, +`supports_forced_tool_use: false` and `thinking_always_on: true`, and Anthropic answers 400 to a +request that ignores any of them, so the proxy has to rewrite or refuse those requests before they +leave. The scripted upstream here is what proves the outbound body, a leaked field is invisible in +the response.""" + +import json +import uuid +from typing import Final + +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +_MODEL: Final = "claude-fable-5-2" +_API_KEY: Final = "synthetic-anthropic-key" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_WEATHER_TOOL: Final[JsonValue] = { + "type": "function", + "function": { + "name": "get_weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + }, +} +_NAMED_TOOL_CHOICE: Final[JsonValue] = {"type": "function", "function": {"name": "get_weather"}} + + +def _anthropic_reply(identity: str, usage: dict[str, JsonValue] | None = None) -> bytes: + return json.dumps( + { + "id": identity, + "type": "message", + "role": "assistant", + "model": _MODEL, + "content": [ + {"type": "thinking", "thinking": "scripted thought", "signature": "scripted-signature"}, + {"type": "text", "text": "done"}, + ], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": usage or {"input_tokens": 12, "output_tokens": 3}, + } + ).encode() + + +def _chat_body(model: str, **extra: JsonValue) -> dict[str, JsonValue]: + return { + "model": model, + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi"}], + "cache": {"no-cache": True}, + **extra, + } + + +def _outbound(request: Request) -> dict[str, JsonValue]: + assert request.method == "POST" and request.target == "/v1/messages", (request.method, request.target) + assert request.headers["x-api-key"] == _API_KEY + body: Final = _JSON_OBJECT.validate_json(request.body) + assert body["model"] == _MODEL, body + return body + + +def test_drop_params_removes_temperature_and_top_p_before_anthropic(gateway: Gateway) -> None: + identity: Final = f"fable-sampling-{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + body: Final = _outbound(request) + assert "temperature" not in body and "top_p" not in body, body + return Reply(body=_anthropic_reply(identity)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY, drop_params=True + ) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, temperature=0.2, top_p=0.5)) + assert response.status_code == 200, response.text + assert len(wire.drain()) == 1 + + +def test_sampling_params_are_refused_before_reaching_anthropic_without_drop_params(gateway: Gateway) -> None: + def respond(request: Request) -> Reply: + raise AssertionError(f"a request with sampling params reached Anthropic: {request.body!r}") + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, temperature=0.2)) + assert response.status_code == 400, response.text + assert "temperature" in response.text and "drop_params" in response.text, response.text + assert wire.drain() == () + + +def test_drop_params_downgrades_forced_and_named_tool_choice_to_auto(gateway: Gateway) -> None: + identity: Final = f"fable-tool-choice-{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + body: Final = _outbound(request) + assert body["tool_choice"] == {"type": "auto"}, body + tools: Final = body["tools"] + assert isinstance(tools, list) and len(tools) == 1, body + first: Final = tools[0] + assert isinstance(first, dict) and first["name"] == "get_weather", body + return Reply(body=_anthropic_reply(identity)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY, drop_params=True + ) + responses: Final = tuple( + gateway.request( + "POST", "/v1/chat/completions", _chat_body(model, tools=[_WEATHER_TOOL], tool_choice=tool_choice) + ) + for tool_choice in ("required", _NAMED_TOOL_CHOICE) + ) + assert [response.status_code for response in responses] == [200, 200], [r.text for r in responses] + assert len(wire.drain()) == 2 + + +def test_forced_tool_choice_is_refused_before_reaching_anthropic_without_drop_params(gateway: Gateway) -> None: + def respond(request: Request) -> Reply: + raise AssertionError(f"a forced tool_choice reached Anthropic: {request.body!r}") + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request( + "POST", "/v1/chat/completions", _chat_body(model, tools=[_WEATHER_TOOL], tool_choice="required") + ) + assert response.status_code == 400, response.text + assert "tool" in response.text and "drop_params" in response.text, response.text + assert wire.drain() == () + + +def test_thinking_disabled_is_dropped_for_an_always_on_model(gateway: Gateway) -> None: + identity: Final = f"fable-thinking-off-{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + body: Final = _outbound(request) + assert "thinking" not in body, body + return Reply(body=_anthropic_reply(identity)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request( + "POST", "/v1/chat/completions", _chat_body(model, thinking={"type": "disabled"}) + ) + assert response.status_code == 200, response.text + payload: Final = _JSON_OBJECT.validate_json(response.content) + choices: Final = payload["choices"] + assert isinstance(choices, list) and isinstance(choices[0], dict), response.text + message: Final = choices[0]["message"] + assert isinstance(message, dict) and message["reasoning_content"] == "scripted thought", response.text + assert len(wire.drain()) == 1 + + +def test_reasoning_effort_becomes_adaptive_thinking_with_output_config_effort(gateway: Gateway) -> None: + identity: Final = f"fable-effort-{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + body: Final = _outbound(request) + assert body["thinking"] == {"type": "adaptive", "display": "summarized"}, body + assert body["output_config"] == {"effort": "max"}, body + assert "budget_tokens" not in json.dumps(body), body + return Reply(body=_anthropic_reply(identity)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, reasoning_effort="max")) + assert response.status_code == 200, response.text + assert len(wire.drain()) == 1 + + +def test_messages_one_hour_cache_control_reaches_anthropic_and_its_usage_comes_back(gateway: Gateway) -> None: + identity: Final = f"fable-cache-1h-{uuid.uuid4().hex}" + policy: Final = f"policy {identity}" + upstream_usage: Final[dict[str, JsonValue]] = { + "input_tokens": 12, + "output_tokens": 3, + "cache_creation_input_tokens": 9216, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 9216}, + } + + def respond(request: Request) -> Reply: + body: Final = _outbound(request) + assert body["system"] == [ + {"type": "text", "text": policy, "cache_control": {"type": "ephemeral", "ttl": "1h"}} + ], body + return Reply(body=_anthropic_reply(identity, upstream_usage)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/messages", + { + "model": model, + "max_tokens": 16, + "system": [{"type": "text", "text": policy, "cache_control": {"type": "ephemeral", "ttl": "1h"}}], + "messages": [{"role": "user", "content": "hi"}], + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == 200, response.text + usage: Final = _JSON_OBJECT.validate_json(response.content)["usage"] + assert isinstance(usage, dict), response.text + assert usage["cache_creation_input_tokens"] == 9216, usage + assert usage["cache_creation"] == {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 9216}, usage + assert len(wire.drain()) == 1 diff --git a/tests/integration/providers/test_bedrock_converse_fable_capability_flags_wire.py b/tests/integration/providers/test_bedrock_converse_fable_capability_flags_wire.py new file mode 100644 index 00000000000..815c5ba750f --- /dev/null +++ b/tests/integration/providers/test_bedrock_converse_fable_capability_flags_wire.py @@ -0,0 +1,234 @@ +"""us.anthropic.claude-fable-5-2 on Bedrock Converse: the row's capability flags shape the Converse body. + +Fable rejects sampling params and forced tool choice and always thinks, so the proxy must strip +or refuse those before the request reaches Bedrock, the same way it does on the Anthropic API. +""" + +import json +from collections.abc import Callable +from typing import Final + +from integration._support.client import Gateway, Scenario +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +_MODEL_ID: Final = "us.anthropic.claude-fable-5-2" +_TARGET: Final = f"/model/{_MODEL_ID}/converse" +_TOKEN: Final = "synthetic-bedrock-bearer" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_FORCED_TOOL_CHOICES: Final[tuple[JsonValue, ...]] = ("required", {"type": "function", "function": {"name": "lookup"}}) +_TOOLS: Final[list[JsonValue]] = [ + { + "type": "function", + "function": { + "name": "lookup", + "description": "Look up a synthetic record.", + "parameters": { + "type": "object", + "properties": {"record_id": {"type": "string"}}, + "required": ["record_id"], + }, + }, + } +] + + +def _reply(usage: dict[str, JsonValue] | None = None) -> bytes: + return json.dumps( + { + "output": { + "message": { + "role": "assistant", + "content": [ + { + "reasoningContent": { + "reasoningText": {"text": "scripted thought", "signature": "scripted-signature"} + } + }, + {"text": "done"}, + ], + } + }, + "stopReason": "end_turn", + "usage": usage or {"inputTokens": 12, "outputTokens": 3, "totalTokens": 15}, + "metrics": {"latencyMs": 1}, + } + ).encode() + + +def _deployment(scenario: Scenario, api_base: str, *, drop_params: bool = False) -> str: + if drop_params: + return scenario.model( + model=f"bedrock/converse/{_MODEL_ID}", + api_key=_TOKEN, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint=api_base, + num_retries=0, + drop_params=True, + ) + return scenario.model( + model=f"bedrock/converse/{_MODEL_ID}", + api_key=_TOKEN, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint=api_base, + num_retries=0, + ) + + +def _peer(on_body: Callable[[dict[str, JsonValue]], bytes]) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == _TARGET, (request.method, request.target) + assert request.headers["authorization"] == f"Bearer {_TOKEN}", request.headers + return Reply(body=on_body(_JSON_OBJECT.validate_json(request.body))) + + return respond + + +def _chat_body(model: str, **extra: JsonValue) -> dict[str, JsonValue]: + return { + "model": model, + "messages": [{"role": "user", "content": "synthetic fable request"}], + "max_tokens": 64, + **extra, + } + + +def test_drop_params_removes_temperature_and_top_p_before_bedrock(gateway: Gateway) -> None: + def on_body(body: dict[str, JsonValue]) -> bytes: + assert body["inferenceConfig"] == {"maxTokens": 64}, body + return _reply() + + with wire_server(_peer(on_body)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url, drop_params=True) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, temperature=0.2, top_p=0.9)) + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["content"] == "done", response.text + assert len(wire.drain()) == 1 + + +def test_sampling_params_are_refused_before_reaching_bedrock_without_drop_params(gateway: Gateway) -> None: + def on_body(body: dict[str, JsonValue]) -> bytes: + raise AssertionError(f"temperature must not reach Bedrock: {body}") + + with wire_server(_peer(on_body)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, temperature=0.2)) + assert response.status_code == 400, response.text + assert "temperature" in response.text and "drop_params" in response.text, response.text + assert len(wire.drain()) == 0 + + +def test_drop_params_downgrades_forced_and_named_tool_choice_to_auto(gateway: Gateway) -> None: + def on_body(body: dict[str, JsonValue]) -> bytes: + tool_config: Final = body["toolConfig"] + assert isinstance(tool_config, dict), body + assert tool_config.get("toolChoice") in (None, {"auto": {}}), body + tools: Final = tool_config["tools"] + assert isinstance(tools, list) and len(tools) == 1, body + return _reply() + + with wire_server(_peer(on_body)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url, drop_params=True) + responses: Final = tuple( + gateway.request("POST", "/v1/chat/completions", _chat_body(model, tools=_TOOLS, tool_choice=tool_choice)) + for tool_choice in _FORCED_TOOL_CHOICES + ) + assert [response.status_code for response in responses] == [200, 200], [response.text for response in responses] + assert len(wire.drain()) == 2 + + +def test_forced_tool_choice_is_refused_before_reaching_bedrock_without_drop_params(gateway: Gateway) -> None: + def on_body(body: dict[str, JsonValue]) -> bytes: + raise AssertionError(f"forced tool_choice must not reach Bedrock: {body}") + + with wire_server(_peer(on_body)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", "/v1/chat/completions", _chat_body(model, tools=_TOOLS, tool_choice="required") + ) + assert response.status_code == 400, response.text + assert "tool_choice" in response.text, response.text + assert len(wire.drain()) == 0 + + +def test_thinking_disabled_is_dropped_for_an_always_on_model(gateway: Gateway) -> None: + def on_body(body: dict[str, JsonValue]) -> bytes: + assert "thinking" not in json.dumps(body.get("additionalModelRequestFields", {})), body + return _reply() + + with wire_server(_peer(on_body)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", "/v1/chat/completions", _chat_body(model, thinking={"type": "disabled"}) + ) + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["reasoning_content"] == "scripted thought", response.text + assert len(wire.drain()) == 1 + + +def test_reasoning_effort_becomes_adaptive_thinking_with_output_config_effort(gateway: Gateway) -> None: + def on_body(body: dict[str, JsonValue]) -> bytes: + extra: Final = body["additionalModelRequestFields"] + assert isinstance(extra, dict), body + assert extra["thinking"] == {"type": "adaptive", "display": "summarized"}, body + assert extra["output_config"] == {"effort": "max"}, body + assert "budget_tokens" not in json.dumps(body), body + return _reply() + + with wire_server(_peer(on_body)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, reasoning_effort="max")) + assert response.status_code == 200, response.text + assert len(wire.drain()) == 1 + + +def test_one_hour_cache_point_reaches_bedrock_and_its_usage_comes_back(gateway: Gateway) -> None: + usage: Final[dict[str, JsonValue]] = { + "inputTokens": 9, + "outputTokens": 3, + "totalTokens": 7237, + "cacheReadInputTokens": 0, + "cacheWriteInputTokens": 7225, + "cacheDetails": [{"ttl": "1h", "inputTokens": 7225}], + } + + def on_body(body: dict[str, JsonValue]) -> bytes: + assert body["system"] == [ + {"text": "synthetic cached system prompt"}, + {"cachePoint": {"type": "default", "ttl": "1h"}}, + ], body + return _reply(usage) + + with wire_server(_peer(on_body)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + _chat_body( + model, + messages=[ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "synthetic cached system prompt", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + }, + {"role": "user", "content": "synthetic fable request"}, + ], + ), + ) + assert response.status_code == 200, response.text + usage_out: Final = _JSON_OBJECT.validate_json(response.text)["usage"] + assert isinstance(usage_out, dict), response.text + prompt_details: Final = usage_out["prompt_tokens_details"] + assert isinstance(prompt_details, dict), response.text + assert prompt_details["cache_creation_token_details"] == { + "ephemeral_5m_input_tokens": 0, + "ephemeral_1h_input_tokens": 7225, + }, response.text + assert usage_out["cache_creation_input_tokens"] == 7225, response.text + assert len(wire.drain()) == 1 diff --git a/tests/integration/providers/test_vertex_anthropic_fable_capability_flags_wire.py b/tests/integration/providers/test_vertex_anthropic_fable_capability_flags_wire.py new file mode 100644 index 00000000000..8376cd92de7 --- /dev/null +++ b/tests/integration/providers/test_vertex_anthropic_fable_capability_flags_wire.py @@ -0,0 +1,242 @@ +"""claude-fable-5-2 on Vertex AI: the row's capability flags shape the rawPredict body the proxy sends. + +Fable rejects sampling params and forced tool choice and always thinks, so the proxy must strip +or refuse those before the request reaches Vertex, the same way it does on the Anthropic API. +""" + +import json +import uuid +from collections.abc import Callable +from typing import Final +from urllib.parse import parse_qs + +from integration._support.client import Gateway, Scenario +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +_MODEL: Final = "claude-fable-5-2" +_PROJECT: Final = "scripted-project" +_LOCATION: Final = "global" +_TARGET: Final = f"/v1/projects/{_PROJECT}/locations/{_LOCATION}/publishers/anthropic/models/{_MODEL}:rawPredict" +_TOKEN: Final = "scripted-vertex-access-token" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_FORCED_TOOL_CHOICES: Final[tuple[JsonValue, ...]] = ("required", {"type": "function", "function": {"name": "lookup"}}) +_TOOLS: Final[list[JsonValue]] = [ + { + "type": "function", + "function": { + "name": "lookup", + "description": "Look up a synthetic record.", + "parameters": { + "type": "object", + "properties": {"record_id": {"type": "string"}}, + "required": ["record_id"], + }, + }, + } +] + + +def _reply(identity: str, usage: dict[str, JsonValue] | None = None) -> bytes: + return json.dumps( + { + "id": identity, + "type": "message", + "role": "assistant", + "model": _MODEL, + "content": [ + {"type": "thinking", "thinking": "scripted thought", "signature": "scripted-signature"}, + {"type": "text", "text": "done"}, + ], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": usage or {"input_tokens": 12, "output_tokens": 3}, + } + ).encode() + + +def _deployment(scenario: Scenario, api_base: str, *, drop_params: bool = False) -> str: + credentials: Final = service_account_json(_PROJECT, api_base) + if drop_params: + return scenario.model( + model=f"vertex_ai/{_MODEL}", + api_base=api_base, + api_key=None, + vertex_project=_PROJECT, + vertex_location=_LOCATION, + vertex_credentials=credentials, + num_retries=0, + drop_params=True, + ) + return scenario.model( + model=f"vertex_ai/{_MODEL}", + api_base=api_base, + api_key=None, + vertex_project=_PROJECT, + vertex_location=_LOCATION, + vertex_credentials=credentials, + num_retries=0, + ) + + +def _peer(on_body: Callable[[dict[str, JsonValue]], bytes]) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.target == "/_oauth/token": + form: Final = parse_qs(request.body.decode()) + assert form["grant_type"] == ["urn:ietf:params:oauth:grant-type:jwt-bearer"], form + return Reply(body=json.dumps({"access_token": _TOKEN, "token_type": "Bearer", "expires_in": 3600}).encode()) + assert request.method == "POST" and request.target == _TARGET, (request.method, request.target) + assert request.headers["authorization"] == f"Bearer {_TOKEN}", request.headers + body: Final = _JSON_OBJECT.validate_json(request.body) + assert body["anthropic_version"] == "vertex-2023-10-16", body + return Reply(body=on_body(body)) + + return respond + + +def _chat_body(model: str, **extra: JsonValue) -> dict[str, JsonValue]: + return { + "model": model, + "messages": [{"role": "user", "content": "synthetic fable request"}], + "max_tokens": 64, + **extra, + } + + +def _model_calls(requests: tuple[Request, ...]) -> int: + return sum(request.target == _TARGET for request in requests) + + +def test_drop_params_removes_temperature_and_top_p_before_vertex(gateway: Gateway) -> None: + identity: Final = f"msg_{uuid.uuid4().hex}" + + def on_body(body: dict[str, JsonValue]) -> bytes: + assert "temperature" not in body and "top_p" not in body, body + assert body["max_tokens"] == 64, body + return _reply(identity) + + with wire_server(_peer(on_body)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url, drop_params=True) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, temperature=0.2, top_p=0.9)) + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["content"] == "done", response.text + assert _model_calls(wire.drain()) == 1 + + +def test_sampling_params_are_refused_before_reaching_vertex_without_drop_params(gateway: Gateway) -> None: + def on_body(body: dict[str, JsonValue]) -> bytes: + raise AssertionError(f"temperature must not reach Vertex: {body}") + + with wire_server(_peer(on_body)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, temperature=0.2)) + assert response.status_code == 400, response.text + assert "temperature" in response.text and "drop_params" in response.text, response.text + assert _model_calls(wire.drain()) == 0 + + +def test_drop_params_downgrades_forced_and_named_tool_choice_to_auto(gateway: Gateway) -> None: + identity: Final = f"msg_{uuid.uuid4().hex}" + + def on_body(body: dict[str, JsonValue]) -> bytes: + assert body["tool_choice"] == {"type": "auto"}, body + tools: Final = body["tools"] + assert isinstance(tools, list) and [tool["name"] for tool in tools if isinstance(tool, dict)] == ["lookup"], ( + body + ) + return _reply(identity) + + with wire_server(_peer(on_body)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url, drop_params=True) + responses: Final = tuple( + gateway.request("POST", "/v1/chat/completions", _chat_body(model, tools=_TOOLS, tool_choice=tool_choice)) + for tool_choice in _FORCED_TOOL_CHOICES + ) + assert [response.status_code for response in responses] == [200, 200], [response.text for response in responses] + assert _model_calls(wire.drain()) == 2 + + +def test_forced_tool_choice_is_refused_before_reaching_vertex_without_drop_params(gateway: Gateway) -> None: + def on_body(body: dict[str, JsonValue]) -> bytes: + raise AssertionError(f"forced tool_choice must not reach Vertex: {body}") + + with wire_server(_peer(on_body)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", "/v1/chat/completions", _chat_body(model, tools=_TOOLS, tool_choice="required") + ) + assert response.status_code == 400, response.text + assert "tool_choice" in response.text, response.text + assert _model_calls(wire.drain()) == 0 + + +def test_thinking_disabled_is_dropped_for_an_always_on_model(gateway: Gateway) -> None: + identity: Final = f"msg_{uuid.uuid4().hex}" + + def on_body(body: dict[str, JsonValue]) -> bytes: + assert "thinking" not in body, body + return _reply(identity) + + with wire_server(_peer(on_body)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", "/v1/chat/completions", _chat_body(model, thinking={"type": "disabled"}) + ) + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["reasoning_content"] == "scripted thought", response.text + assert _model_calls(wire.drain()) == 1 + + +def test_reasoning_effort_becomes_adaptive_thinking_with_output_config_effort(gateway: Gateway) -> None: + identity: Final = f"msg_{uuid.uuid4().hex}" + + def on_body(body: dict[str, JsonValue]) -> bytes: + assert body["thinking"] == {"type": "adaptive", "display": "summarized"}, body + assert body["output_config"] == {"effort": "max"}, body + assert "budget_tokens" not in json.dumps(body), body + return _reply(identity) + + with wire_server(_peer(on_body)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, reasoning_effort="max")) + assert response.status_code == 200, response.text + assert _model_calls(wire.drain()) == 1 + + +def test_messages_one_hour_cache_control_reaches_vertex_and_its_usage_comes_back(gateway: Gateway) -> None: + identity: Final = f"msg_{uuid.uuid4().hex}" + system: Final[list[JsonValue]] = [ + {"type": "text", "text": "synthetic cached system prompt", "cache_control": {"type": "ephemeral", "ttl": "1h"}} + ] + usage: Final[dict[str, JsonValue]] = { + "input_tokens": 9, + "output_tokens": 3, + "cache_creation_input_tokens": 7225, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 7225}, + } + + def on_body(body: dict[str, JsonValue]) -> bytes: + assert body["system"] == system, body + return _reply(identity, usage) + + with wire_server(_peer(on_body)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request( + "POST", + "/v1/messages", + { + "model": model, + "max_tokens": 64, + "system": system, + "messages": [{"role": "user", "content": "synthetic fable request"}], + }, + ) + assert response.status_code == 200, response.text + assert response.json()["usage"]["cache_creation"] == { + "ephemeral_5m_input_tokens": 0, + "ephemeral_1h_input_tokens": 7225, + }, response.text + assert response.json()["usage"]["cache_creation_input_tokens"] == 7225, response.text + assert _model_calls(wire.drain()) == 1 diff --git a/tests/integration/providers/test_vertex_labels_router_default_metadata_wire.py b/tests/integration/providers/test_vertex_labels_router_default_metadata_wire.py new file mode 100644 index 00000000000..89747a5bed8 --- /dev/null +++ b/tests/integration/providers/test_vertex_labels_router_default_metadata_wire.py @@ -0,0 +1,85 @@ +import json +from collections.abc import Mapping +from pathlib import Path +from typing import Final + +import pytest +from integration._support.client import Gateway +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_MODEL_NAME: Final = "labelled-gemini" +_PROJECT: Final = "scripted-project" +_LOCATION: Final = "us-central1" +_ROUTER_LABELS: Final = {"env": "prod"} +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_REPLY: Final = json.dumps( + { + "candidates": [{"content": {"role": "model", "parts": [{"text": "OK"}]}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 3, "candidatesTokenCount": 1, "totalTokenCount": 4}, + } +).encode() + + +def _gemini_peer(request: Request) -> Reply: + if request.target.endswith(":generateContent"): + return Reply(body=_REPLY) + return Reply(status=404, body=json.dumps({"error": request.target}).encode()) + + +def _config(directory: Path, gateway: Gateway, api_base: str) -> Path: + config: Final = directory / "router-default-requester-metadata.yaml" + config.write_text( + json.dumps( + { + "model_list": [ + { + "model_name": _MODEL_NAME, + "litellm_params": { + "model": "vertex_ai/gemini-3.8-flash", + "api_base": api_base, + "vertex_project": _PROJECT, + "vertex_location": _LOCATION, + "vertex_credentials": service_account_json(_PROJECT, gateway.upstream_url.rstrip("/")), + }, + } + ], + "router_settings": {"default_litellm_params": {"metadata": {"requester_metadata": _ROUTER_LABELS}}}, + } + ), + encoding="utf-8", + ) + return config + + +def _labels_sent(proxy: Gateway, wire: Wire, metadata: Mapping[str, JsonValue] | None) -> JsonValue: + response: Final = proxy.request( + "POST", + "/v1/chat/completions", + { + "model": _MODEL_NAME, + "messages": [{"role": "user", "content": "Return exactly OK."}], + **({"metadata": dict(metadata)} if metadata is not None else {}), + }, + ) + assert response.status_code == 200, response.text + generations: Final = tuple(request for request in wire.drain() if request.target.endswith(":generateContent")) + assert len(generations) == 1, generations + return _JSON_OBJECT.validate_json(generations[0].body).get("labels") + + +def test_router_default_requester_metadata_merges_under_each_callers_vertex_labels( + gateway: Gateway, tmp_path: Path +) -> None: + with ( + wire_server(_gemini_peer) as wire, + owned_proxy_process(gateway, tmp_path, {}, config=_config(tmp_path, gateway, wire.url)) as owned, + ): + proxy: Final = Gateway(owned.gateway.client, owned.gateway.key, gateway.upstream_url) + assert _labels_sent(proxy, wire, {"team": "ai-platform"}) == {**_ROUTER_LABELS, "team": "ai-platform"} + assert _labels_sent(proxy, wire, {"env": "staging"}) == {"env": "staging"} + assert _labels_sent(proxy, wire, None) == _ROUTER_LABELS diff --git a/tests/integration/routing/test_usage_based_routing_redis_reads.py b/tests/integration/routing/test_usage_based_routing_redis_reads.py index 8a31eac7f91..edb24cc7930 100644 --- a/tests/integration/routing/test_usage_based_routing_redis_reads.py +++ b/tests/integration/routing/test_usage_based_routing_redis_reads.py @@ -3,6 +3,7 @@ from __future__ import annotations import json import shlex import threading +import time import uuid from collections.abc import Iterator from contextlib import contextmanager @@ -19,6 +20,7 @@ from integration._support.client import Gateway, eventually, object_value, strin from integration._support.process import owned_proxy from integration._support.redis_process import owned_redis from integration._support.wire import Reply, Request, wire_server +from litellm.constants import DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS from pydantic import JsonValue, TypeAdapter from redis import Redis from redis.exceptions import TimeoutError as RedisTimeoutError @@ -217,11 +219,12 @@ def test_proxy_usage_routing_reads_cooldown_tpm_then_rpm_from_redis( lambda registered: registered == frozenset(deployment_ids), seconds=15, ) + cooldown_read_throttle_ends: Final = time.monotonic() + DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS wire.drain() eventually( - lambda: datetime.now(UTC), - lambda current: current.second < 40, - seconds=65, + lambda: (datetime.now(UTC), time.monotonic()), + lambda clock: clock[0].second < 40 and clock[1] >= cooldown_read_throttle_ends, + seconds=65 + DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS, ) minute: Final = datetime.now(UTC).strftime("%H-%M") markers: Final = tuple(f"{run_id}-{index}" for index in range(3)) @@ -256,18 +259,18 @@ def test_proxy_usage_routing_reads_cooldown_tpm_then_rpm_from_redis( observed_bodies: Final = tuple(_request_object(request.body) for request in received) expected_bodies: Final = tuple(_expected_wire_body(endpoint, marker) for marker in markers[:2]) assert observed_bodies == expected_bodies, observed_bodies - expected_mget: Final = ( + cooldown_mget: Final = ( "MGET", - f"deployment:{deployment_ids[0]}:cooldown", - f"deployment:{deployment_ids[1]}:cooldown", - f"{deployment_ids[0]}:openai/{OPENAI_MODEL}:tpm:{minute}", - f"{deployment_ids[1]}:openai/{OPENAI_MODEL}:tpm:{minute}", - *( - f"{deployment_id}:openai/{OPENAI_MODEL}:rpm:{minute}" - for deployment_id in deployment_ids - ), + *(f"deployment:{deployment_id}:cooldown" for deployment_id in deployment_ids), + ) + usage_mget: Final = ( + "MGET", + *(f"{deployment_id}:openai/{OPENAI_MODEL}:tpm:{minute}" for deployment_id in deployment_ids), + *(f"{deployment_id}:openai/{OPENAI_MODEL}:rpm:{minute}" for deployment_id in deployment_ids), ) mgets: Final = _drain_mgets(commands) - assert any(arguments == expected_mget for _, arguments in mgets), mgets + mget_arguments: Final = tuple(arguments for _, arguments in mgets) + assert cooldown_mget in mget_arguments, mgets + assert usage_mget in mget_arguments, mgets raw_mgets: Final = tuple(line for line, _ in mgets) print(f"proxy {endpoint} MGETs: {raw_mgets}") diff --git a/tests/integration/sdk/test_router_concurrent_deployment_delete_sdk.py b/tests/integration/sdk/test_router_concurrent_deployment_delete_sdk.py index 3867808aa5f..9f8aa966335 100644 --- a/tests/integration/sdk/test_router_concurrent_deployment_delete_sdk.py +++ b/tests/integration/sdk/test_router_concurrent_deployment_delete_sdk.py @@ -19,6 +19,8 @@ from pydantic import JsonValue import litellm from litellm import Router +from litellm.router_rust.surface.model_list import HostedModelList +from litellm.router_rust.surface.registry import HostRegistry _DEPLOYMENTS: Final = 4 _TARGET: Final = _DEPLOYMENTS - 1 @@ -28,12 +30,13 @@ _BURST_THREADS: Final = 8 _BURST_CALLS: Final = 24 _CHURN_UPSERTS: Final = 120 _UNREPORTED_MODEL_ID: Final = frozenset({"aresponses-stream", "aanthropic-messages"}) -_REINDEX: Final = Router._update_deployment_indices_after_removal.__code__ +_LIST_REFRESH: Final = HostedModelList.refresh.__code__ +_HAND_CHANGES: Final = HostRegistry._hand_changes.__code__ def _router_source_files() -> frozenset[str]: facade_files: Final = (sys.modules[cls.__module__].__file__ for cls in Router.__mro__ if cls is not object) - return frozenset({_REINDEX.co_filename, *(path for path in facade_files if path is not None)}) + return frozenset({_HAND_CHANGES.co_filename, *(path for path in facade_files if path is not None)}) def _lock_lines(path: str) -> frozenset[int]: @@ -88,9 +91,9 @@ def _router(wire: Wire, team: bool = False) -> Router: def _read_during_held_write(read: Callable[[], T], write: Callable[[], object]) -> T: - """Run `read` on its own thread while `write` waits where the list has shrunk but the index maps have not caught - up. The write resumes once `read` returns or reaches the router's index lock, so a read that skips the lock meets - the half-updated list and a read that takes the lock waits for the write""" + """Run `read` on its own thread while `write` waits where router-core has applied the change but the Python model + list has not been read again. The write resumes once `read` returns or reaches the router's index lock, so a read + that skips the lock meets the half-updated list and a read that takes the lock waits for the write""" held: Final = threading.Event() release: Final = threading.Event() settled: Final = threading.Event() @@ -99,7 +102,11 @@ def _read_during_held_write(read: Callable[[], T], write: Callable[[], object]) failure: Final[list[BaseException]] = [] def hold_at_reindex(frame: FrameType, event: str, _arg: object) -> None: - if event == "call" and frame.f_code is _REINDEX and not held.is_set(): + caller: Final = frame.f_back + refreshes_after_a_write: Final = ( + frame.f_code is _LIST_REFRESH and caller is not None and caller.f_code is _HAND_CHANGES + ) + if event == "call" and refreshes_after_a_write and not held.is_set(): held.set() released_in_time.append(release.wait(timeout=_SYNC_SECONDS)) diff --git a/tests/integration/sdk/test_usage_based_routing_sdk_redis_reads.py b/tests/integration/sdk/test_usage_based_routing_sdk_redis_reads.py index da9a8309c8d..ff4761b1da3 100644 --- a/tests/integration/sdk/test_usage_based_routing_sdk_redis_reads.py +++ b/tests/integration/sdk/test_usage_based_routing_sdk_redis_reads.py @@ -140,7 +140,7 @@ async def _exercise_router(router: Router, model_name: str, markers: tuple[str, return _model_id(first), _model_id(second) -def test_sdk_usage_routing_reads_tpm_then_rpm_from_redis(tmp_path: Path) -> None: +def test_sdk_usage_routing_reads_cooldown_tpm_then_rpm_from_redis(tmp_path: Path) -> None: with owned_redis(tmp_path) as cache, wire_server(_reply) as wire: run_id: Final = uuid.uuid4().hex model_name: Final = f"usage-redis-{run_id}" @@ -176,15 +176,19 @@ def test_sdk_usage_routing_reads_tpm_then_rpm_from_redis(tmp_path: Path) -> None for marker in markers[:2] ) assert observed_bodies == expected_bodies, observed_bodies - expected_mget: Final = ( + cooldown_mget: Final = ( + "MGET", + *(f"deployment:{deployment_id}:cooldown" for deployment_id in deployment_ids), + ) + usage_mget: Final = ( "MGET", - f"deployment:{deployment_ids[0]}:cooldown", - f"deployment:{deployment_ids[1]}:cooldown", *(f"{deployment_id}:openai/{OPENAI_MODEL}:tpm:{minute}" for deployment_id in deployment_ids), *(f"{deployment_id}:openai/{OPENAI_MODEL}:rpm:{minute}" for deployment_id in deployment_ids), ) mgets: Final = _drain_mgets(commands) - assert any(arguments == expected_mget for _, arguments in mgets), mgets + mget_arguments: Final = tuple(arguments for _, arguments in mgets) + assert cooldown_mget in mget_arguments, mgets + assert usage_mget in mget_arguments, mgets raw_mgets: Final = tuple(line for line, _ in mgets) print(f"sdk MGETs: {raw_mgets}") finally: diff --git a/tests/integration/spend/test_fail_closed_team_member_budget.py b/tests/integration/spend/test_fail_closed_team_member_budget.py new file mode 100644 index 00000000000..52b2e8dd0f3 --- /dev/null +++ b/tests/integration/spend/test_fail_closed_team_member_budget.py @@ -0,0 +1,66 @@ +import uuid +from collections.abc import Iterator +from contextlib import ExitStack +from pathlib import Path +from typing import Final + +import pytest +import yaml +from integration._support.client import JSON_OBJECT, Gateway, gateway_from_environment, object_value, string_value +from integration._support.process import owned_proxy_process +from pydantic import JsonValue + + +def _fail_closed_config(directory: Path) -> Path: + base: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + config: Final = { + **base, + "general_settings": {**object_value(base["general_settings"]), "fail_closed_budget_enforcement": True}, + } + path: Final = directory / f"fail-closed-{uuid.uuid4().hex[:8]}.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@pytest.fixture(scope="module") +def fail_closed_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: + directory: Final = tmp_path_factory.mktemp("fail-closed-team-member") + with ExitStack() as stack: + rig_gateway: Final = stack.enter_context(gateway_from_environment()) + owned: Final = stack.enter_context( + owned_proxy_process(rig_gateway, directory, {}, config=_fail_closed_config(directory)) + ) + yield owned.gateway + + +def _chat(gateway: Gateway, model: str, key: str) -> tuple[int, dict[str, JsonValue]]: + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "integration control"}]}, + key=key, + ) + return response.status_code, JSON_OBJECT.validate_json(response.content) + + +def test_team_member_without_membership_row_is_told_to_be_added_then_admitted(fail_closed_gateway: Gateway) -> None: + with fail_closed_gateway.scenario() as scenario: + model: Final = scenario.model() + team_id: Final = scenario.team(models=[model], team_member_budget=100.0) + user_id: Final = scenario.user(max_budget=100.0) + key: Final = scenario.key(team_id=team_id, user_id=user_id, models=[model]) + + status, body = _chat(fail_closed_gateway, model, key) + + assert status == 403, body + message: Final = string_value(object_value(body["error"])["message"]) + assert user_id in message and team_id in message and "/team/member_add" in message + + fail_closed_gateway.post( + "/team/member_add", {"team_id": team_id, "member": {"role": "user", "user_id": user_id}} + ) + + admitted_status, admitted = _chat(fail_closed_gateway, model, key) + + assert admitted_status == 200, admitted + assert object_value(admitted["usage"])["total_tokens"] diff --git a/tests/integration/spend/test_passthrough_failure.py b/tests/integration/spend/test_passthrough_failure.py new file mode 100644 index 00000000000..cb102c406f7 --- /dev/null +++ b/tests/integration/spend/test_passthrough_failure.py @@ -0,0 +1,53 @@ +import json +from typing import Final +from pathlib import Path + +import pytest +from tests.integration._support.client import Gateway, eventually +from tests.integration._support.database import read_rows +from tests.integration._support.process import owned_proxy +from tests.integration._support.wire import Reply, Request, wire_server + + +@pytest.mark.timeout(180) +def test_config_pass_through_endpoint_failure_writes_a_spend_row(gateway: Gateway, tmp_path: Path) -> None: + def respond(request: Request) -> Reply: + assert request.target == "/" + assert json.loads(request.body) == {"search": "hello"} + return Reply(status=500, body=json.dumps({"error": "upstream search failure"}).encode()) + + with wire_server(respond) as wire: + config: Final = tmp_path / "proxy_config.yaml" + config.write_text( + "model_list: []\n" + "general_settings:\n" + " master_key: os.environ/LITELLM_MASTER_KEY\n" + " database_url: os.environ/DATABASE_URL\n" + " store_model_in_db: true\n" + " disable_spend_logs: false\n" + " proxy_batch_write_at: 1\n" + " pass_through_endpoints:\n" + " - path: /azure-search\n" + f" target: {wire.url}\n" + " headers:\n" + " Authorization: Bearer synthetic-search-key\n" + "router_settings:\n" + " disable_cooldowns: true\n" + ) + with owned_proxy(gateway, tmp_path, {}, config=config) as candidate: + response: Final = candidate.request("POST", "/azure-search", {"search": "hello"}) + assert response.status_code == 500, response.text + call_id: Final = response.headers["x-litellm-call-id"] + rows: Final = eventually( + lambda: read_rows( + 'SELECT status, call_type, api_base, spend FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (call_id,), + database_url=candidate.database_url, + ), + lambda values: len(values) == 1, + seconds=70, + ) + assert rows[0]["status"] == "failure", rows + assert rows[0]["call_type"] == "pass_through_endpoint", rows + assert "127.0.0.1" in rows[0]["api_base"], rows + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/")] diff --git a/tests/integration/streaming/test_stream_clean_eof_without_finish_reason.py b/tests/integration/streaming/test_stream_clean_eof_without_finish_reason.py index f3b427d91c4..933c152eb26 100644 --- a/tests/integration/streaming/test_stream_clean_eof_without_finish_reason.py +++ b/tests/integration/streaming/test_stream_clean_eof_without_finish_reason.py @@ -464,7 +464,7 @@ def test_concurrent_burst_ends_each_stream_by_whether_its_provider_finished_it(g ) -_SPEND_LOG_INSERT: Final = b'INSERT INTO "public"."LiteLLM_SpendLogs"' +_SPEND_LOG_INSERT: Final = b'"LiteLLM_SpendLogs" (' def _spend_statuses(database_url: str | None) -> Counter[tuple[str, str]]: diff --git a/tests/test_litellm_rust/cache/test_python_cache.py b/tests/test_litellm_rust/cache/test_python_cache.py index 9e76dd468c4..aa1dfbc70c0 100644 --- a/tests/test_litellm_rust/cache/test_python_cache.py +++ b/tests/test_litellm_rust/cache/test_python_cache.py @@ -9,7 +9,11 @@ from litellm.caching.caching import Cache, CacheMode from litellm.caching.caching_handler import ( _PENDING_CACHE_WRITES, # pyright: ignore[reportPrivateUsage] # await the existing background cache writer before the next request ) -from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict +from litellm.litellm_core_utils.add_retry_fallback_headers import ( + _add_headers_to_response, # pyright: ignore[reportPrivateUsage] # stamp the response the way the router does + add_retry_headers_to_response, + get_hidden_params_dict, +) from litellm.proxy._types import Litellm_EntityType from litellm.proxy.hooks.model_max_budget_limiter import ( PROXY_VirtualKeyModelMaxBudgetLimiter, @@ -19,6 +23,8 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelReq from litellm.proxy.utils import InternalUsageCache from litellm.redis.stores.rate_limits import RateLimitScope from litellm.rust_bridge.catalog import Route +from litellm.types.llms.anthropic import AnthropicResponseContentBlockText +from litellm.types.utils import CallTypes, LLMResponseTypes from tests.test_litellm_rust.cache.support import cache_key, collect, invoke, payload from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_EVENTS, MESSAGES_MODEL, MESSAGES_RESPONSE from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging, request_body @@ -251,9 +257,11 @@ async def test_rust_responses_rejects_a_configured_python_cache( @pytest.mark.asyncio @pytest.mark.parametrize("cancel_lookup", (False, True)) -async def test_python_cache_operations_stay_in_the_rust_callers_task( +@pytest.mark.parametrize("stream", (False, True), ids=("response", "stream")) +async def test_python_cache_operations_stay_in_the_rust_callers_task_except_the_unary_store( recording_server: RecordingServer, cancel_lookup: bool, + stream: bool, ) -> None: from litellm.caching.base_cache import BaseCache from litellm.caching.in_memory_cache import InMemoryCache @@ -278,14 +286,19 @@ async def test_python_cache_operations_stay_in_the_rust_callers_task( return storage.get_cache(key, **kwargs) async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None: - assert asyncio.current_task() is caller + writer: Final = asyncio.current_task() + if stream: + assert writer is caller + else: + assert writer is not caller + assert writer in _PENDING_CACHE_WRITES await asyncio.sleep(0) storage.set_cache(key, value, **kwargs) litellm.cache = Cache(_backend=CallerCache()) if cancel_lookup: recording_server.expected_requests = 0 - task: Final = asyncio.create_task(invoke("messages", recording_server, {})) + task: Final = asyncio.create_task(invoke("messages", recording_server, {"stream": stream})) await asyncio.wait_for(entered.wait(), timeout=5) task.cancel() with pytest.raises(asyncio.CancelledError): @@ -295,13 +308,82 @@ async def test_python_cache_operations_stay_in_the_rust_callers_task( assert len(recording_server.requests) == 0 assert storage.cache_dict == {} return - first: Final = await invoke("messages", recording_server, {}) - second: Final = await invoke("messages", recording_server, {}) - assert payload(first) == payload(second) + first: Final = await invoke("messages", recording_server, {"stream": stream}) + first_payload: Final = await collect(first) if stream else payload(first) + await asyncio.gather(*tuple(_PENDING_CACHE_WRITES)) + second: Final = await invoke("messages", recording_server, {"stream": stream}) + second_payload: Final = await collect(second) if stream else payload(second) + assert first_payload == second_payload assert cache_key(second) assert len(recording_server.requests) == 1 +@pytest.mark.asyncio +async def test_rust_unary_messages_store_keeps_the_hidden_params_stamped_after_the_call( + recording_server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.caching.in_memory_cache import InMemoryCache + + storage: Final = InMemoryCache() + monkeypatch.setattr(litellm, "cache", Cache(_backend=storage)) + first: Final = await invoke("messages", recording_server, {}) + assert isinstance(first, dict) + assert storage.cache_dict == {} + add_retry_headers_to_response( + _add_headers_to_response(first, {"x-litellm-model-group": "group"}), attempted_retries=0 + ) + get_hidden_params_dict(first)["model_id"] = "deployment" + await asyncio.gather(*tuple(_PENDING_CACHE_WRITES)) + (entry,) = storage.cache_dict.values() + stored: Final = TypeAdapter(dict[str, object]).validate_python( + TypeAdapter(dict[str, object]).validate_python(entry)["response"] + ) + assert stored["_hidden_params"] == { + "additional_headers": {"x-litellm-model-group": "group", "x-litellm-attempted-retries": 0}, + "model_id": "deployment", + } + assert get_hidden_params_dict(first)["additional_headers"] == { + "x-litellm-rust": "true", + "x-litellm-model-group": "group", + "x-litellm-attempted-retries": 0, + } + hit: Final = await invoke("messages", recording_server, {}) + assert get_hidden_params_dict(hit)["cache_key"] == cache_key(hit) + assert len(recording_server.requests) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("native", (False, True), ids=("python", "rust")) +async def test_messages_cache_stores_the_response_a_yielding_deployment_hook_returned( + recording_server: RecordingServer, native: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.caching.in_memory_cache import InMemoryCache + from litellm.integrations.custom_logger import CustomLogger + + blocks: Final = TypeAdapter(list[AnthropicResponseContentBlockText]) + + class Redact(CustomLogger): + async def async_post_call_success_deployment_hook( + self, request_data: dict[str, object], response: LLMResponseTypes, call_type: CallTypes | None + ) -> LLMResponseTypes | None: + await asyncio.sleep(0.05) + assert isinstance(response, dict) + return {**response, "content": [{"type": "text", "text": "redacted"}]} # pyright: ignore[reportReturnType] # messages hooks see and return the plain dict the route serves + + monkeypatch.setenv("LITELLM_RUST", "1" if native else "0") + storage: Final = InMemoryCache() + monkeypatch.setattr(litellm, "cache", Cache(_backend=storage)) + monkeypatch.setattr(litellm, "callbacks", [Redact()]) + first: Final = await invoke("messages", recording_server, {}, native=native) + assert isinstance(first, dict) + assert [block.text for block in blocks.validate_python(first["content"])] == ["redacted"] + await asyncio.gather(*tuple(_PENDING_CACHE_WRITES)) + hit: Final = await invoke("messages", recording_server, {}, native=native) + assert isinstance(hit, dict) + assert [block.text for block in blocks.validate_python(hit["content"])] == ["redacted"] + assert len(recording_server.requests) == 1 + + @pytest.mark.asyncio @pytest.mark.parametrize("namespace_source", ("cache", "metadata")) async def test_rust_messages_legacy_cache_honors_request_namespaces( diff --git a/tests/test_litellm_rust/messages/test_callbacks.py b/tests/test_litellm_rust/messages/test_callbacks.py index a3c03898a12..0550b530913 100644 --- a/tests/test_litellm_rust/messages/test_callbacks.py +++ b/tests/test_litellm_rust/messages/test_callbacks.py @@ -10,14 +10,15 @@ from pydantic import TypeAdapter import litellm from litellm.integrations.argilla import ArgillaLogger from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.types.integrations.argilla import ArgillaItem from litellm.types.utils import ModelResponse, StandardLoggingPayload -from tests.test_litellm_rust.messages.data import BAD_REQUEST, MESSAGES_RESPONSE, RATE_LIMITED, call +from tests.test_litellm_rust.messages.data import BAD_REQUEST, MESSAGES_RESPONSE, RATE_LIMITED, REQUEST, call from tests.test_litellm_rust.support.callback_recorder import HookRecorder, RecordingLogger, request_body from tests.test_litellm_rust.support.isolation import rebound from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, relaxed from tests.test_litellm_rust.support.results import response_id, with_id -from tests.test_litellm_rust.support.routing import assert_served_natively +from tests.test_litellm_rust.support.routing import CALL_ID, assert_served_natively pytestmark = pytest.mark.requires_rust_extension @@ -110,6 +111,33 @@ async def test_pre_api_call_sees_the_wire_body(server: RecordingServer) -> None: assert request_body(pre_call) == server.requests[0].body +@pytest.mark.asyncio +async def test_pre_api_call_of_a_retried_call_keeps_the_provider_out_of_the_litellm_params( + server: RecordingServer, +) -> None: + recorder: Final = RecordingLogger() + logging_obj: Final = Logging( + model=REQUEST["model"], + messages=[], + stream=False, + start_time=datetime.now(), + call_type="anthropic_messages", + litellm_call_id=CALL_ID, + function_id=CALL_ID, + dynamic_input_callbacks=[recorder], + ) + + await call(server, litellm_logging_obj=logging_obj) + await call(server, litellm_logging_obj=logging_obj) + + pre_calls: Final = recorder.wait_for("log_pre_api_call", count=2) + + assert_served_natively(server, count=2) + for pre_call in pre_calls: + assert isinstance(pre_call.kwargs, dict) + assert "custom_llm_provider" not in _OBJECT.validate_python(pre_call.kwargs["litellm_params"]) + + @pytest.mark.asyncio async def test_pre_api_call_edits_reach_the_provider(server: RecordingServer) -> None: class Edit(CustomLogger): @@ -140,6 +168,47 @@ async def test_pre_api_call_sees_the_shaped_optional_params(server: RecordingSer assert optional_params["max_tokens"] == sent["max_tokens"] +@pytest.mark.asyncio +async def test_success_event_logs_the_requested_reasoning_effort_not_its_translation(server: RecordingServer) -> None: + recorder: Final = RecordingLogger() + + with rebound(litellm, "callbacks", [recorder]): + response: Final = await call(server, reasoning_effort="low") + success: Final = await recorder.wait_for_async("async_log_success_event") + + assert_served_natively(server, response) + sent: Final = server.requests[0].body + assert isinstance(sent, dict) + assert sent["output_config"] == {"effort": "low"} + assert isinstance(success[0].kwargs, dict) + logged: Final = _OBJECT.validate_python(success[0].kwargs["standard_logging_object"]) + assert logged["model_parameters"] == {"max_tokens": 64, "stream": False, "reasoning_effort": "low"} + + +@pytest.mark.asyncio +async def test_success_event_logs_the_sent_messages_and_the_callers_params(server: RecordingServer) -> None: + recorder: Final = RecordingLogger() + tool_use: Final = {"type": "tool_use", "id": "functions.Bash:0", "name": "bash", "input": {"command": "ls"}} + messages: Final = [ + {"role": "user", "content": "List /tmp"}, + {"role": "assistant", "content": [{"type": "text", "text": " \n"}, tool_use]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "functions.Bash:0", "content": "a"}]}, + ] + + with rebound(litellm, "callbacks", [recorder]): + response: Final = await call(server, messages=messages, temperature=1) + success: Final = await recorder.wait_for_async("async_log_success_event") + + assert_served_natively(server, response) + sent: Final = server.requests[0].body + assert isinstance(sent, dict) + assert sent["messages"][1]["content"] == [{**tool_use, "id": "functions_Bash_0"}] + assert isinstance(success[0].kwargs, dict) + logged: Final = _OBJECT.validate_python(success[0].kwargs["standard_logging_object"]) + assert logged["messages"] == sent["messages"] + assert repr(_OBJECT.validate_python(logged["model_parameters"])["temperature"]) == "1" + + @pytest.mark.asyncio async def test_logging_hook_and_redaction_edits_reach_the_async_success_event(server: RecordingServer) -> None: class Rewrite(HookRecorder): @@ -215,6 +284,23 @@ async def test_pre_call_deployment_rejection_logs_failure_without_provider_io(se assert not {"log_pre_api_call", "log_success_event", "async_log_success_event"} & set(recorder.names) +@pytest.mark.asyncio +async def test_non_string_header_is_a_bad_request_refused_before_the_provider_call(server: RecordingServer) -> None: + recorder: Final = RecordingLogger() + + with rebound(litellm, "callbacks", [recorder]): + with pytest.raises( + litellm.BadRequestError, match=r"invalid extra_headers: header x-org must be a string, got int" + ) as raised: + await call(server, extra_headers={"x-org": 12345}) + failures: Final = await recorder.wait_for_async("async_log_failure_event") + + assert raised.value.status_code == 400 + assert not server.requests + assert "log_pre_api_call" not in recorder.names + assert _OBJECT.validate_python(failures[0].kwargs)["exception"] is raised.value + + @pytest.mark.asyncio async def test_pre_call_deployment_hook_can_delete_keywords(server: RecordingServer) -> None: class DropKeywords(CustomLogger): diff --git a/tests/test_litellm_rust/messages/test_dispatch.py b/tests/test_litellm_rust/messages/test_dispatch.py index 6d65f289524..b580eca2e42 100644 --- a/tests/test_litellm_rust/messages/test_dispatch.py +++ b/tests/test_litellm_rust/messages/test_dispatch.py @@ -2,10 +2,12 @@ import inspect from collections.abc import Awaitable, Callable, Mapping from typing import Final, cast # noqa: TID251 # narrows legacy callable signatures for inspect +import httpx import pytest import litellm from litellm.caching.caching import Cache +from litellm.constants import COMPLETION_HTTP_FALLBACK_SECONDS from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.websearch_interception.handler import WebSearchInterceptionLogger from litellm.llms.anthropic.messages.native import handler as python_messages @@ -172,8 +174,25 @@ def test_native_receives_normalized_request_and_original_call_shape() -> None: assert call_kwargs["litellm_metadata"] is metadata -@pytest.mark.parametrize("stream, expected_timeout", ((True, 0.5), (False, 30.0))) -def test_native_streams_are_bounded_by_stream_timeout(stream: bool, expected_timeout: float) -> None: +@pytest.mark.parametrize( + ("kwargs", "expected_timeout"), + ( + pytest.param({"stream": True, "timeout": 30.0, "stream_timeout": 0.5}, 0.5, id="stream-timeout-bounds-stream"), + pytest.param( + {"stream": False, "timeout": 30.0, "stream_timeout": 0.5}, 30.0, id="stream-timeout-ignored-unary" + ), + pytest.param({"request_timeout": 5}, 5.0, id="request-timeout"), + pytest.param({"timeout": 30.0, "request_timeout": 5}, 30.0, id="timeout-before-request-timeout"), + pytest.param( + {"stream": True, "request_timeout": 5, "stream_timeout": 0.5}, 0.5, id="stream-timeout-over-request" + ), + pytest.param({"timeout": httpx.Timeout(60.0, connect=2.0)}, 60.0, id="httpx-timeout-read"), + pytest.param( + {"timeout": httpx.Timeout(None, connect=2.0)}, COMPLETION_HTTP_FALLBACK_SECONDS, id="httpx-timeout-no-read" + ), + ), +) +def test_native_timeout_resolves_like_python(kwargs: Mapping[str, object], expected_timeout: float) -> None: captured: Final[list[NativeCall]] = [] def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: rejected fallback @@ -185,7 +204,7 @@ def test_native_streams_are_bounded_by_stream_timeout(stream: bool, expected_tim _ = _DISPATCH.run( (16, MESSAGES, "anthropic/claude-sonnet-4-5"), - {"stream": stream, "timeout": 30.0, "stream_timeout": 0.5}, + kwargs, python=python, binding=messages_binding(native), native=lambda hook, request, call_args, call_kwargs: hook(request), @@ -456,7 +475,6 @@ STRICT_TOOL: Final = {"name": "a", "input_schema": {"type": "object"}, "strict": {}, ), (Feature.CALLER_CLIENT, "anthropic/claude-sonnet-4-5", {"client": object()}, {}), - (Feature.TIMEOUT_VARIANTS, "anthropic/claude-sonnet-4-5", {"request_timeout": 5}, {}), (Feature.AGENTIC_LOOP_HOOK, "anthropic/claude-sonnet-4-5", {}, {"callbacks": [_Agentic()]}), (Feature.RESPONSE_CACHE, "anthropic/claude-sonnet-4-5", {}, {"cache": Cache()}), (Feature.TLS_POLICY, "anthropic/claude-sonnet-4-5", {}, {"ssl_security_level": "DEFAULT@SECLEVEL=1"}), @@ -528,6 +546,38 @@ def test_each_detected_feature_reaches_its_rule( assert served(kwargs) is python_result +@pytest.mark.parametrize( + "kwargs", + ( + pytest.param({"shared_session": object()}, id="proxy-shared-session"), + pytest.param({"stream": True, "timeout": 600.0, "stream_timeout": 600.0}, id="proxy-stream-timeout"), + pytest.param({"request_timeout": 5}, id="request-timeout"), + pytest.param({"timeout": httpx.Timeout(60.0, connect=2.0)}, id="httpx-timeout"), + ), +) +def test_proxy_request_shape_stays_native(kwargs: Mapping[str, object]) -> None: + rules: Final[Rules] = ( + RouteRule( + Route.MESSAGES, Rollout.PYTHON_ONLY, features=frozenset({Feature.CALLER_CLIENT, Feature.TIMEOUT_VARIANTS}) + ), + RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED), + ) + native_result: Final = response() + + def python(*call_args: object, **call_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: rejected fallback + pytest.fail("A proxy-shaped request must not fall back to Python") + + served: Final = _DISPATCH.run( + (16, MESSAGES, "anthropic/claude-sonnet-4-5"), + kwargs, + python=python, + binding=messages_binding(lambda request: native_result), + native=lambda hook, request, call_args, native_kwargs: hook(request), + rules=rules, + ) + assert served is native_result + + @pytest.mark.asyncio @pytest.mark.parametrize( "kwargs", diff --git a/tests/test_litellm_rust/messages/test_requests.py b/tests/test_litellm_rust/messages/test_requests.py index 92031e72aee..3e486bbe9e6 100644 --- a/tests/test_litellm_rust/messages/test_requests.py +++ b/tests/test_litellm_rust/messages/test_requests.py @@ -4,7 +4,7 @@ import pytest import litellm from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_RESPONSE, call -from tests.test_litellm_rust.support.recording_server import RecordingServer +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec from tests.test_litellm_rust.support.routing import assert_served_natively pytestmark = pytest.mark.requires_rust_extension @@ -57,6 +57,21 @@ async def test_replayed_tool_ids_are_sanitized_before_they_reach_the_provider(se ] +@pytest.mark.asyncio +async def test_a_success_body_that_is_not_a_message_passes_through(server: RecordingServer) -> None: + overloaded: Final = {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}} + server.default_response = ResponseSpec(body=overloaded) + + response: Final = await call(server) + + assert_served_natively(server, response) + assert isinstance(response, dict) + assert {name: value for name, value in response.items() if name != "_hidden_params"} == { + **overloaded, + "stop_sequence": None, + } + + @pytest.mark.asyncio async def test_provider_specific_headers_scoped_to_anthropic_reach_the_wire(server: RecordingServer) -> None: await call( diff --git a/tests/test_litellm_rust/messages/test_streaming.py b/tests/test_litellm_rust/messages/test_streaming.py index e57631f48b9..99b6101da4e 100644 --- a/tests/test_litellm_rust/messages/test_streaming.py +++ b/tests/test_litellm_rust/messages/test_streaming.py @@ -1,10 +1,14 @@ +import asyncio +import json from collections.abc import AsyncIterator, Iterator from typing import Final import pytest +import litellm from tests.test_litellm_rust.messages.data import MESSAGES_EVENTS, MESSAGES_RESPONSE, call from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging +from tests.test_litellm_rust.support.isolation import rebound from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, relaxed from tests.test_litellm_rust.support.routing import assert_served_natively @@ -68,3 +72,59 @@ async def test_stream_closed_early_logs_success_once(server: RecordingServer) -> assert len(await recorder.wait_for_async("async_log_success_event")) == 1 with pytest.raises(StopAsyncIteration): await anext(stream) + + +@pytest.mark.asyncio +async def test_stream_cancelled_while_waiting_on_the_provider_logs_success_for_what_it_read( + server: RecordingServer, +) -> None: + server.enqueue(ResponseSpec(body=None, events=MESSAGES_EVENTS, chunk_delay=1)) + recorder: Final = RecordingLogger() + + stream: Final = await call(server, stream=True, callbacks=[recorder]) + assert isinstance(stream, AsyncIterator) + await anext(stream) + with pytest.raises(TimeoutError): + await asyncio.wait_for(anext(stream), timeout=0.2) + + success: Final = await recorder.wait_for_async("async_log_success_event") + assert len(success) == 1 + assert success[0].response.usage.prompt_tokens == MESSAGES_RESPONSE["usage"]["input_tokens"] + assert "log_failure_event" not in recorder.names + + +async def _streamed_pre_call_input(server: RecordingServer) -> object: + server.enqueue(STREAM) + recorder: Final = RecordingLogger() + with rebound(litellm, "callbacks", [recorder]): + stream: Final = await call( + server, stream=True, system="Be brief", metadata={"user_id": "user-1"}, stop_sequences=["END"] + ) + assert isinstance(stream, AsyncIterator) + _ = [chunk async for chunk in stream] + pre_call: Final = recorder.wait_for("log_pre_api_call")[0].kwargs + assert isinstance(pre_call, dict) + return pre_call["input"] + + +@pytest.mark.asyncio +async def test_stream_logs_the_request_body_as_input_like_python(server: RecordingServer) -> None: + logged: Final = await _streamed_pre_call_input(server) + + assert_served_natively(server) + assert logged == [ + { + "role": "user", + "content": json.dumps( + { + "messages": [{"role": "user", "content": "Hello"}], + "max_tokens": 64, + "model": "claude-sonnet-5", + "metadata": {"user_id": "user-1"}, + "stop_sequences": ["END"], + "stream": True, + "system": "Be brief", + } + ), + } + ] diff --git a/tests/test_litellm_rust/ocr/test_lifecycle.py b/tests/test_litellm_rust/ocr/test_lifecycle.py index 611316fb52c..c3846d894bb 100644 --- a/tests/test_litellm_rust/ocr/test_lifecycle.py +++ b/tests/test_litellm_rust/ocr/test_lifecycle.py @@ -211,8 +211,12 @@ async def test_deferred_logging_requires_release_and_runs_at_most_once(server: R response: Final = await call_aocr(server, litellm_logging_obj=logger) await drain_logging() assert "async_log_success_event" not in recorder.names + assert logger.deferred_success_result is not None + assert logger.deferred_success_result[0] is response + assert not logger.async_success_log_enqueued ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(logger, blocked) ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(logger, blocked) + assert logger.async_success_log_enqueued is not blocked await drain_logging() events: Final = tuple(event for event in recorder.events if event.name == "async_log_success_event") assert len(events) == int(not blocked) @@ -220,6 +224,39 @@ async def test_deferred_logging_requires_release_and_runs_at_most_once(server: R assert events[0].response is response +@pytest.mark.asyncio +@pytest.mark.parametrize("blocked", [False, True]) +async def test_deferred_sync_success_callbacks_wait_for_release(server: RecordingServer, blocked: bool) -> None: + finished: Final = threading.Event() + observations: Final = [] + + class Observe(CustomLogger): + def log_success_event(self, kwargs, response_obj, start_time, end_time): + observations.append(response_obj) + finished.set() + + logger: Final = Logging( + model="mistral-ocr-latest", + messages=[], + stream=False, + call_type="aocr", + start_time=datetime.datetime.now(), + litellm_call_id="deferred-sync", + function_id="deferred-sync", + dynamic_success_callbacks=[Observe()], + ) + logger._defer_async_logging = True + response: Final = await call_aocr(server, litellm_logging_obj=logger) + await drain_logging() + assert not await asyncio.to_thread(finished.wait, 0.2) + assert observations == [] + ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(logger, blocked) + ProxyBaseLLMRequestProcessing._flush_deferred_async_logging(logger, blocked) + await drain_logging() + assert await asyncio.to_thread(finished.wait, 0.5 if blocked else 5) is not blocked + assert observations == ([] if blocked else [response]) + + def test_sync_success_uses_executor_and_copied_caller_context(server: RecordingServer) -> None: context: Final = ContextVar("sync-lifecycle", default="missing") context.set("caller") @@ -454,7 +491,7 @@ async def test_empty_callbacks_run_deployment_hooks_and_defer_like_the_python_cl assert logger.model_call_details["first_api_call_start_time"] <= logger.model_call_details["end_time"] assert logger.model_call_details["response_cost"] == 0 else: - assert getattr(logger, "_native_pending_logging", None) is not None + assert logger.enqueue_deferred_logging is not None assert "end_time" not in logger.model_call_details diff --git a/tests/test_litellm_rust/responses/test_dispatch.py b/tests/test_litellm_rust/responses/test_dispatch.py index 4498dbdf39d..19255379dcf 100644 --- a/tests/test_litellm_rust/responses/test_dispatch.py +++ b/tests/test_litellm_rust/responses/test_dispatch.py @@ -260,7 +260,6 @@ ENCODED_ID: Final = ResponsesAPIRequestUtils._build_encrypted_item_id("deploymen (Feature.MOCK_RESPONSE, {"mock_response": "hi"}, {}), (Feature.PROXY_MCP_TOOLS, {"tools": [{"type": "mcp", "server_label": "g", "server_url": "litellm_proxy"}]}, {}), (Feature.CALLER_CLIENT, {"client": object()}, {}), - pytest.param(Feature.CALLER_CLIENT, {"shared_session": object()}, {}, id="shared-session"), pytest.param(Feature.EXTRA_BODY, {"extra_body": {"stream": True}}, {}, id="extra-body-stream"), pytest.param( Feature.PROVIDER_AFFINITY, diff --git a/tests/test_litellm_rust/test_public_entries.py b/tests/test_litellm_rust/test_public_entries.py index 2f5c33e3ced..24a3419c5a0 100644 --- a/tests/test_litellm_rust/test_public_entries.py +++ b/tests/test_litellm_rust/test_public_entries.py @@ -1,3 +1,4 @@ +import asyncio import inspect from collections.abc import Awaitable, Callable, Iterator, Mapping from contextlib import contextmanager @@ -29,7 +30,10 @@ Served: TypeAlias = Callable[[object], bool] @dataclass(frozen=True, slots=True) class Entry: - """``python`` is ``None`` for a route with no Python implementation, which only RUST_REQUIRED rules serve.""" + """``python`` is ``None`` for a route with no Python implementation, which only RUST_REQUIRED rules serve. + + ``executor_hop`` is whether the route's Python async entry left the loop through run_in_executor before it built the request. + """ route: Route sync: Callable[..., object] @@ -40,6 +44,7 @@ class Entry: arguments: Mapping[str, object] positional: int python: Served | None + executor_hop: bool def args(self) -> tuple[object, ...]: return tuple(self.arguments.values())[: self.positional] @@ -79,6 +84,7 @@ ENTRIES: Final = ( arguments={"model": "openai/gpt-4o", "messages": MESSAGES, "mock_response": "python"}, positional=2, python=_chat, + executor_hop=True, ), Entry( route=Route.EMBEDDINGS, @@ -90,6 +96,7 @@ ENTRIES: Final = ( arguments={"model": "openai/text-embedding-3-small", "input": ["hi"], "mock_response": [0.5]}, positional=2, python=_embedding, + executor_hop=True, ), Entry( route=Route.RESPONSES, @@ -101,6 +108,7 @@ ENTRIES: Final = ( arguments={"input": "hi", "model": "openai/gpt-4o", "mock_response": "python"}, positional=2, python=_responses, + executor_hop=True, ), Entry( route=Route.MESSAGES, @@ -117,6 +125,7 @@ ENTRIES: Final = ( }, positional=0, python=_messages, + executor_hop=True, ), Entry( route=Route.OCR, @@ -128,6 +137,7 @@ ENTRIES: Final = ( arguments={"model": "mistral/mistral-ocr-latest", "document": {"type": "file", "file": b"pdf"}}, positional=2, python=None, + executor_hop=False, ), ) @@ -285,3 +295,27 @@ async def test_sync_entry_with_flag_true_raises_setup_errors_at_await(entry: Ent assert inspect.iscoroutine(pending) with pytest.raises(RuntimeError, match="is unavailable"): await pending + + +@pytest.mark.parametrize("entry", ENTRIES, ids=entry_id) +async def test_async_entry_runs_native_through_the_client_wrapper_steps(entry: Entry) -> None: + scheduled: Final[list[str]] = [] + + async def failing(native_call: NativeCall) -> ModelResponse: + scheduled.append("native") + raise litellm.Timeout(message="timed out", model="model", llm_provider="openai") + + async def started_before() -> None: + scheduled.append("task") + + entry.async_binding.override(failing) + task: Final = asyncio.create_task(started_before()) + try: + with scenario_rules(entry, SCENARIOS["required"]), pytest.raises(litellm.Timeout) as raised: + await entry.asynchronous(*entry.args(), **entry.kwargs(timeout=30.0, num_retries=0)) + finally: + entry.async_binding.reset() + await task + + assert (getattr(raised.value, "timeout", None), getattr(raised.value, "num_retries", None)) == (30.0, 0) + assert scheduled == (["task", "native"] if entry.executor_hop else ["native", "task"]) diff --git a/tests/unit/caching/test_request_redis_batch_post_call.py b/tests/unit/caching/test_request_redis_batch_post_call.py index 41d508a0eea..01feee0be54 100644 --- a/tests/unit/caching/test_request_redis_batch_post_call.py +++ b/tests/unit/caching/test_request_redis_batch_post_call.py @@ -424,7 +424,7 @@ async def test_spend_counter_increments_ride_the_native_post_call_batch(monkeypa with request_redis_batch_scope(native_factory=natives, native_send_factory=fake_send): # pyright: ignore[reportArgumentType] # test double await proxy_server._apply_spend_counter_increments(pending) - declared = list(spend.post_call) + declared = [(request, increments) for request, increments, _ in spend.post_call] assert natives.built[0].events == [] await flush_post_call_redis_batches() @@ -433,6 +433,27 @@ async def test_spend_counter_increments_ride_the_native_post_call_batch(monkeypa assert spend.commands == [] +@pytest.mark.asyncio +async def test_the_request_sees_its_committed_spend_charge_settle_whoever_sends_the_batch(monkeypatch): + from litellm.proxy import proxy_server + + spend = FakeSpendCounters(redis={}) + monkeypatch.setattr(proxy_server, "spend_counter_store", spend.store()) + + with request_redis_batch_scope(native_factory=FakeNatives(), native_send_factory=fake_send) as request: # pyright: ignore[reportArgumentType] # test double + await proxy_server._apply_spend_counter_increments([PendingSpendIncrement("spend:key:k1", 0.5)]) + without_committed_spend = request.post_call_settlements + await proxy_server._apply_spend_counter_increments( + [PendingSpendIncrement("spend:key:k1", 0.5), PendingSpendIncrement("spend:committed:key:k1", 0.5)] + ) + declared = [settlement.settled for settlement in request.post_call_settlements] + spend.settle_post_call() + sent = [settlement.settled for settlement in request.post_call_settlements] + + assert without_committed_spend == () + assert (declared, sent) == ([False], [True]) + + @pytest.mark.asyncio async def test_native_post_call_ops_go_out_with_the_post_call_flush_not_the_request_boundary(): natives = FakeNatives() diff --git a/tests/unit/integration_support/test_database_relay.py b/tests/unit/integration_support/test_database_relay.py index 98501f84058..de2bf1f5d2b 100644 --- a/tests/unit/integration_support/test_database_relay.py +++ b/tests/unit/integration_support/test_database_relay.py @@ -4,7 +4,7 @@ from typing import Final import pytest -from tests.integration._support.database_relay import TriggerScanner +from tests.integration._support.database_relay import StatementMarkerScanner, TriggerScanner TRIGGER: Final = b'SELECT "startTime" FROM "LiteLLM_SpendLogs"' @@ -26,3 +26,55 @@ def test_trigger_scanner_reports_a_match_once() -> None: scanner: Final = TriggerScanner(TRIGGER) assert scanner.feed(b"x" + TRIGGER + b"y") assert not scanner.feed(b"z") + + +STARTUP: Final = b"\x00\x00\x00\x08\x00\x03\x00\x00" + + +def _message(kind: bytes, payload: bytes) -> bytes: + return kind + (len(payload) + 4).to_bytes(4, "big") + payload + + +PARSE: Final = _message(b"P", b"s1\x00" + TRIGGER + b"\x00\x00\x00") +BIND: Final = _message(b"B", b"\x00s1\x00\x00\x00\x00\x00\x00\x00") + + +@pytest.mark.parametrize("split_at", range(1, len(STARTUP + PARSE + BIND))) +def test_statement_markers_survive_every_protocol_frame_split(split_at: int) -> None: + scanner: Final = StatementMarkerScanner((TRIGGER,)) + stream: Final = STARTUP + PARSE + BIND + assert (*scanner.feed(stream[:split_at]), *scanner.feed(stream[split_at:])) == ( + frozenset({TRIGGER}), + frozenset({TRIGGER}), + ) + + +def test_statement_markers_are_local_to_a_connection_and_forget_closed_statements() -> None: + scanner: Final = StatementMarkerScanner((TRIGGER,)) + other: Final = StatementMarkerScanner((TRIGGER,)) + assert tuple(scanner.feed(STARTUP + PARSE)) == (frozenset({TRIGGER}),) + assert tuple(other.feed(STARTUP + BIND)) == (frozenset(),) + assert tuple(scanner.feed(_message(b"C", b"Ss1\x00") + BIND)) == (frozenset(),) + + +def test_replaced_and_unnamed_statements_do_not_reuse_old_markers() -> None: + scanner: Final = StatementMarkerScanner((TRIGGER,)) + assert tuple(scanner.feed(STARTUP + PARSE)) == (frozenset({TRIGGER}),) + replacement: Final = _message(b"P", b"s1\x00SELECT 1\x00\x00\x00") + unnamed: Final = _message(b"P", b"\x00" + TRIGGER + b"\x00\x00\x00") + assert tuple(scanner.feed(replacement + BIND + unnamed)) == (frozenset(), frozenset(), frozenset({TRIGGER})) + + +def test_statement_scanner_rejects_invalid_frame_lengths() -> None: + scanner: Final = StatementMarkerScanner((TRIGGER,)) + assert tuple(scanner.feed(STARTUP)) == () + with pytest.raises(ValueError, match="Invalid PostgreSQL message length"): + tuple(scanner.feed(b"P\x00\x00\x00\x03")) + + +def test_statement_scanner_matches_bound_parameter_markers() -> None: + scanner: Final = StatementMarkerScanner((TRIGGER,)) + bind: Final = _message( + b"B", b"\x00s2\x00\x00\x00\x00\x01" + len(TRIGGER).to_bytes(4, "big") + TRIGGER + b"\x00\x00" + ) + assert tuple(scanner.feed(STARTUP + bind)) == (frozenset({TRIGGER}),) diff --git a/tests/unit/integrations/generic_api/test_generic_api_callback.py b/tests/unit/integrations/generic_api/test_generic_api_callback.py index afaae02321d..8f1e79562db 100644 --- a/tests/unit/integrations/generic_api/test_generic_api_callback.py +++ b/tests/unit/integrations/generic_api/test_generic_api_callback.py @@ -20,6 +20,8 @@ from litellm.integrations.gcs_pubsub.pub_sub import * from litellm.integrations.generic_api.generic_api_callback import LOG_FORMAT_TYPES, GenericAPILogger from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys from litellm.types.utils import StandardLoggingPayload from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -795,3 +797,27 @@ async def test_generic_api_concurrent_flushes_deliver_each_event_once(log_format task.cancel() await asyncio.gather(*owned_tasks, return_exceptions=True) await handler.close() + + +@pytest.mark.asyncio +async def test_v1_payload_of_a_chat_row_never_shows_the_committed_spend_auth_read() -> None: + counter_keys: Final = EntitySpendCounterKeys.of( + token="hashed-token", team_id="t1", user_id=None, org_id=None, end_user_id=None + ) + auth: Final = UserAPIKeyAuth( + api_key="hashed-token", + team_id="t1", + committed_budget_spend=CommittedBudgetSpend(counter_keys=counter_keys, key=0.123456, team=0.654321), + ) + generic_logger: Final = GenericAPILogger(endpoint="https://example.com/api/logs", flush_interval=3600) + + payload: Final = generic_logger._get_v1_logging_payload( + kwargs={"litellm_params": {"metadata": {"user_api_key_auth": auth}}}, + response_obj={"usage": {}}, + start_time=None, + end_time=None, + ) + + assert "team_id='t1'" in payload["metadata"] + assert "0.123456" not in payload["metadata"] + assert str(counter_keys.team) not in payload["metadata"] diff --git a/tests/unit/integrations/test_prometheus_carried_budget_state.py b/tests/unit/integrations/test_prometheus_carried_budget_state.py index e49e6baa8b0..9e006c20079 100644 --- a/tests/unit/integrations/test_prometheus_carried_budget_state.py +++ b/tests/unit/integrations/test_prometheus_carried_budget_state.py @@ -4,13 +4,22 @@ from request metadata. get_*_object only runs when that state is missing (custom auth, SDK callers, failure paths) """ +import asyncio +from collections.abc import Sequence from datetime import datetime, timezone +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest from prometheus_client import REGISTRY +from litellm.caching.redis_batch import ( + active_post_call_redis_batch, + flush_post_call_redis_batches, + request_redis_batch_scope, +) from litellm.integrations.prometheus import PrometheusLogger +from litellm.proxy import proxy_server from litellm.proxy._types import ( LiteLLM_BudgetTable, LiteLLM_OrganizationTable, @@ -24,11 +33,16 @@ from litellm.proxy.spend_tracking.carried_budget_state import ( carry_organization_budget_state, carry_team_and_user_budget_state, ) +from litellm.proxy.spend_tracking.spend_counter_batch import PendingSpendIncrement from litellm.types.proxy.carried_budget_state import ( + CommittedBudgetSpend, + EntitySpendCounterKeys, KeyBudgetSnapshot, TeamBudgetSnapshot, UserBudgetSnapshot, ) +from tests.unit.caching.test_redis_batch import FakeNatives, fake_send, make +from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters TEAM_RESET_AT = datetime(2026, 10, 1, tzinfo=timezone.utc) USER_RESET_AT = datetime(2026, 11, 1, tzinfo=timezone.utc) @@ -55,6 +69,15 @@ def prometheus_logger(): return PrometheusLogger() +@pytest.fixture(autouse=True) +def spend_counters(): + spend: Final = FakeSpendCounters(redis=None) + with patch( # test-quality-ok: prometheus reads the proxy_server spend counters at emit time, no injection seam + "litellm.proxy.proxy_server.spend_counter_store", spend.store() + ): + yield spend + + @pytest.fixture def getters(): """Every response-path object getter, patched where prometheus imports them from.""" @@ -137,6 +160,7 @@ def _request_metadata(token: UserAPIKeyAuth) -> dict: return { **LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(token), **carried_budget_metadata(token), + "user_api_key_auth": token, } @@ -198,6 +222,499 @@ async def test_authed_request_sets_every_gauge_without_any_object_getter(prometh prometheus_logger.litellm_user_budget_remaining_hours_metric.labels().set.assert_called_once() +_COUNTER_KEYS: Final = ("spend:key:hashed", "spend:team:t1", "spend:user:u1", "spend:org:o1") +_TWIN_KEYS: Final = ( + "spend:committed:key:hashed", + "spend:committed:team:t1", + "spend:committed:user:u1", + "spend:committed:org:o1", +) + + +def _store(spend: FakeSpendCounters, counter_keys: Sequence[str], values: Sequence[float]) -> None: + tier: Final = spend.memory if spend.redis is None else spend.redis + tier.update(zip(counter_keys, values, strict=True)) + + +def _store_counters(spend: FakeSpendCounters, key: float, team: float, user: float, org: float) -> None: + _store(spend, _COUNTER_KEYS, (key, team, user, org)) + + +def _store_committed(spend: FakeSpendCounters, key: float, team: float, user: float, org: float) -> None: + _store(spend, _TWIN_KEYS, (key, team, user, org)) + + +def _remaining(prometheus_logger: PrometheusLogger) -> dict[str, float]: + return { + "key": prometheus_logger.litellm_remaining_api_key_budget_metric.labels().set.call_args[0][0], + "team": prometheus_logger.litellm_remaining_team_budget_metric.labels().set.call_args[0][0], + "user": prometheus_logger.litellm_remaining_user_budget_metric.labels().set.call_args[0][0], + "org": prometheus_logger.litellm_remaining_org_budget_metric.labels().set.call_args[0][0], + } + + +def _token_committed_at_auth() -> UserAPIKeyAuth: + token = _authed_token() + token.committed_budget_spend = CommittedBudgetSpend( + counter_keys=EntitySpendCounterKeys.of( + token="hashed", team_id="t1", user_id="u1", org_id="o1", end_user_id=None + ), + key=3.0, + team=30.0, + user=8.0, + org=60.0, + ) + return token + + +@pytest.mark.asyncio +async def test_gauges_use_the_committed_spend_at_emit_when_other_requests_finished_after_auth( + prometheus_logger, getters, spend_counters +): + """This request's 2.0 plus 1.0 from another request that finished first, so the gauge is never set back.""" + _stub_gauges(prometheus_logger) + _store_committed(spend_counters, key=6.0, team=33.0, user=11.0, org=63.0) + + await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth())) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(4.0), + "team": pytest.approx(267.0), + "user": pytest.approx(39.0), + "org": pytest.approx(437.0), + } + + +@pytest.mark.asyncio +async def test_gauges_add_the_cost_to_another_requests_charge_while_the_own_charge_is_not_in_yet( + prometheus_logger, getters, spend_counters +): + """1.0 from another request reached the counters during this long request, this request's 2.0 has not.""" + _stub_gauges(prometheus_logger) + _store_committed(spend_counters, key=4.0, team=31.0, user=9.0, org=61.0) + + await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth())) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(4.0), + "team": pytest.approx(267.0), + "user": pytest.approx(39.0), + "org": pytest.approx(437.0), + } + + +@pytest.mark.asyncio +async def test_gauges_read_no_committed_counter_auth_did_not_name_as_budgeted( + prometheus_logger, getters, spend_counters +): + _stub_gauges(prometheus_logger) + _store_committed(spend_counters, key=9.0, team=99.0, user=19.0, org=99.0) + + await _emit(prometheus_logger, _request_metadata(_authed_token())) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(7.0), + "team": pytest.approx(278.0), + "user": pytest.approx(43.0), + "org": 458.0, + } + + +def _billed_row_metadata() -> dict: + return {key: value for key, value in _request_metadata(_authed_token()).items() if key != "user_api_key_auth"} + + +@pytest.mark.asyncio +async def test_row_with_no_auth_adds_its_cost_to_the_committed_spend_of_the_entities_it_is_charged_to( + prometheus_logger, getters, spend_counters +): + """Counter increments flush after the success callbacks, so the row's own charge is not in the counters yet.""" + _stub_gauges(prometheus_logger) + _store_committed(spend_counters, key=6.0, team=33.0, user=11.0, org=63.0) + + await _emit(prometheus_logger, _billed_row_metadata()) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(2.0), + "team": pytest.approx(265.0), + "user": pytest.approx(37.0), + "org": 435.0, + } + + +@pytest.mark.asyncio +async def test_row_with_no_auth_and_no_committed_counters_keeps_the_stored_spend( + prometheus_logger, getters, spend_counters +): + _stub_gauges(prometheus_logger) + + await _emit(prometheus_logger, _billed_row_metadata()) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(7.0), + "team": pytest.approx(278.0), + "user": pytest.approx(43.0), + "org": 458.0, + } + + +@pytest.mark.asyncio +async def test_slow_committed_read_on_a_failed_request_is_bounded_and_keeps_the_auth_time_spend( + prometheus_logger, getters, monkeypatch +): + async def never_answers(counter: str) -> None: + await asyncio.sleep(30) + + monkeypatch.setattr( + "litellm.proxy.proxy_server.spend_counter_store", + FakeSpendCounters(redis={}, batch_reads=False, read_hook=never_answers).store(), + ) + monkeypatch.setenv("PROMETHEUS_BUDGET_METRICS_PER_REQUEST_TIMEOUT", "0.01") + _stub_gauges(prometheus_logger) + + await asyncio.wait_for( + prometheus_logger.async_log_failure_event( + _failed_request(_request_metadata(_token_committed_at_auth())), None, KEY_RESET_AT, KEY_RESET_AT + ), + timeout=5, + ) + + prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(pytest.approx(440.0)) + + +@pytest.mark.asyncio +async def test_slow_committed_read_on_a_successful_request_still_sets_the_gauges_from_the_auth_time_spend( + prometheus_logger, getters, monkeypatch +): + async def never_answers(counter: str) -> None: + await asyncio.sleep(30) + + monkeypatch.setattr( + "litellm.proxy.proxy_server.spend_counter_store", + FakeSpendCounters(redis={}, batch_reads=False, read_hook=never_answers).store(), + ) + monkeypatch.setenv("PROMETHEUS_BUDGET_METRICS_PER_REQUEST_TIMEOUT", "0.5") + _stub_gauges(prometheus_logger) + + await asyncio.wait_for(_emit(prometheus_logger, _request_metadata(_token_committed_at_auth())), timeout=5) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(5.0), + "team": pytest.approx(268.0), + "user": pytest.approx(40.0), + "org": pytest.approx(438.0), + } + + +@pytest.fixture +def redis_spend_counters(): + spend: Final = FakeSpendCounters(redis={}) + with patch( # test-quality-ok: prometheus reads the proxy_server spend counters at emit time, no injection seam + "litellm.proxy.proxy_server.spend_counter_store", spend.store() + ): + yield spend + + +def _request_scope() -> request_redis_batch_scope: + return request_redis_batch_scope( + native_factory=FakeNatives(), # pyright: ignore[reportArgumentType] # test double + native_send_factory=fake_send, # pyright: ignore[reportArgumentType] # test double + ) + + +async def _declare_the_own_two_dollar_charge() -> None: + await proxy_server._apply_spend_counter_increments( + tuple(PendingSpendIncrement(counter_key, 2.0) for counter_key in (*_COUNTER_KEYS, *_TWIN_KEYS)) + ) + + +_BOTH_TWO_DOLLAR_CHARGES_COUNTED_ONCE: Final = { + "key": pytest.approx(3.0), + "team": pytest.approx(266.0), + "user": pytest.approx(38.0), + "org": pytest.approx(436.0), +} + + +@pytest.mark.asyncio +async def test_gauges_add_the_own_charge_to_a_same_size_charge_of_another_request_while_the_own_waits_in_the_pipeline( + prometheus_logger, getters, redis_spend_counters +): + """Another request's 2.0 reached the counters and this request's 2.0 is queued, so both count.""" + _stub_gauges(prometheus_logger) + _store_committed(redis_spend_counters, key=5.0, team=32.0, user=10.0, org=62.0) + + with _request_scope(): + await _declare_the_own_two_dollar_charge() + await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth())) + await flush_post_call_redis_batches() + + assert _remaining(prometheus_logger) == _BOTH_TWO_DOLLAR_CHARGES_COUNTED_ONCE + + +@pytest.mark.asyncio +async def test_gauges_do_not_add_the_own_charge_twice_once_the_pipeline_went_out( + prometheus_logger, getters, redis_spend_counters +): + """The 2.0 over the auth time spend is this request's own charge, settled before the gauges read it.""" + _stub_gauges(prometheus_logger) + _store_committed(redis_spend_counters, key=3.0, team=30.0, user=8.0, org=60.0) + + with _request_scope(): + await _declare_the_own_two_dollar_charge() + redis_spend_counters.settle_post_call() + await flush_post_call_redis_batches() + await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth())) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(5.0), + "team": pytest.approx(268.0), + "user": pytest.approx(40.0), + "org": pytest.approx(438.0), + } + + +@pytest.mark.asyncio +async def test_an_unrelated_post_call_flush_does_not_make_the_gauge_drop_its_own_unsent_charge( + prometheus_logger, getters, redis_spend_counters +): + """An unrelated write went out first, then another request's 2.0 reached the counters. This request's 2.0 counts + once whether it is undeclared, declared but unsent, or settled.""" + _stub_gauges(prometheus_logger) + _store_committed(redis_spend_counters, key=5.0, team=32.0, user=10.0, org=62.0) + unrelated_cache, _ = make() + metadata: Final = _request_metadata(_token_committed_at_auth()) + + with _request_scope(): + active_post_call_redis_batch(unrelated_cache).increment("unrelated-write", 1.0) + await flush_post_call_redis_batches() + await _emit(prometheus_logger, metadata) + undeclared: Final = _remaining(prometheus_logger) + await _declare_the_own_two_dollar_charge() + await _emit(prometheus_logger, metadata) + declared_unsent: Final = _remaining(prometheus_logger) + redis_spend_counters.settle_post_call() + await _emit(prometheus_logger, metadata) + settled: Final = _remaining(prometheus_logger) + await flush_post_call_redis_batches() + + assert {"undeclared": undeclared, "declared_unsent": declared_unsent, "settled": settled} == { + "undeclared": _BOTH_TWO_DOLLAR_CHARGES_COUNTED_ONCE, + "declared_unsent": _BOTH_TWO_DOLLAR_CHARGES_COUNTED_ONCE, + "settled": _BOTH_TWO_DOLLAR_CHARGES_COUNTED_ONCE, + } + + +@pytest.mark.asyncio +async def test_gauges_count_the_own_charge_once_when_it_settles_while_the_gauges_read_the_counters( + prometheus_logger, getters, redis_spend_counters +): + """Another request's 2.0 reached the counters, and this request's 2.0 goes out during the gauges' read.""" + _stub_gauges(prometheus_logger) + _store_committed(redis_spend_counters, key=5.0, team=32.0, user=10.0, org=62.0) + + async def settle_during_read(_counter: str) -> None: + redis_spend_counters.settle_post_call() + + redis_spend_counters.batch_reads = False + redis_spend_counters.read_hook = settle_during_read + + with _request_scope(): + await _declare_the_own_two_dollar_charge() + await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth())) + await flush_post_call_redis_batches() + + assert _remaining(prometheus_logger) == _BOTH_TWO_DOLLAR_CHARGES_COUNTED_ONCE + + +@pytest.mark.asyncio +async def test_gauges_without_redis_treat_a_charge_equal_to_the_own_as_the_own( + prometheus_logger, getters, spend_counters +): + """Without Redis the charge increments in memory at once, so a request scope says nothing about it.""" + _stub_gauges(prometheus_logger) + _store_committed(spend_counters, key=5.0, team=32.0, user=10.0, org=62.0) + + with request_redis_batch_scope(): + await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth())) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(5.0), + "team": pytest.approx(268.0), + "user": pytest.approx(40.0), + "org": pytest.approx(438.0), + } + + +@pytest.mark.asyncio +async def test_gauges_add_the_cost_to_the_auth_time_spend_before_the_charge_reaches_the_counters( + prometheus_logger, getters, spend_counters +): + _stub_gauges(prometheus_logger) + _store_committed(spend_counters, key=3.0, team=30.0, user=8.0, org=60.0) + + await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth())) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(5.0), + "team": pytest.approx(268.0), + "user": pytest.approx(40.0), + "org": pytest.approx(438.0), + } + + +@pytest.mark.asyncio +async def test_gauges_keep_the_auth_time_spend_when_the_proxy_spend_counters_cannot_be_read( + prometheus_logger, getters, monkeypatch +): + _stub_gauges(prometheus_logger) + monkeypatch.delattr("litellm.proxy.proxy_server.spend_counter_store") + + await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth())) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(5.0), + "team": pytest.approx(268.0), + "user": pytest.approx(40.0), + "org": pytest.approx(438.0), + } + + +@pytest.mark.asyncio +async def test_a_budget_hold_on_the_enforcement_counters_never_shows_in_the_gauges( + prometheus_logger, getters, spend_counters +): + _stub_gauges(prometheus_logger) + _store_committed(spend_counters, key=5.0, team=32.0, user=10.0, org=62.0) + _store_counters(spend_counters, key=9.0, team=36.0, user=14.0, org=66.0) + + await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth())) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(5.0), + "team": pytest.approx(268.0), + "user": pytest.approx(40.0), + "org": pytest.approx(438.0), + } + + +def _failed_request(metadata: dict, response_cost: float = 0.0) -> dict: + return { + "model": "openai/gpt-5.4-nano", + "litellm_params": {"metadata": metadata}, + "standard_logging_object": { + "model": "openai/gpt-5.4-nano", + "model_id": "dep-1", + "model_group": "chat", + "api_base": "https://api.openai.com", + "custom_llm_provider": "openai", + "status": "failure", + "response_cost": response_cost, + "request_tags": [], + "hidden_params": {"litellm_overhead_time_ms": None, "additional_headers": None}, + "metadata": { + "user_api_key_hash": "hashed", + "user_api_key_alias": "key-alias", + "user_api_key_team_id": "t1", + "user_api_key_team_alias": "team-alias", + "user_api_key_user_id": "u1", + "user_api_key_org_id": "o1", + "requester_metadata": None, + "user_api_key_end_user_id": None, + }, + }, + "exception": Exception("upstream failed"), + } + + +@pytest.mark.asyncio +async def test_failed_request_sets_the_org_gauge_from_the_carried_org_and_committed_spend( + prometheus_logger, getters, spend_counters +): + _stub_gauges(prometheus_logger) + _store_committed(spend_counters, key=3.0, team=30.0, user=8.0, org=65.0) + _store_counters(spend_counters, key=9.0, team=36.0, user=14.0, org=90.0) + + await prometheus_logger.async_log_failure_event( + _failed_request(_request_metadata(_token_committed_at_auth())), None, KEY_RESET_AT, KEY_RESET_AT + ) + + assert getters["get_org_object"].await_count == 0 + prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(pytest.approx(435.0)) + + +@pytest.mark.asyncio +async def test_failed_request_without_the_auth_object_still_reads_the_committed_org_spend( + prometheus_logger, getters, spend_counters +): + _stub_gauges(prometheus_logger) + _store_committed(spend_counters, key=3.0, team=30.0, user=8.0, org=65.0) + metadata = _request_metadata(_token_committed_at_auth()) + del metadata["user_api_key_auth"] + + await prometheus_logger.async_log_failure_event(_failed_request(metadata), None, KEY_RESET_AT, KEY_RESET_AT) + + prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(pytest.approx(435.0)) + + +@pytest.mark.asyncio +async def test_failed_request_without_a_carried_org_reads_no_committed_org_spend( + prometheus_logger, getters, spend_counters +): + _stub_gauges(prometheus_logger) + _store_committed(spend_counters, key=3.0, team=30.0, user=8.0, org=65.0) + token = _token_committed_at_auth().model_copy(update={"org_budget_snapshot": None}) + metadata = {key: value for key, value in _request_metadata(token).items() if key != "user_api_key_auth"} + + await prometheus_logger.async_log_failure_event(_failed_request(metadata), None, KEY_RESET_AT, KEY_RESET_AT) + + prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(pytest.approx(460.0)) + + +@pytest.mark.asyncio +async def test_failed_request_that_was_charged_counts_its_own_charge_while_it_waits_in_the_pipeline( + prometheus_logger, getters, redis_spend_counters +): + """A stream cut after 40 events still costs 2.0, and that charge is queued when the failure logs.""" + _stub_gauges(prometheus_logger) + _store_committed(redis_spend_counters, key=5.0, team=32.0, user=10.0, org=62.0) + + with _request_scope(): + await _declare_the_own_two_dollar_charge() + await prometheus_logger.async_log_failure_event( + _failed_request(_request_metadata(_token_committed_at_auth()), response_cost=2.0), + None, + KEY_RESET_AT, + KEY_RESET_AT, + ) + await flush_post_call_redis_batches() + + prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(pytest.approx(436.0)) + + +@pytest.mark.asyncio +async def test_ordinary_metadata_keys_named_like_snapshot_fields_change_no_gauge(prometheus_logger, getters): + _stub_gauges(prometheus_logger) + metadata = { + "user_api_key_team_spend": 20.0, + "user_api_key_team_max_budget": 300.0, + "spend": 1000.0, + "max_budget": 1.0, + "budget_reset_at": None, + "user_alias": "not-the-user", + } + + await _emit(prometheus_logger, metadata) + + assert {name: getter.await_count for name, getter in getters.items()} == { + "get_key_object": 1, + "get_team_object": 1, + "get_user_object": 1, + "get_org_object": 1, + } + prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(458.0) + + @pytest.mark.asyncio async def test_metadata_without_carried_state_still_fetches_each_object_once(prometheus_logger, getters): _stub_gauges(prometheus_logger) diff --git a/tests/unit/integrations/test_prometheus_user_team_metrics.py b/tests/unit/integrations/test_prometheus_user_team_metrics.py index ab0fb67b52f..acec1c882ef 100644 --- a/tests/unit/integrations/test_prometheus_user_team_metrics.py +++ b/tests/unit/integrations/test_prometheus_user_team_metrics.py @@ -1065,6 +1065,79 @@ async def test_set_customer_budget_metrics_after_api_request_uses_cached_default assert _customer_sample("litellm_customer_max_budget_metric", "cust-default") == pytest.approx(3.0) +@pytest.mark.asyncio +async def test_customer_gauge_subtracts_the_committed_spend_instead_of_the_cached_spend( + prometheus_logger, customer_metrics_enabled +): + """Auth saw 1.5 committed, and by emit the counter holds 2.5 with this request's 0.5 and another request's.""" + import sys + + from litellm.models.budget import LiteLLM_BudgetTable + from litellm.models.end_user import LiteLLM_EndUserTable + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys + from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters + + mock_proxy_server = MagicMock() + mock_proxy_server.user_api_key_cache.async_get_cache = AsyncMock( + return_value=LiteLLM_EndUserTable( + user_id="cust-live", + blocked=False, + spend=0.5, + litellm_budget_table=LiteLLM_BudgetTable(budget_id="b", max_budget=3.0), + ) + ) + mock_proxy_server.spend_counter_store = FakeSpendCounters( + redis=None, memory={"spend:committed:end_user:cust-live": 2.5} + ).store() + token = UserAPIKeyAuth(token="hashed") + token.committed_budget_spend = CommittedBudgetSpend( + counter_keys=EntitySpendCounterKeys(end_user="spend:end_user:cust-live"), end_user=1.5 + ) + + with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}): + await prometheus_logger._increment_remaining_budget_metrics( + user_api_team=None, + user_api_team_alias=None, + user_api_key=None, + user_api_key_alias=None, + litellm_params={"metadata": {"user_api_key_auth": token}}, + response_cost=0.5, + end_user_id="cust-live", + ) + + assert _customer_sample("litellm_remaining_customer_budget_metric", "cust-live") == pytest.approx(0.5) + + +@pytest.mark.asyncio +async def test_customer_gauge_counts_a_charge_the_cached_customer_row_already_holds_once( + prometheus_logger, customer_metrics_enabled +): + import sys + + from litellm.models.budget import LiteLLM_BudgetTable + from litellm.models.end_user import LiteLLM_EndUserTable + + mock_proxy_server = MagicMock() + mock_proxy_server.user_api_key_cache.async_get_cache = AsyncMock( + return_value=LiteLLM_EndUserTable( + user_id="cust-charged", + blocked=False, + spend=0.5, + litellm_budget_table=LiteLLM_BudgetTable(budget_id="b", max_budget=3.0), + ) + ) + + with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}): + await prometheus_logger._set_customer_budget_metrics_after_api_request( + end_user_id="cust-charged", + response_cost=0.5, + committed_spend=0.0, + ) + + assert _customer_sample("litellm_remaining_customer_budget_metric", "cust-charged") == pytest.approx(2.5) + + @pytest.mark.asyncio async def test_set_customer_budget_metrics_after_api_request_without_budget_only_emits_remaining( prometheus_logger, customer_metrics_enabled diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_utils.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_utils.py index 1d2a5e4cc9b..f007d3740aa 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_utils.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_utils.py @@ -1,4 +1,5 @@ import asyncio +import math import uuid from collections.abc import Mapping from types import MappingProxyType @@ -483,3 +484,42 @@ async def test_router_prices_balanced_columns_by_the_same_rules_as_flex( assert balanced_cost == pytest.approx(flex_cost) if not overridden_columns: assert balanced_cost == pytest.approx(_cost_at(PARITY_ROW, "_balanced")) + + +def _vertex_cost_with_cache_writes(model: str, prompt_tokens: int, cache_creation_tokens: int) -> float: + response: Final = ModelResponse( + model=model, + usage=Usage( + prompt_tokens=prompt_tokens, + completion_tokens=100, + total_tokens=prompt_tokens + 100, + cache_creation_input_tokens=cache_creation_tokens, + ), + ) + return litellm.completion_cost(completion_response=response, model=model, custom_llm_provider="vertex_ai") + + +@pytest.mark.parametrize( + "model", + ( + "gemini-2.5-pro", + "gemini-3.1-pro-preview", + "gemini-3.1-pro-preview-customtools", + "vertex_ai/gemini-3-pro-preview", + "vertex_ai/gemini-3.1-pro-preview", + "vertex_ai/gemini-3.1-pro-preview-customtools", + ), +) +def test_vertex_gemini_pro_bills_cache_writes_above_200k_at_the_above_200k_input_rate( + local_model_cost_map: None, model: str +) -> None: + rates: Final = litellm.get_model_info(model=model, custom_llm_provider="vertex_ai") + input_rate_above_200k: Final = rates.get("input_cost_per_token_above_200k_tokens") + output_rate_above_200k: Final = rates.get("output_cost_per_token_above_200k_tokens") + cache_creation_tokens: Final = 150_000 + below: Final = _vertex_cost_with_cache_writes(model, 190_000, cache_creation_tokens) + above: Final = _vertex_cost_with_cache_writes(model, 210_000, cache_creation_tokens) + + assert input_rate_above_200k is not None and output_rate_above_200k is not None + assert above > below + assert math.isclose(above, 210_000 * input_rate_above_200k + 100 * output_rate_above_200k) diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index 9d3c19b7460..777bef7afcc 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -1,10 +1,15 @@ +import contextlib import copy +import importlib import json +import warnings from collections.abc import Mapping from typing import Final, cast +from unittest import mock import pytest +from litellm.litellm_core_utils.prompt_templates import common_utils from litellm.litellm_core_utils.prompt_templates.common_utils import ( TOOL_RESULT_IMAGE_BOUNDARY, TOOL_RESULT_IMAGE_PLACEHOLDER, @@ -21,6 +26,10 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( strip_litellm_response_only_fields, system_messages_first, ) +from litellm.litellm_core_utils.prompt_templates.json_schema_utils import ( + flatten_combinators_and_drop_non_python_regex_patterns, +) +from litellm.litellm_core_utils.prompt_templates.tool_call_json_repair import parse_tool_call_arguments from litellm.types.llms.openai import ChatCompletionAudioObject, ChatCompletionFileObject, ChatCompletionImageObject from litellm.types.utils import ChatCompletionMessageToolCall, Choices, Function, Message, ModelResponse @@ -987,6 +996,180 @@ def test_content_list_flattening_returns_new_messages_and_leaves_the_callers_unt assert messages == messages_before +_MOVED_PRIVATE_NAMES_BY_MODULE: Final = ( + ( + "file_id_utils", + ("_decode_vector_store_ids_in_tools",), + ), + ( + "json_schema_utils", + ( + "_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", + ), + ), + ( + "prompt_cache_breakpoint_utils", + ( + "_BREAKPOINT_BLOCKS", + "_BREAKPOINT_BLOCK_LIST_KEYS", + "_BREAKPOINT_FIELDS", + "_MarkedT", + "_PROMPT_CACHE_BREAKPOINT", + "_breakpoints_stripped_from", + "_without_prompt_cache_breakpoint", + ), + ), + ( + "reasoning_utils", + ( + "_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", + ), + ), + ( + "tool_call_json_repair", + ( + "_DECODED_JSON", + "_attempt_json_repair", + ), + ), + ( + "tpm_media_utils", + ( + "_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", + ), + ), +) + + +@pytest.mark.parametrize(("module_name", "names"), _MOVED_PRIVATE_NAMES_BY_MODULE) +def test_moved_private_helpers_still_import_from_common_utils(module_name: str, names: tuple[str, ...]) -> None: + new_home: Final = importlib.import_module(f"litellm.litellm_core_utils.prompt_templates.{module_name}") + unreachable: Final = [name for name in names if getattr(common_utils, name, None) is not getattr(new_home, name)] + assert unreachable == [] + + +def test_unknown_private_name_still_raises_attribute_error_on_common_utils() -> None: + missing_name: Final = "_not_a_moved_helper" + with pytest.raises(AttributeError, match=missing_name): + getattr(common_utils, missing_name) + + +def test_old_path_access_warns_from_the_callers_line_with_the_new_home() -> None: + moved_name: Final = "_any_of" + with pytest.warns(DeprecationWarning, match=r"common_utils\._any_of moved to .*\.json_schema_utils") as caught: + getattr(common_utils, moved_name) + assert [warning.filename for warning in caught] == [__file__] + + +def test_patching_a_moved_helper_at_its_old_path_reaches_the_callers_in_its_new_home() -> None: + patched_result: Final = {"patched": True} + with mock.patch( + "litellm.litellm_core_utils.prompt_templates.common_utils._attempt_json_repair", return_value=patched_result + ): + assert parse_tool_call_arguments('{"a": 1') == patched_result + assert parse_tool_call_arguments('{"a": 1') == {"a": 1} + + +def test_patching_a_re_exported_helper_at_its_old_path_reaches_the_callers_in_its_new_home() -> None: + schema: Final = {"anyOf": [{"type": "object", "properties": {"a": {"type": "string"}}}]} + unpatched: Final = flatten_combinators_and_drop_non_python_regex_patterns(schema) + patched_result: Final = {"patched": True} + with mock.patch( + "litellm.litellm_core_utils.prompt_templates.common_utils.flatten_top_level_schema_combinators", + return_value=patched_result, + ): + assert flatten_combinators_and_drop_non_python_regex_patterns(schema) == patched_result + assert flatten_combinators_and_drop_non_python_regex_patterns(schema) == unpatched + + +def test_old_path_patch_restores_the_moved_helper_when_warnings_become_errors_inside_it() -> None: + with warnings.catch_warnings(), contextlib.suppress(DeprecationWarning): + with mock.patch( + "litellm.litellm_core_utils.prompt_templates.common_utils._attempt_json_repair", return_value=None + ): + warnings.simplefilter("error", DeprecationWarning) + assert parse_tool_call_arguments('{"a": 1') == {"a": 1} + + +def test_old_path_patch_with_create_leaves_the_moved_helper_in_place() -> None: + with mock.patch( + "litellm.litellm_core_utils.prompt_templates.common_utils._attempt_json_repair", return_value=None, create=True + ): + assert parse_tool_call_arguments("{}") == {} + assert parse_tool_call_arguments('{"a": 1') == {"a": 1} + + +def test_overriding_a_moved_helper_through_the_old_path_warns_and_deleting_it_restores_the_released_one() -> None: + moved_name: Final = "_attempt_json_repair" + with pytest.warns(DeprecationWarning, match=r"common_utils\._attempt_json_repair moved to .*\.tool_call_json_repair"): + setattr(common_utils, moved_name, mock.Mock(return_value={"overridden": True})) + assert parse_tool_call_arguments('{"a": 1') == {"overridden": True} + delattr(common_utils, moved_name) + assert parse_tool_call_arguments('{"a": 1') == {"a": 1} + + +def test_patching_a_name_common_utils_still_defines_stays_on_common_utils() -> None: + with mock.patch.object(common_utils, "invalid_type_message", return_value="patched"): + assert common_utils.invalid_type_message("field", 1) == "patched" + assert common_utils.invalid_type_message("field", 1) != "patched" + + def test_get_content_from_model_response_includes_server_tool_calls_before_client_calls() -> None: server_call: Final = ChatCompletionMessageToolCall( id="srvtoolu_01", type="function", function=Function(name="web_search", arguments='{"query": "persimmon"}') diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_file_id_utils.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_file_id_utils.py similarity index 100% rename from tests/unit/litellm_core_utils/prompt_templates/test_file_id_utils.py rename to tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_file_id_utils.py diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_json_schema_utils.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_json_schema_utils.py similarity index 100% rename from tests/unit/litellm_core_utils/prompt_templates/test_json_schema_utils.py rename to tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_json_schema_utils.py diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_prompt_cache_breakpoint_utils.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_prompt_cache_breakpoint_utils.py similarity index 100% rename from tests/unit/litellm_core_utils/prompt_templates/test_prompt_cache_breakpoint_utils.py rename to tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_prompt_cache_breakpoint_utils.py diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_reasoning_utils.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_reasoning_utils.py similarity index 100% rename from tests/unit/litellm_core_utils/prompt_templates/test_reasoning_utils.py rename to tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_reasoning_utils.py diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_tool_call_json_repair.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_tool_call_json_repair.py similarity index 66% rename from tests/unit/litellm_core_utils/prompt_templates/test_tool_call_json_repair.py rename to tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_tool_call_json_repair.py index 75f4e686137..4bd86d80c36 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_tool_call_json_repair.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_tool_call_json_repair.py @@ -1,7 +1,11 @@ import json +import logging +from typing import Final import pytest +import litellm +from litellm.constants import REDACTED_BY_LITELLM from litellm.litellm_core_utils.prompt_templates.tool_call_json_repair import ( parse_tool_call_arguments, salvage_concatenated_tool_arguments, @@ -115,3 +119,53 @@ def _distinct_json_objects(count: int) -> str: ) def test_salvage_concatenated_tool_arguments(raw: str, expected: tuple[dict[str, object], ...]) -> None: assert salvage_concatenated_tool_arguments(raw) == expected + + +def test_parse_tool_call_arguments_repair_warning_previews_the_first_200_characters( + caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "turn_off_message_logging", False) + arguments: Final = '{"patient": "Jane Doe", "note": "' + "x" * 300 + '"' + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + result: Final = parse_tool_call_arguments(arguments, tool_name="lookup", context="test") + + assert result == {"patient": "Jane Doe", "note": "x" * 300} + assert ( + f"Repaired truncated tool call arguments for tool 'lookup' (test). Original ({len(arguments)} chars): " + f"{arguments[:200]}..." in caplog.text + ) + assert arguments not in caplog.text + + +@pytest.mark.parametrize( + ("arguments", "expected"), + ( + ('{"patient": "Jane Doe"', {"patient": "Jane Doe"}), + ('{"note": "' + "x" * 189 + '"', {"note": "x" * 189}), + ), +) +def test_parse_tool_call_arguments_repair_warning_logs_arguments_up_to_200_characters_whole( + caplog: pytest.LogCaptureFixture, + monkeypatch: pytest.MonkeyPatch, + arguments: str, + expected: dict[str, str], +) -> None: + monkeypatch.setattr(litellm, "turn_off_message_logging", False) + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + result: Final = parse_tool_call_arguments(arguments, tool_name="lookup", context="test") + + assert result == expected + assert f"Original ({len(arguments)} chars): {arguments}\n" in caplog.text + + +def test_parse_tool_call_arguments_repair_warning_redacts_arguments_when_message_logging_is_off( + caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "turn_off_message_logging", True) + arguments: Final = '{"patient": "Jane Doe"' + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + result: Final = parse_tool_call_arguments(arguments, tool_name="lookup", context="test") + + assert result == {"patient": "Jane Doe"} + assert f"Original ({len(arguments)} chars): {REDACTED_BY_LITELLM}" in caplog.text + assert "Jane Doe" not in caplog.text diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_tpm_media_utils.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_tpm_media_utils.py similarity index 100% rename from tests/unit/litellm_core_utils/prompt_templates/test_tpm_media_utils.py rename to tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_tpm_media_utils.py diff --git a/tests/unit/litellm_core_utils/test_audio_utils.py b/tests/unit/litellm_core_utils/test_audio_utils.py index 122445e8e30..520eeab12e9 100644 --- a/tests/unit/litellm_core_utils/test_audio_utils.py +++ b/tests/unit/litellm_core_utils/test_audio_utils.py @@ -6,6 +6,7 @@ import base64 import io import os import tempfile +from pathlib import Path from typing import Final from unittest.mock import mock_open, patch @@ -329,16 +330,12 @@ class TestCalculateRequestDuration: duration = calculate_request_duration(file_obj) # Verify it succeeded (returns a duration, not None) - assert ( - duration is not None - ), "Duration should be calculated even when BytesIO is at end" + assert duration is not None, "Duration should be calculated even when BytesIO is at end" assert isinstance(duration, float), "Duration should be a float" assert duration > 0, "Duration should be positive" # Verify the file position was restored - assert file_obj.tell() == len( - wav_header - ), "File position should be restored to original position" + assert file_obj.tell() == len(wav_header), "File position should be restored to original position" class TestGetAudioFileContentHash: @@ -364,9 +361,7 @@ class TestGetAudioFileContentHash: hash1 = get_audio_file_content_hash((filename1, content)) hash2 = get_audio_file_content_hash((filename2, content)) - assert ( - hash1 == hash2 - ), "Same content should produce same hash regardless of filename" + assert hash1 == hash2, "Same content should produce same hash regardless of filename" def test_bytes_input(self): """Test that raw bytes input works""" @@ -473,3 +468,17 @@ class TestSpeechMediaTypeFromAudioBytes: from litellm.litellm_core_utils.audio_utils.utils import speech_media_type_from_audio_bytes assert speech_media_type_from_audio_bytes(audio) == expected + + +@pytest.mark.parametrize("include_content_type", [False, True]) +def test_duration_reads_pathlike_content_in_multipart_tuple(tmp_path: Path, include_content_type: bool) -> None: + audio_path: Final = tmp_path / "sample.wav" + audio_path.write_bytes(base64.b64decode(_encoded_audio(2, 8000, 1, "WAV"))) + upload: Final = ("sample.wav", audio_path, "audio/wav") if include_content_type else ("sample.wav", audio_path) + assert calculate_request_duration(upload) == 2.0 + + +def test_duration_does_not_open_bare_string_tuple_content(tmp_path: Path) -> None: + audio_path: Final = tmp_path / "sample.wav" + audio_path.write_bytes(base64.b64decode(_encoded_audio(2, 8000, 1, "WAV"))) + assert calculate_request_duration(("sample.wav", str(audio_path))) is None diff --git a/tests/unit/litellm_core_utils/test_fallback_generalizations.py b/tests/unit/litellm_core_utils/test_fallback_generalizations.py index d8e56c0ae51..b22edaa171a 100644 --- a/tests/unit/litellm_core_utils/test_fallback_generalizations.py +++ b/tests/unit/litellm_core_utils/test_fallback_generalizations.py @@ -810,6 +810,35 @@ def test_shipped_wandb_rule_is_anchored_to_the_wandb_namespace(shipped_cost_map) assert matched is None or not matched.get("supports_reasoning"), foreign +@pytest.mark.parametrize("model", ["minimax/MiniMax-M3.2", "minimax/MiniMax-Text-02"]) +def test_shipped_rules_flag_unmapped_minimax_ids_as_reasoning(shipped_cost_map, model): + """MiniMax ships reasoning models faster than the registry names them. The rule + carries no mode and no pricing, so cost stays on the standard unpriced behavior.""" + assert model not in litellm.model_cost + + info = litellm.get_model_info(model, custom_llm_provider="minimax") + assert info["litellm_provider"] == "minimax" + assert info["supports_reasoning"] is True + assert info.get("mode") is None + assert not info.get("input_cost_per_token") + assert not info.get("output_cost_per_token") + + +def test_shipped_minimax_rule_is_anchored_to_the_minimax_namespace(shipped_cost_map): + """``^minimax/`` is anchored, so it cannot leak onto another provider's ids.""" + assert match_capability_generalizations("minimax/some-new-model") == {"supports_reasoning": True} + for foreign in ("openrouter/minimax/minimax-m3", "notminimax/MiniMax-M3", "together_ai/minimax/MiniMax-M3"): + matched = match_capability_generalizations(foreign) + assert matched is None or not matched.get("supports_reasoning"), foreign + + +@pytest.mark.parametrize("model", ["speech-02-hd", "MiniMax-H3"]) +def test_shipped_minimax_rule_leaves_mapped_entries_alone(shipped_cost_map, model): + assert f"minimax/{model}" in litellm.model_cost + assert match_fill_missing_generalizations(f"minimax/{model}", "minimax") is None + assert litellm.supports_reasoning(model=model, custom_llm_provider="minimax") is False + + def test_shipped_wandb_rule_keeps_reasoning_effort_on_an_unmapped_model(shipped_cost_map): """End to end through the provider config: the gate WandbConfig applies reads the rule, so reasoning_effort is advertised and survives get_optional_params rather than diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 58fd03e0cbc..2f5edbd6350 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -14690,3 +14690,66 @@ def test_standard_logging_payload_reports_the_per_call_prices_the_request_was_bi model_map_value: Final = payload["model_map_information"]["model_map_value"] assert model_map_value is not None assert (model_map_value["input_cost_per_token"], model_map_value["output_cost_per_token"]) == (3e-06, 7e-06) + + +@pytest.mark.asyncio +async def test_a_response_discarded_after_its_success_log_was_parked_is_billed_once_on_the_serving_attempt( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + from litellm.responses.streaming_iterator import drain_logging_tasks + from litellm.utils import _schedule_async_success_logging + + class Recorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.events: tuple[tuple[dict[str, object], object], ...] = () + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self.events = (*self.events, (dict(kwargs), response_obj)) + + recorder: Final = Recorder() + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "success_callback", []) + logging_obj: Final = LitellmLogging( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="acompletion", + start_time=time.time(), + litellm_call_id="refused-then-served", + function_id="refused-then-served", + dynamic_async_success_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model="gpt-4o-mini", user="", optional_params={}, litellm_params={"metadata": {}}, custom_llm_provider="openai" + ) + logging_obj.defer_async_logging = True + refused: Final = ModelResponse( + id="refused", model="gpt-4o-mini", usage=Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + ) + served: Final = ModelResponse( + id="served", model="gpt-4o-mini", usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) + ) + now: Final = datetime.datetime.now() + + _schedule_async_success_logging(logging_obj, refused, now, now, is_completion_with_fallbacks=False) + logging_obj.carry_failed_attempt_usage() + _schedule_async_success_logging(logging_obj, served, now, now, is_completion_with_fallbacks=False) + prepared: Final = logging_obj.prepare_response_cost(served, now) + enqueue: Final = logging_obj.enqueue_deferred_logging + assert enqueue is not None + enqueue() + await drain_logging_tasks() + await GLOBAL_LOGGING_WORKER.flush() + + both_attempts: Final = litellm.completion_cost(refused, model="gpt-4o-mini") + litellm.completion_cost( + served, model="gpt-4o-mini" + ) + ((kwargs, response_obj),) = recorder.events + assert getattr(response_obj, "id", None) == "served" + assert kwargs["response_cost"] == pytest.approx(both_attempts) + assert prepared is not None + assert prepared.amount == pytest.approx(both_attempts) diff --git a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py index 8059cab6500..024d4b2b39b 100644 --- a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py +++ b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_handler.py @@ -2178,6 +2178,89 @@ async def test_anthropic_messages_forwards_safeguards_and_dangerous_tool_use_bet assert response["safeguard_results"] == safeguard_results +_OPUS_4_8_ON_BEDROCK_INVOKE: Final = "bedrock/invoke/us.anthropic.claude-opus-4-8" + + +def _bedrock_invoke_recording(bodies: list[dict[str, object]]) -> AsyncHTTPHandler: + def record_and_answer(request: httpx.Request) -> httpx.Response: + bodies.append(json.loads(request.content)) + return httpx.Response( + 200, + json={ + "id": "msg_drop_top_k", + "type": "message", + "role": "assistant", + "model": "claude-opus-4-8", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 3, "output_tokens": 1}, + }, + request=request, + ) + + upstream: Final = AsyncHTTPHandler() + upstream.client = httpx.AsyncClient(transport=httpx.MockTransport(record_and_answer)) + return upstream + + +async def _opus_4_8_messages_with_top_k(bodies: list[dict[str, object]], **deployment: object) -> object: + return await messages.acreate( + max_tokens=16, + messages=[{"role": "user", "content": "hi"}], + model=_OPUS_4_8_ON_BEDROCK_INVOKE, + top_k=5, + aws_access_key_id="test-access-key", + aws_secret_access_key="test-secret-key", + aws_region_name="us-east-1", + client=_bedrock_invoke_recording(bodies), + **deployment, + ) + + +@pytest.mark.asyncio +async def test_bedrock_invoke_messages_drops_top_k_the_deployment_lists_on_a_model_without_sampling_params( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + bodies: Final[list[dict[str, object]]] = [] + + response: Final = await _opus_4_8_messages_with_top_k(bodies, additional_drop_params=["top_k"]) + + assert response["content"] == [{"type": "text", "text": "ok"}] + (body,) = bodies + assert "top_k" not in body + assert body["max_tokens"] == 16 + + +@pytest.mark.asyncio +async def test_bedrock_invoke_messages_still_rejects_top_k_the_deployment_does_not_list( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + bodies: Final[list[dict[str, object]]] = [] + + with pytest.raises(litellm.UnsupportedParamsError, match="does not support top_k=5") as raised: + await _opus_4_8_messages_with_top_k(bodies, additional_drop_params=["top_p"]) + + assert raised.value.status_code == 400 + assert bodies == [] + + +@pytest.mark.asyncio +async def test_bedrock_invoke_messages_answers_the_top_k_400_for_a_drop_list_that_is_not_a_list( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + bodies: Final[list[dict[str, object]]] = [] + + with pytest.raises(litellm.UnsupportedParamsError, match="does not support top_k=5") as raised: + await _opus_4_8_messages_with_top_k(bodies, additional_drop_params="top_k") + + assert raised.value.status_code == 400 + assert bodies == [] + + @pytest.mark.asyncio @pytest.mark.parametrize("client_headers", _CLIENT_BETA_HEADERS) async def test_anthropic_messages_forwards_safeguards_and_dangerous_tool_use_beta_to_vertex( diff --git a/tests/unit/llms/anthropic/messages/native/test_request_optional_param_utils.py b/tests/unit/llms/anthropic/messages/native/test_request_optional_param_utils.py index d525ccfec37..8104c0cef56 100644 --- a/tests/unit/llms/anthropic/messages/native/test_request_optional_param_utils.py +++ b/tests/unit/llms/anthropic/messages/native/test_request_optional_param_utils.py @@ -6,6 +6,8 @@ Regression tests for the /v1/messages request-parse fast paths: while resolving the (static) type hints only once per process. """ +from typing import Final + import pytest import litellm @@ -133,3 +135,37 @@ def test_sampling_param_raises_clean_400_without_drop_params(monkeypatch): model="claude-opus-4-7", drop_params=False, ) + + +@pytest.mark.parametrize( + ("model", "param", "value"), + [ + ("claude-opus-4-8", "temperature", 0.3), + ("claude-opus-4-8", "top_p", 0.9), + ("claude-opus-4-8", "top_k", 5), + ("claude-sonnet-4-6", "speed", "fast"), + ], +) +def test_deployment_listed_param_is_dropped_before_the_unsupported_param_gate(monkeypatch, model, param, value): + monkeypatch.setattr(litellm, "drop_params", False) + result: Final = AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param( + params={param: value, "stream": True}, + model=model, + drop_params=False, + additional_drop_params=(param,), + ) + + assert result == {"stream": True} + + +def test_param_missing_from_the_drop_list_still_raises_the_unsupported_param_400(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.utils.UnsupportedParamsError, match="does not support top_p") as raised: + AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param( + params={"top_k": 5, "top_p": 0.9}, + model="claude-opus-4-8", + drop_params=False, + additional_drop_params=("top_k", "top_p.nested"), + ) + + assert raised.value.status_code == 400 diff --git a/tests/unit/llms/anthropic/test_anthropic_prompt_cache_prediction.py b/tests/unit/llms/anthropic/test_anthropic_prompt_cache_prediction.py index d95f831df3e..89ea412af0c 100644 --- a/tests/unit/llms/anthropic/test_anthropic_prompt_cache_prediction.py +++ b/tests/unit/llms/anthropic/test_anthropic_prompt_cache_prediction.py @@ -618,6 +618,26 @@ def test_native_baseline_body_drops_a_non_numeric_budget_without_logging_a_drop( assert [record.getMessage() for record in caplog.records if "Dropping" in record.getMessage()] == [] +@pytest.mark.usefixtures("local_model_cost_map") +def test_native_baseline_body_leaves_out_a_listed_sampling_param_the_model_does_not_support( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + request: Final = { + "model": "claude-opus-4-8", + "max_tokens": 64, + "messages": [{"role": "user", "content": "hi"}], + "top_k": 5, + "additional_drop_params": ["top_k"], + } + + body: Final = prepare_native_baseline_body(request, "anthropic/claude-opus-4-8") + + assert body is not None + assert "top_k" not in body + assert body["max_tokens"] == 64 + + @pytest.mark.usefixtures("local_model_cost_map") def test_native_baseline_body_drops_a_non_string_reasoning_effort_without_logging_a_drop( caplog: pytest.LogCaptureFixture, diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py index 591205b4602..19c9a042afb 100644 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_handler.py @@ -6,6 +6,7 @@ import httpx import pytest from botocore.credentials import RefreshableCredentials +from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler from litellm.llms.bedrock.count_tokens.transformation import BedrockCountTokensConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @@ -85,6 +86,42 @@ async def test_handle_count_tokens_request_looks_up_its_region_off_the_event_loo assert result == {"input_tokens": 7} +_COUNTED_MODEL: Final = "us.anthropic.claude-opus-4-8" + + +@pytest.mark.parametrize( + ("request_data", "problem"), + [ + ({"model": _COUNTED_MODEL}, "Request must contain content to count tokens"), + ({"model": _COUNTED_MODEL, "messages": []}, "messages parameter is required for Converse input"), + ({"model": _COUNTED_MODEL, "messages": [{"role": "user"}]}, "Message 0 must have a 'content' field"), + ({"model": _COUNTED_MODEL, "messages": [{"content": "hi"}]}, "Message 0 must have a 'role' field"), + ({"model": _COUNTED_MODEL, "messages": ["hi"]}, "Message 0 must be a dictionary"), + ( + {"model": _COUNTED_MODEL, "messages": [{"role": "user", "content": "hi"}, {"role": "assistant"}]}, + "Message 1 must have a 'content' field", + ), + ({"messages": [{"role": "user", "content": "hi"}]}, "model parameter is required"), + ], +) +@pytest.mark.asyncio +async def test_handle_count_tokens_request_answers_400_for_a_malformed_body_without_calling_bedrock( + request_data: dict[str, object], problem: str +) -> None: + client: Final = AsyncMock(spec=AsyncHTTPHandler) + + with pytest.raises(BedrockError) as exc_info: + await BedrockCountTokensHandler().handle_count_tokens_request( + request_data=request_data, + litellm_params={"aws_region_name": "us-east-1"}, + resolved_model=_COUNTED_MODEL, + client=client, + ) + + assert (exc_info.value.status_code, exc_info.value.message) == (400, problem) + client.post.assert_not_awaited() + + class TestBedrockCountTokensEndpoint: def _make_handler(self) -> BedrockCountTokensConfig: return BedrockCountTokensConfig() diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py index d5b57e42146..018b0fde8a0 100644 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_count_tokens_transformation.py @@ -311,3 +311,37 @@ def test_count_tokens_endpoint_encodes_model_id(monkeypatch): endpoint == "https://bedrock-runtime.us-east-1.amazonaws.com/model/..%2F..%2Fmodel%2Fother%3Fx%3D1%23frag/count-tokens" ) + + +@pytest.mark.parametrize( + "request_data", + [ + {"model": "us.anthropic.claude-opus-4-8", "messages": [{"role": "user", "content": "hi"}]}, + {"model": "us.anthropic.claude-opus-4-8", "messages": "hi"}, + { + "model": "us.anthropic.claude-opus-4-8", + "messages": [{"role": "user", "content": [{"type": "text", "text": "hi"}]}, {"content": "hi"}], + }, + ], +) +def test_count_tokens_request_problem_leaves_bodies_bedrock_can_judge_to_bedrock(request_data): + assert BedrockCountTokensConfig().count_tokens_request_problem(request_data) is None + + +def test_count_tokens_request_problem_names_tool_calls_instead_of_asking_for_content(): + problem = BedrockCountTokensConfig().count_tokens_request_problem( + { + "model": "us.amazon.nova-pro-v1:0", + "messages": [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "tool_calls": [{"id": "1", "type": "function", "function": {"name": "t", "arguments": "{}"}}], + }, + ], + } + ) + + assert problem is not None + assert "tool_calls" in problem + assert "must have a 'content' field" not in problem diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py index 30063936b2b..de0c14558b3 100644 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py @@ -74,6 +74,59 @@ async def test_claude_model_bedrock_runtime_cannot_count_is_counted_on_mantle() assert mantle_body["system"] == "You are a terse assistant." +@pytest.mark.asyncio +async def test_claude_model_message_without_content_is_rejected_without_calling_runtime_or_mantle() -> None: + bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 1353}), mantle=(200, {"input_tokens": 1336})) + + result: Final = await _counter(bedrock).count_tokens( + model_to_use="global.anthropic.claude-opus-4-8", + messages=[{"role": "user"}], + contents=None, + deployment=DEPLOYMENT, + request_model="claude-opus-4-8", + ) + + assert result is not None + assert (result.error, result.status_code, result.error_message, result.rejected_before_provider_call) == ( + True, + 400, + "Message 0 must have a 'content' field", + True, + ) + assert bedrock.posted_hosts() == () + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("message", "problem"), + [ + ({"role": "user"}, "Message 0 must have a 'content' field"), + ({"content": "hi"}, "Message 0 must have a 'role' field"), + ], +) +async def test_mantle_counted_model_rejects_a_malformed_message_without_calling_mantle( + message: dict[str, str], problem: str +) -> None: + bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 1353}), mantle=(200, {"input_tokens": 1336})) + + result: Final = await _counter(bedrock).count_tokens( + model_to_use="bedrock_mantle/anthropic.claude-opus-5-5", + messages=[message], + contents=None, + deployment=DEPLOYMENT, + request_model="claude-opus-5-5", + ) + + assert result is not None + assert (result.error, result.status_code, result.error_message, result.rejected_before_provider_call) == ( + True, + 400, + problem, + True, + ) + assert bedrock.posted_hosts() == () + + @pytest.mark.asyncio async def test_bedrock_runtime_count_is_kept_when_it_answers() -> None: bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 1353}), mantle=(200, {"input_tokens": 1336})) diff --git a/tests/unit/llms/bedrock/test_bedrock_common_utils.py b/tests/unit/llms/bedrock/test_bedrock_common_utils.py index 31999ee0379..0799ec15225 100644 --- a/tests/unit/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/unit/llms/bedrock/test_bedrock_common_utils.py @@ -1663,6 +1663,21 @@ class TestGetBedrockBaseModel: assert get_bedrock_base_model(input_model) == expected +def test_region_only_deployment_takes_no_traffic_from_a_valid_model_group() -> None: + valid_model: Final = "bedrock/us.anthropic.claude-opus-4-6-v1" + router: Final = litellm.Router( + model_list=[ + {"model_name": "group", "litellm_params": {"model": "bedrock/us", "aws_region_name": "us-east-1"}}, + {"model_name": "group", "litellm_params": {"model": valid_model, "aws_region_name": "us-east-1"}}, + ], + ignore_invalid_deployments=True, + ) + + assert [deployment["litellm_params"]["model"] for deployment in router.get_model_list(model_name="group")] == [ + valid_model + ] + + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") class TestBedrockModelInfoWrappers: """Tests that BedrockModelInfo methods correctly wrap standalone functions.""" diff --git a/tests/unit/llms/minimax/chat/test_transformation.py b/tests/unit/llms/minimax/chat/test_transformation.py index 2645b2832aa..d289bdba9bf 100644 --- a/tests/unit/llms/minimax/chat/test_transformation.py +++ b/tests/unit/llms/minimax/chat/test_transformation.py @@ -2,10 +2,18 @@ Test MiniMax OpenAI-compatible API support """ +import json +from collections.abc import Mapping +from pathlib import Path +from typing import Final, get_args from unittest.mock import MagicMock, patch +import pytest +import respx + import litellm from litellm.llms.minimax.chat.transformation import MinimaxChatConfig +from litellm.types.llms.openai import REASONING_EFFORT def test_minimax_chat_config(): @@ -94,9 +102,7 @@ def test_minimax_provider_config_manager(): from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager - config = ProviderConfigManager.get_provider_chat_config( - model="MiniMax-M2.1", provider=LlmProviders.MINIMAX - ) + config = ProviderConfigManager.get_provider_chat_config(model="MiniMax-M2.1", provider=LlmProviders.MINIMAX) assert config is not None assert isinstance(config, MinimaxChatConfig) @@ -121,3 +127,90 @@ if __name__ == "__main__": print("✓ Provider config manager test passed") print("\n✅ All basic tests passed!") + + +@pytest.fixture +def minimax_request_mock( + local_model_cost_map: None, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> respx.Route: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + return respx_mock.post("https://api.minimax.io/v1/chat/completions").respond( + json={ + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 1677652288, + "model": "MiniMax-M3.1-Flash", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "Done"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 9, "completion_tokens": 12, "total_tokens": 21}, + } + ) + + +def _sent_body(model: str, request_mock: respx.Route, effort: str, drop_params: bool = True) -> Mapping[str, object]: + litellm.completion( + model=f"minimax/{model}", + messages=[{"role": "user", "content": "What is the weather in Oslo?"}], + api_key="fake-minimax-key", + reasoning_effort=effort, + drop_params=drop_params, + ) + body: Final[Mapping[str, object]] = json.loads(request_mock.calls.last.request.content) + return body + + +def _flash_effort_levels() -> tuple[str, ...]: + levels: Final = litellm.get_model_info("MiniMax-M3.1-Flash", custom_llm_provider="minimax")[ + "reasoning_effort_levels" + ] + assert levels + return tuple(levels) + + +def test_minimax_flash_forwards_every_declared_reasoning_effort(minimax_request_mock: respx.Route) -> None: + """Issue #49574: reasoning_effort sent to MiniMax-M3.1-Flash was stripped before the upstream call.""" + levels: Final = _flash_effort_levels() + sent: Final = tuple( + _sent_body("MiniMax-M3.1-Flash", minimax_request_mock, effort)["reasoning_effort"] for effort in levels + ) + assert sent == levels + assert minimax_request_mock.call_count == len(levels) + + +def test_minimax_flash_cost_map_entry_is_a_priced_reasoning_model(local_model_cost_map: None) -> None: + repo_root: Final = Path(litellm.__file__).parent.parent + entries: Final = tuple( + json.loads(path.read_text())["minimax/MiniMax-M3.1-Flash"] + for path in ( + repo_root / "model_prices_and_context_window.json", + repo_root / "litellm/model_prices_and_context_window_backup.json", + ) + ) + assert entries[0] == entries[1] + + info: Final = litellm.get_model_info("MiniMax-M3.1-Flash", custom_llm_provider="minimax") + assert info["supports_reasoning"] is True + assert set(_flash_effort_levels()) <= set(get_args(REASONING_EFFORT)) + assert 0 < info["cache_read_input_token_cost"] < info["input_cost_per_token"] < info["output_cost_per_token"] + + +@pytest.mark.parametrize("model", ["MiniMax-M3.2", "MiniMax-Text-02"]) +def test_minimax_forwards_reasoning_effort_for_an_unmapped_model(minimax_request_mock: respx.Route, model: str) -> None: + """A MiniMax model newer than the cost map keeps reasoning_effort and thinking through the fallback rule.""" + assert f"minimax/{model}" not in litellm.model_cost + assert "thinking" in MinimaxChatConfig().get_supported_openai_params(model=model) + assert _sent_body(model, minimax_request_mock, "low", drop_params=False)["reasoning_effort"] == "low" + + +@pytest.mark.respx(assert_all_called=False) +def test_minimax_drops_reasoning_effort_for_a_known_non_reasoning_model( + minimax_request_mock: respx.Route, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setitem(litellm.model_cost["minimax/MiniMax-M3"], "supports_reasoning", False) + litellm.get_model_info.cache_clear() + + assert "reasoning_effort" not in MinimaxChatConfig().get_supported_openai_params(model="MiniMax-M3") + assert "reasoning_effort" not in _sent_body("MiniMax-M3", minimax_request_mock, "low") + + with pytest.raises(litellm.UnsupportedParamsError, match="reasoning_effort"): + _sent_body("MiniMax-M3", minimax_request_mock, "low", drop_params=False) + assert minimax_request_mock.call_count == 1 diff --git a/tests/unit/llms/ollama/test_ollama_chat_transformation.py b/tests/unit/llms/ollama/test_ollama_chat_transformation.py index 98ca5aee407..36df47deaab 100644 --- a/tests/unit/llms/ollama/test_ollama_chat_transformation.py +++ b/tests/unit/llms/ollama/test_ollama_chat_transformation.py @@ -1,8 +1,10 @@ import asyncio import inspect +import itertools import logging import os import sys +import traceback from typing import Final, cast import pytest @@ -26,7 +28,14 @@ from unittest.mock import ANY, MagicMock, patch import litellm from litellm._logging import verbose_logger -from litellm.types.utils import Choices, EmbeddingResponse, Message, ModelResponse, ModelResponseStream +from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + Choices, + EmbeddingResponse, + Message, + ModelResponse, + ModelResponseStream, +) from unittest import mock @@ -1043,6 +1052,132 @@ class TestOllamaToolCallTransformation: assert "tool_call_id" in tool_msg, "tool_call_id must be forwarded to Ollama" assert tool_msg["tool_call_id"] == "call_abc123" + @staticmethod + def _replayed_history(arguments: object) -> list[AllMessageValues]: + return cast( + list[AllMessageValues], + [ + {"role": "user", "content": "Read a.rs and b.rs"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_abc123", + "type": "function", + "function": {"name": "read_file", "arguments": arguments}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_abc123", "content": "fn main() {}"}, + ], + ) + + @pytest.mark.parametrize( + ("arguments", "expected"), + [ + ("", {}), + (None, {}), + ("null", {}), + ('{"path": "a.rs"', {"path": "a.rs"}), + ('{"path": "a.rs", "limit": 2}', {"path": "a.rs", "limit": 2}), + ({"path": "a.rs"}, {"path": "a.rs"}), + ], + ids=["empty", "null", "null_string", "truncated", "non_string_value", "object"], + ) + def test_replayed_tool_call_arguments_are_forwarded_as_an_object(self, arguments: object, expected: object): + result: Final = OllamaChatConfig().transform_request( + model="qwen3:8b", + messages=self._replayed_history(arguments), + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert result["messages"][1]["tool_calls"] == [ + {"id": "call_abc123", "function": {"name": "read_file", "arguments": expected}} + ] + + @pytest.mark.parametrize( + "arguments", + ['{"path": "a.rs"}{"path": "b.rs"}', '["a.rs", "b.rs"]', '[["path", "a.rs"]]', "[]"], + ids=["glued", "array", "pairs", "empty_array"], + ) + def test_replayed_tool_call_arguments_that_are_not_an_object_answer_400_without_echoing_them( + self, arguments: str + ): + with pytest.raises(litellm.BadRequestError) as exc_info: + litellm.completion( + model="ollama_chat/qwen3:8b", + messages=self._replayed_history(arguments), + api_base="http://127.0.0.1:9", + ) + + assert exc_info.value.status_code == 400 + assert "call_abc123" in str(exc_info.value) + assert "a.rs" not in "".join(traceback.format_exception(exc_info.value)) + + +class TestOllamaStreamingParallelToolCallIndex: + """Regression for https://github.com/BerriAI/litellm/issues/33678. + + Ollama streams each parallel tool call in its own chunk without an OpenAI-style top-level + index. Every delta used to carry index 0, so index-accumulating clients appended the + arguments of every later call onto the first one. + """ + + @staticmethod + def _tool_call(index: int, command: str) -> dict[str, object]: + return { + "id": f"call_{index}", + "function": {"index": index, "name": "bash", "arguments": {"command": command}}, + } + + @staticmethod + def _chunk(*tool_calls: dict[str, object]) -> dict[str, object]: + return { + "model": "qwen3:8b", + "message": {"role": "assistant", "content": "", "tool_calls": list(tool_calls)}, + "done": False, + } + + @staticmethod + def _parsed_tool_calls( + iterator: OllamaChatCompletionResponseIterator, *chunks: dict[str, object] + ) -> tuple[ChatCompletionDeltaToolCall, ...]: + return tuple( + itertools.chain.from_iterable( + iterator.chunk_parser(chunk).choices[0].delta.tool_calls or () for chunk in chunks + ) + ) + + def test_tool_calls_in_separate_chunks_get_distinct_sequential_indices(self): + iterator: Final = OllamaChatCompletionResponseIterator(streaming_response=iter([]), sync_stream=True) + chunks: Final = [ + self._chunk(self._tool_call(index, command)) + for index, command in enumerate(("date", "uname -a", "whoami")) + ] + + tool_calls: Final = self._parsed_tool_calls(iterator, *chunks) + + assert [tool_call.index for tool_call in tool_calls] == [0, 1, 2] + assert [json.loads(tool_call.function.arguments) for tool_call in tool_calls] == [ + {"command": "date"}, + {"command": "uname -a"}, + {"command": "whoami"}, + ] + + def test_index_counter_continues_across_chunks(self): + iterator: Final = OllamaChatCompletionResponseIterator(streaming_response=iter([]), sync_stream=True) + chunks: Final = [ + self._chunk(self._tool_call(0, "date"), self._tool_call(1, "uname -a")), + self._chunk(self._tool_call(2, "whoami")), + ] + + tool_calls: Final = self._parsed_tool_calls(iterator, *chunks) + + assert [tool_call.index for tool_call in tool_calls] == [0, 1, 2] + class TestOllamaStreamingUsage: @staticmethod diff --git a/tests/unit/llms/openai/transcriptions/test_gpt_transcribe.py b/tests/unit/llms/openai/transcriptions/test_gpt_transcribe.py index 9bb5fb6d300..4eff0651941 100644 --- a/tests/unit/llms/openai/transcriptions/test_gpt_transcribe.py +++ b/tests/unit/llms/openai/transcriptions/test_gpt_transcribe.py @@ -885,7 +885,7 @@ async def _spend_for_one_second_wav( @pytest.mark.asyncio -async def test_custom_named_gpt_transcribe_deployment_bills_the_measured_upload_over_reported_seconds( +async def test_custom_named_gpt_transcribe_deployment_bills_the_reported_seconds_over_the_measured_upload( local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch ) -> None: router: Final = litellm.Router( @@ -902,7 +902,7 @@ async def test_custom_named_gpt_transcribe_deployment_bills_the_measured_upload_ monkeypatch, {"text": "hello", "usage": {"type": "duration", "seconds": 2}}, router ) - assert spend.cost == pytest.approx(_gpt_transcribe_cost_per_second()) + assert spend.cost == pytest.approx(2 * _gpt_transcribe_cost_per_second()) @pytest.mark.asyncio diff --git a/tests/unit/llms/openai/transcriptions/test_transcription_duration_hidden.py b/tests/unit/llms/openai/transcriptions/test_transcription_duration_hidden.py index 61798b49be0..9c0b9bcec13 100644 --- a/tests/unit/llms/openai/transcriptions/test_transcription_duration_hidden.py +++ b/tests/unit/llms/openai/transcriptions/test_transcription_duration_hidden.py @@ -129,7 +129,7 @@ class TestCostCalculatorReadsDurationFromHiddenParams: @patch( # test-quality-ok: isolates duration selection from the model pricing table "litellm.cost_calculator.openai_cost_per_second" ) - def test_completion_cost_prefers_measured_duration_over_provider_usage_duration(self, mock_cost_fn): + def test_completion_cost_prefers_provider_usage_duration_over_measured_duration(self, mock_cost_fn): mock_cost_fn.return_value = (0.001, 0.0) response = TranscriptionResponse( text="test", @@ -148,7 +148,7 @@ class TestCostCalculatorReadsDurationFromHiddenParams: mock_cost_fn.assert_called_once() _, kwargs = mock_cost_fn.call_args - assert kwargs["duration"] == 17.5 + assert kwargs["duration"] == 60.0 @pytest.mark.parametrize( ("model", "provider"), diff --git a/tests/unit/llms/opensearch/vector_stores/test_opensearch_transformation.py b/tests/unit/llms/opensearch/vector_stores/test_opensearch_transformation.py index 07836d57af0..c55393f2e06 100644 --- a/tests/unit/llms/opensearch/vector_stores/test_opensearch_transformation.py +++ b/tests/unit/llms/opensearch/vector_stores/test_opensearch_transformation.py @@ -4,6 +4,7 @@ import re from collections.abc import Callable, Mapping from datetime import datetime, timezone from pathlib import Path +from types import MappingProxyType from typing import Final from unittest.mock import MagicMock @@ -12,7 +13,7 @@ import openai import pytest from botocore.auth import SigV4Auth from botocore.awsrequest import AWSRequest -from botocore.credentials import Credentials +from botocore.credentials import Credentials, RefreshableCredentials from botocore.exceptions import CredentialRetrievalError import litellm @@ -20,6 +21,12 @@ from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.opensearch.vector_stores.index_passthrough import ( + IndexPath, + IndexRefusal, + index_operation, + prepare_index_request, +) from litellm.llms.opensearch.vector_stores.transformation import OpenSearchVectorStoreConfig from litellm.types.llms.bedrock import AwsAuthParams from litellm.types.router import GenericLiteLLMParams @@ -286,6 +293,17 @@ _AMBIENT_AWS_ENV: Final = ( ) + +def _isolate_from_ambient_aws(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + for name in _AMBIENT_AWS_ENV: + monkeypatch.delenv(name, raising=False) + for name in ("AWS_CONFIG_FILE", "AWS_SHARED_CREDENTIALS_FILE", "BOTO_CONFIG"): + config_file: Path = tmp_path / name + config_file.write_text("") + monkeypatch.setenv(name, str(config_file)) + monkeypatch.setenv("AWS_EC2_METADATA_DISABLED", "true") + + @pytest.mark.parametrize("asynchronous", [False, True]) @pytest.mark.parametrize( "credentials,env", @@ -313,13 +331,7 @@ async def test_missing_aws_credentials_are_a_401_before_embedding_or_any_request monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: - for name in _AMBIENT_AWS_ENV: - monkeypatch.delenv(name, raising=False) - for name in ("AWS_CONFIG_FILE", "AWS_SHARED_CREDENTIALS_FILE", "BOTO_CONFIG"): - config_file: Path = tmp_path / name - config_file.write_text("") - monkeypatch.setenv(name, str(config_file)) - monkeypatch.setenv("AWS_EC2_METADATA_DISABLED", "true") + _isolate_from_ambient_aws(monkeypatch, tmp_path) for name, value in env.items(): monkeypatch.setenv(name, value) executor: Final = RecordingEmbeddingExecutor() @@ -353,6 +365,107 @@ def test_failing_credential_source_is_a_401_that_keeps_its_output_out_of_the_rep assert CREDENTIAL_SOURCE_OUTPUT not in str(error.value) +@pytest.mark.parametrize( + "sign", + [ + pytest.param( + lambda config: config.validate_environment({}, GenericLiteLLMParams(api_base=SERVERLESS)), + id="vector-store-search", + ), + pytest.param( + lambda config: config.sign_http_request( + "GET", f"{SERVERLESS}/docs/_mapping", b"", "application/json", {"api_base": SERVERLESS} + ), + id="index-pass-through", + ), + ], +) +def test_a_missing_container_credentials_token_file_is_a_500( + sign: Callable[[OpenSearchVectorStoreConfig], object], monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + _isolate_from_ambient_aws(monkeypatch, tmp_path) + monkeypatch.setenv("AWS_CONTAINER_CREDENTIALS_FULL_URI", "http://127.0.0.1:9/v1/credentials") + monkeypatch.setenv("AWS_CONTAINER_AUTHORIZATION_TOKEN_FILE", str(tmp_path / "missing-token")) + + with pytest.raises(litellm.InternalServerError, match="could not load or refresh the AWS credentials") as error: + sign(OpenSearchVectorStoreConfig()) + + assert error.value.status_code == 500 + + +EXPIRED_METADATA: Final = MappingProxyType( + {"access_key": "AKIDEXAMPLE", "secret_key": "secret", "token": "token", "expiry_time": "2020-01-01T00:00:00Z"} +) + + +def _refresh_returning_expired_credentials() -> dict[str, str]: + return dict(EXPIRED_METADATA) + + +def _refresh_after_the_token_file_is_gone() -> dict[str, str]: + raise CredentialRetrievalError(provider="container-role", error_msg=CREDENTIAL_SOURCE_OUTPUT) + + +class _CredentialSourceThatCannotRefresh(BaseAWSLLM): + def __init__(self, refresh: Callable[[], dict[str, str]]) -> None: + super().__init__() + self._refresh = refresh + + def resolve_credentials(self, auth_params: AwsAuthParams, aws_region_name: str | None) -> Credentials: + return RefreshableCredentials.create_from_metadata( + metadata=dict(EXPIRED_METADATA), refresh_using=self._refresh, method="container-role" + ) + + +@pytest.mark.parametrize( + "refresh", [_refresh_returning_expired_credentials, _refresh_after_the_token_file_is_gone] +) +@pytest.mark.parametrize( + "sign", + [ + pytest.param( + lambda config: config.validate_environment({}, GenericLiteLLMParams(api_base=SERVERLESS)), + id="vector-store-search", + ), + pytest.param( + lambda config: config.sign_http_request( + "GET", f"{SERVERLESS}/docs/_mapping", b"", "application/json", {"api_base": SERVERLESS} + ), + id="index-pass-through", + ), + ], +) +def test_credentials_that_fail_to_refresh_at_signing_are_a_500( + sign: Callable[[OpenSearchVectorStoreConfig], object], refresh: Callable[[], dict[str, str]] +) -> None: + config: Final = OpenSearchVectorStoreConfig(aws=_CredentialSourceThatCannotRefresh(refresh)) + + with pytest.raises(litellm.InternalServerError, match="could not load or refresh the AWS credentials") as error: + sign(config) + + assert error.value.status_code == 500 + assert CREDENTIAL_SOURCE_OUTPUT not in str(error.value) + + +@pytest.mark.asyncio +async def test_the_index_pass_through_answers_500_when_credentials_fail_to_refresh() -> None: + refused: Final = await prepare_index_request( + config=OpenSearchVectorStoreConfig(aws=_CredentialSourceThatCannotRefresh(_refresh_after_the_token_file_is_gone)), + method="GET", + path=IndexPath(index_name="docs", api=("_mapping",)), + operation=index_operation("GET", ("_mapping",)), + provider_index="prod-docs", + content_type=None, + query=(), + body=b"", + store_params={"api_base": SERVERLESS}, + can_read=True, + ) + + assert isinstance(refused, IndexRefusal) + assert (refused.status_code, CREDENTIAL_SOURCE_OUTPUT in refused.message) == (500, False) + + @pytest.mark.parametrize("asynchronous", [False, True]) @pytest.mark.asyncio async def test_opensearch_index_name_picks_the_index_so_the_store_id_can_differ(asynchronous: bool) -> None: diff --git a/tests/unit/llms/test_cache_control_and_reasoning.py b/tests/unit/llms/test_cache_control_and_reasoning.py index be1ba1e7dbd..82406d114a9 100644 --- a/tests/unit/llms/test_cache_control_and_reasoning.py +++ b/tests/unit/llms/test_cache_control_and_reasoning.py @@ -76,6 +76,19 @@ def test_minimax_supports_thinking_param(): assert "reasoning_split" in supported_params +def test_minimax_forwards_reasoning_effort(): + config = MinimaxChatConfig() + + optional_params = config.map_openai_params( + non_default_params={"reasoning_effort": "low"}, + optional_params={}, + model="minimax/MiniMax-M3.1-Flash", + drop_params=True, + ) + + assert optional_params == {"reasoning_effort": "low"} + + def test_zai_preserves_cache_control_in_messages(): """ZAI should NOT strip cache_control from messages.""" config = ZAIChatConfig() diff --git a/tests/unit/llms/test_file_content_block.py b/tests/unit/llms/test_file_content_block.py index 297b3b6f7d6..6e605eee033 100644 --- a/tests/unit/llms/test_file_content_block.py +++ b/tests/unit/llms/test_file_content_block.py @@ -6,7 +6,9 @@ Regression tests for: - litellm/llms/vertex_ai/gemini/transformation.py - litellm/llms/gemini/chat/transformation.py - litellm/litellm_core_utils/prompt_templates/common_utils.py - (migrate_file_to_image_url raises on missing `file`; file-id helpers skip non-OpenAI shapes) + (migrate_file_to_image_url raises on missing `file`) +- litellm/litellm_core_utils/prompt_templates/file_id_utils.py + (file-id helpers skip non-OpenAI shapes) - litellm/litellm_core_utils/prompt_templates/factory.py (Bedrock + Anthropic) - litellm/llms/openai/chat/gpt_transformation.py """ diff --git a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py index 80e7e793f6e..117d148c287 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py +++ b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py @@ -13,6 +13,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +import respx from fastapi import Request, Response from fastapi.responses import StreamingResponse from fastapi.testclient import TestClient @@ -1222,6 +1223,49 @@ def test_count_tokens_rejected_before_counting_logs_the_400_it_answers( ) +def test_count_tokens_on_bedrock_answers_400_for_a_message_without_content_without_calling_bedrock( + monkeypatch: pytest.MonkeyPatch, httpx_transport: None, respx_mock: respx.MockRouter +) -> None: + from fastapi import FastAPI + + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.anthropic_endpoints.endpoints import router + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.setattr( + proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "count-bedrock-claude", + "litellm_params": { + "model": "bedrock/us.anthropic.claude-zeta-9", + "aws_region_name": "us-east-1", + "aws_access_key_id": "AKIATESTACCESSKEY", + "aws_secret_access_key": "test-secret", + }, + "model_info": {"id": "count-bedrock-claude-deployment", "token_count_framing": 3}, + } + ] + ), + ) + app: Final = FastAPI() + app.include_router(router) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth() + + response: Final = TestClient(app).post( + "/v1/messages/count_tokens", json={"model": "count-bedrock-claude", "messages": [{"role": "user"}]} + ) + + assert response.status_code == 400, response.text + assert response.json()["error"]["type"] == "invalid_request_error" + assert response.json()["error"]["message"] == "Message 0 must have a 'content' field" + assert len(respx_mock.calls) == 0 + + class TestEventLoggingBatchEndpoint: """Test the stubbed event logging batch endpoint""" diff --git a/tests/unit/proxy/auth/test_access_group_rules.py b/tests/unit/proxy/auth/test_access_group_rules.py new file mode 100644 index 00000000000..e73acda7173 --- /dev/null +++ b/tests/unit/proxy/auth/test_access_group_rules.py @@ -0,0 +1,245 @@ +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import pytest +from pydantic import ValidationError + +from litellm.proxy._types import ConfigGeneralSettings, LiteLLM_AccessGroupTable, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.access_group_rules import AccessGroupRule, access_group_ids_for_identity +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + +@dataclass(frozen=True) +class _Database: + db: object + + +@pytest.mark.parametrize( + "claims,expected", + ( + ({"groups": ["science"], "tenant": {"id": "production"}}, ("science-id", "standard-id")), + ({"groups": ["other"]}, ()), + ({"groups": "science", "tenant": {"id": ["production"]}}, ()), + ({}, ()), + ({"groups": ["science", "science"]}, ("science-id",)), + ({"groups": [{"name": "science"}]}, ()), + ), +) +@pytest.mark.asyncio +async def test_verified_attributes_grant_union_with_strict_claim_shapes( + claims: Mapping[str, object], expected: tuple[str, ...] +) -> None: + rules: Final = ( + AccessGroupRule(attribute="jwt_claims.groups", contains="science", access_group="science"), + AccessGroupRule(attribute="jwt_claims.tenant.id", equals="production", access_group="standard"), + ) + cache: Final = UserApiKeyCache() + for name in ("science", "standard"): + group: Final = LiteLLM_AccessGroupTable(access_group_id=f"{name}-id", access_group_name=name) + cache.set_cache(f"access_group_rule_name:{name}", group, model_type=LiteLLM_AccessGroupTable) + cache.set_cache(f"access_group_id:{name}-id", group, model_type=LiteLLM_AccessGroupTable) + assert ( + await access_group_ids_for_identity( + UserAPIKeyAuth(jwt_claims=dict(claims)), rules=rules, cache=cache, prisma_client=_Database(object()) + ) + == expected + ) + + +@pytest.mark.parametrize("rule", ({"contains": "science", "equals": "science"}, {})) +def test_invalid_rule_fails_general_settings_validation(rule: Mapping[str, str]) -> None: + with pytest.raises(ValidationError, match="exactly one"): + ConfigGeneralSettings.model_validate( + {"access_group_rules": [{"attribute": "jwt_claims.groups", "access_group": "science", **rule}]} + ) + + +@pytest.mark.parametrize("attribute", ("metadata.department", "user_email", "user_email.domain", "team_id", "org_id")) +def test_caller_editable_attributes_are_rejected_as_grant_sources(attribute: str) -> None: + with pytest.raises(ValidationError, match="administrator- or IdP-controlled"): + ConfigGeneralSettings.model_validate( + {"access_group_rules": [{"attribute": attribute, "equals": "science", "access_group": "science"}]} + ) + + +@pytest.mark.asyncio +async def test_missing_group_and_no_database_grant_nothing() -> None: + rules: Final = (AccessGroupRule(attribute="user_id", equals="subject", access_group="absent"),) + table: Final = MagicMock() + table.find_unique = AsyncMock(return_value=None) + cache: Final = UserApiKeyCache() + assert ( + await access_group_ids_for_identity( + UserAPIKeyAuth(user_id="subject"), + rules=rules, + cache=cache, + prisma_client=_Database(MagicMock(litellm_accessgrouptable=table)), + ) + == () + ) + assert await access_group_ids_for_identity(UserAPIKeyAuth(user_id="subject"), rules=rules, cache=cache) == () + assert await access_group_ids_for_identity( + UserAPIKeyAuth(user_id="subject"), rules=rules, prisma_client=_Database(object()) + ) == () + + +@pytest.mark.asyncio +async def test_database_failure_does_not_turn_into_an_unrestricted_grant() -> None: + table: Final = MagicMock() + table.find_unique = AsyncMock(side_effect=RuntimeError("database unavailable")) + with pytest.raises(RuntimeError, match="database unavailable"): + await access_group_ids_for_identity( + UserAPIKeyAuth(user_id="subject"), + rules=(AccessGroupRule(attribute="user_id", equals="subject", access_group="science"),), + cache=UserApiKeyCache(), + prisma_client=_Database(MagicMock(litellm_accessgrouptable=table)), + ) + + +@pytest.mark.asyncio +async def test_virtual_key_user_grants_are_additive_and_do_not_mutate_cached_identity() -> None: + identity: Final = UserAPIKeyAuth(models=["explicit"], user_id="science", access_group_ids=["existing"]) + group: Final = LiteLLM_AccessGroupTable( + access_group_id="science-id", access_group_name="science", access_model_names=["group-model"] + ) + cache: Final = UserApiKeyCache() + cache.set_cache("access_group_rule_name:science", group, model_type=LiteLLM_AccessGroupTable) + cache.set_cache("access_group_id:science-id", group, model_type=LiteLLM_AccessGroupTable) + ids: Final = await access_group_ids_for_identity( + identity, + rules=(AccessGroupRule(attribute="user_id", equals="science", access_group="science"),), + cache=cache, + prisma_client=_Database(object()), + ) + granted: Final = identity.model_copy() + granted.apply_access_group_rule_grants(ids) + assert granted.access_group_ids == ["existing", "science-id"] + assert granted.access_group_rule_grants == frozenset({"science-id"}) + assert identity.access_group_ids == ["existing"] + identity.apply_access_group_rule_grants(()) + assert identity.access_group_ids == ["existing"] + + +@pytest.mark.asyncio +async def test_null_equals_does_not_match_an_absent_claim() -> None: + rules: Final = (AccessGroupRule(attribute="jwt_claims.absent", equals=None, access_group="science"),) + assert ( + await access_group_ids_for_identity(UserAPIKeyAuth(jwt_claims={}), rules=rules, cache=UserApiKeyCache()) == () + ) + + +@pytest.mark.parametrize( + "claims,local_email,expected", + ( + ({"email": "science@example.com"}, "other@example.com", ("science-id",)), + ({"email": "other@example.com"}, "science@example.com", ()), + ({}, "science@example.com", ()), + ({"email": ["science@example.com"]}, "science@example.com", ()), + ), +) +@pytest.mark.asyncio +async def test_email_grants_read_verified_claims_not_the_locally_editable_email( + claims: Mapping[str, object], local_email: str, expected: tuple[str, ...] +) -> None: + group: Final = LiteLLM_AccessGroupTable(access_group_id="science-id", access_group_name="science") + cache: Final = UserApiKeyCache() + cache.set_cache("access_group_rule_name:science", group, model_type=LiteLLM_AccessGroupTable) + cache.set_cache("access_group_id:science-id", group, model_type=LiteLLM_AccessGroupTable) + assert ( + await access_group_ids_for_identity( + UserAPIKeyAuth(jwt_claims=dict(claims), user_email=local_email), + rules=( + AccessGroupRule(attribute="jwt_claims.email", equals="science@example.com", access_group="science"), + ), + cache=cache, + prisma_client=_Database(object()), + ) + == expected + ) + + +@pytest.mark.parametrize("operator", ("contains", "equals")) +@pytest.mark.asyncio +async def test_null_operator_remains_distinct_when_rules_are_reloaded(operator: str) -> None: + table: Final = MagicMock() + table.find_unique = AsyncMock(return_value={"access_group_id": "science-id", "access_group_name": "science"}) + cache: Final = UserApiKeyCache() + database: Final = _Database(MagicMock(litellm_accessgrouptable=table)) + rule: Final = AccessGroupRule.model_validate( + {"attribute": "jwt_claims.value", operator: None, "access_group": "science"} + ) + principal: Final = UserAPIKeyAuth(jwt_claims={"value": [None] if operator == "contains" else None}) + assert await access_group_ids_for_identity(principal, rules=(rule,), cache=cache, prisma_client=database) == ( + "science-id", + ) + + +@pytest.mark.asyncio +async def test_authenticated_identity_fields_and_claims_are_available_without_metadata_flattening() -> None: + table: Final = MagicMock() + table.find_unique = AsyncMock(return_value={"access_group_id": "science-id", "access_group_name": "science"}) + identity: Final = UserAPIKeyAuth( + user_id="verified-user", user_role=LitellmUserRoles.INTERNAL_USER, jwt_claims={"groups": ["science"]} + ) + rules: Final = tuple( + AccessGroupRule.model_validate({"attribute": attribute, operator: value, "access_group": "science"}) + for attribute, operator, value in ( + ("user_id", "equals", "verified-user"), + ("jwt_claims.groups", "contains", "science"), + ("user_role", "equals", "internal_user"), + ) + ) + database: Final = _Database(MagicMock(litellm_accessgrouptable=table)) + for index, expected in ((0, ("science-id",)), (1, ("science-id",)), (2, ("science-id",))): + assert ( + await access_group_ids_for_identity( + identity, rules=(rules[index],), cache=UserApiKeyCache(), prisma_client=database + ) + == expected + ) + + +@pytest.mark.asyncio +async def test_name_resolution_is_cached_and_renamed_group_does_not_keep_granting_old_name() -> None: + table: Final = MagicMock() + table.find_unique = AsyncMock(return_value={"access_group_id": "science-id", "access_group_name": "science"}) + database: Final = _Database(MagicMock(litellm_accessgrouptable=table)) + cache: Final = UserApiKeyCache() + rules: Final = (AccessGroupRule(attribute="user_id", equals="subject", access_group="science"),) + assert await access_group_ids_for_identity( + UserAPIKeyAuth(user_id="subject"), rules=rules, cache=cache, prisma_client=database + ) == ("science-id",) + cache.set_cache( + "access_group_id:science-id", + LiteLLM_AccessGroupTable(access_group_id="science-id", access_group_name="science"), + model_type=LiteLLM_AccessGroupTable, + ) + assert await access_group_ids_for_identity( + UserAPIKeyAuth(user_id="subject"), rules=rules, cache=cache, prisma_client=database + ) == ("science-id",) + assert table.find_unique.await_count == 1 + cache.set_cache( + "access_group_id:science-id", + LiteLLM_AccessGroupTable(access_group_id="science-id", access_group_name="renamed"), + model_type=LiteLLM_AccessGroupTable, + ) + table.find_unique.return_value = None + assert ( + await access_group_ids_for_identity( + UserAPIKeyAuth(user_id="subject"), rules=rules, cache=cache, prisma_client=database + ) + == () + ) + + +@pytest.mark.parametrize( + "attribute", + ("groups", "department", "metadata.department", "models", "access_group_rule_grants", "jwt_claims..groups"), +) +def test_attribute_requires_a_real_principal_field(attribute: str) -> None: + with pytest.raises(ValidationError, match="attribute"): + ConfigGeneralSettings.model_validate( + {"access_group_rules": [{"attribute": attribute, "contains": "science", "access_group": "science"}]} + ) diff --git a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py index 78ce9728f3b..e7bcc696c0d 100644 --- a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py +++ b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py @@ -6080,6 +6080,71 @@ class TestGuardrailModificationCheck: self._call({"metadata": '"just a string"'}) +def _member_without_row_prisma(max_budget: float) -> MagicMock: + budget_row: Final = MagicMock() + budget_row.max_budget = max_budget + budget_row.dict = MagicMock(return_value={"budget_id": "budget-default", "max_budget": max_budget}) + prisma_client: Final = MagicMock() + prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=budget_row) + prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=None) + return prisma_client + + +async def _check_member_without_membership_row(prisma_client: MagicMock) -> None: + from litellm.caching.dual_cache import DualCache + from litellm.proxy.utils import ProxyLogging + + await _check_team_member_budget( + team_object=LiteLLM_TeamTable(team_id="team-1", metadata={"team_member_budget_id": "budget-default"}), + user_object=LiteLLM_UserTable(user_id="user-1"), + valid_token=UserAPIKeyAuth(token="tok-1", user_id="user-1", team_id="team-1"), + prisma_client=prisma_client, + user_api_key_cache=DualCache(), + proxy_logging_obj=ProxyLogging(user_api_key_cache=None), + ) + + +@pytest.mark.asyncio +async def test_fail_closed_rejects_member_without_membership_row_with_a_403_naming_member_add( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from fastapi import HTTPException + + import litellm.proxy.proxy_server as ps + from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters + + spend: Final = FakeSpendCounters(redis={}) + prisma_client: Final = _member_without_row_prisma(100.0) + monkeypatch.setattr(ps, "general_settings", {"fail_closed_budget_enforcement": True}) + monkeypatch.setattr(ps, "spend_counter_store", spend.store()) + monkeypatch.setattr(ps, "prisma_client", prisma_client) + + with pytest.raises(HTTPException) as exc_info: + await _check_member_without_membership_row(prisma_client) + + assert exc_info.value.status_code == 403 + message: Final = str(exc_info.value.detail) + assert "user-1" in message and "team-1" in message and "/team/member_add" in message + prisma_client.db.litellm_teammembership.find_unique.assert_awaited_once() + assert spend.commands == [] + + +@pytest.mark.asyncio +async def test_fail_closed_off_still_admits_member_without_membership_row(monkeypatch: pytest.MonkeyPatch) -> None: + import litellm.proxy.proxy_server as ps + from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters + + spend: Final = FakeSpendCounters(redis={}) + prisma_client: Final = _member_without_row_prisma(100.0) + monkeypatch.setattr(ps, "general_settings", {}) + monkeypatch.setattr(ps, "spend_counter_store", spend.store()) + monkeypatch.setattr(ps, "prisma_client", prisma_client) + + await _check_member_without_membership_row(prisma_client) + + assert spend.commands, "the member budget check must still read the spend counter with the flag off" + + @pytest.mark.asyncio async def test_team_member_budget_check_falls_back_to_team_default_budget_id(): """When a member's TeamMembership has no linked budget row, the check @@ -11850,7 +11915,9 @@ async def test_baggage_mapping_rejects_exactly_the_client_tags_rejected_without_ unmapped_request = _client_tag_request(client_tags, baggage=False) unmapped = _client_tag_body(client_tags, client_slot) LiteLLMProxyRequestSetup.pre_seed_litellm_metadata_for_route(unmapped, route) - LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth(unmapped_request, unmapped, UserAPIKeyAuth(api_key="qa-key")) + LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth( + unmapped_request, unmapped, UserAPIKeyAuth(api_key="qa-key") + ) mapped_request = _client_tag_request(client_tags, baggage=baggage) mapped = _client_tag_body(client_tags, client_slot) diff --git a/tests/unit/proxy/auth/test_custom_auth_end_user_budget.py b/tests/unit/proxy/auth/test_custom_auth_end_user_budget.py index 0c3487255f6..e371139dced 100644 --- a/tests/unit/proxy/auth/test_custom_auth_end_user_budget.py +++ b/tests/unit/proxy/auth/test_custom_auth_end_user_budget.py @@ -1,5 +1,7 @@ import pytest +from typing import Final from unittest.mock import AsyncMock, patch +from starlette.requests import Request import litellm from litellm.proxy.auth.user_api_key_auth import ( _run_post_custom_auth_checks, @@ -19,6 +21,7 @@ async def test_custom_auth_run_post_custom_auth_checks_without_end_user_id(): # custom_auth_run_common_checks for custom-auth deployments. The # helper itself no longer calls common_checks. valid_token = UserAPIKeyAuth(token="test_token") + http_request: Final = Request({"type": "http", "method": "POST", "path": "/v1/chat/completions"}) # Default: common_checks should NOT be called inside the helper with patch( @@ -27,7 +30,7 @@ async def test_custom_auth_run_post_custom_auth_checks_without_end_user_id(): mock_common.return_value = True result = await _run_post_custom_auth_checks( valid_token=valid_token, - request=None, + request=http_request, request_data={}, route="/v1/chat/completions", parent_otel_span=None, @@ -50,7 +53,7 @@ async def test_custom_auth_run_post_custom_auth_checks_without_end_user_id(): mock_common.return_value = True result = await _run_post_custom_auth_checks( valid_token=valid_token, - request=None, + request=http_request, request_data={}, route="/v1/chat/completions", parent_otel_span=None, @@ -137,6 +140,7 @@ async def test_custom_auth_defers_end_user_budget_to_common_checks_when_enabled( # With custom_auth_run_common_checks set, the wrapper's common_checks() # enforces the end-user budget, so the helper must not double-enforce it. valid_token = UserAPIKeyAuth(token="test_token", end_user_id="customer-1") + http_request: Final = Request({"type": "http", "method": "POST", "path": "/v1/chat/completions"}) end_user_obj = LiteLLM_EndUserTable( user_id="customer-1", blocked=False, @@ -165,7 +169,7 @@ async def test_custom_auth_defers_end_user_budget_to_common_checks_when_enabled( ): await _run_post_custom_auth_checks( valid_token=valid_token, - request=None, + request=http_request, request_data={"model": "gpt-4"}, route="/v1/chat/completions", parent_otel_span=None, diff --git a/tests/unit/proxy/auth/test_handle_jwt.py b/tests/unit/proxy/auth/test_handle_jwt.py index ae6bf95c13f..36ba8a734b8 100644 --- a/tests/unit/proxy/auth/test_handle_jwt.py +++ b/tests/unit/proxy/auth/test_handle_jwt.py @@ -1,4 +1,5 @@ import asyncio +import json import re import sys import time @@ -8,26 +9,32 @@ from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx +import jwt import pytest +from cryptography.hazmat.primitives.asymmetric import rsa from fastapi import HTTPException +from starlette.requests import Request +import litellm.proxy.proxy_server as proxy_server from litellm.caching.dual_cache import DualCache from litellm.proxy._types import ( DEFAULT_JWKS_STALE_TTL, - JWTLiteLLMRoleMap, JWTIssuerConfig, - LiteLLMRoutes, + JWTLiteLLMRoleMap, + LiteLLM_AccessGroupTable, LiteLLM_JWTAuth, LiteLLM_ModelTable, LiteLLM_TeamMembership, LiteLLM_TeamTable, LiteLLM_UserTable, + LiteLLMRoutes, LitellmUserRoles, Member, ProxyErrorTypes, ProxyException, RoleBasedPermissions, ScopeMapping, + UserAPIKeyAuth, ) from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry from litellm.proxy.auth.auth_checks import TeamNotFoundError @@ -44,10 +51,315 @@ from litellm.proxy.auth.handle_jwt import ( NoMatchingJWTPublicKeyError, ) from litellm.proxy.auth.model_access_denied import ModelAccessDeniedHTTPException +from litellm.proxy.auth.user_api_key_auth import ( + user_api_key_auth, +) from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.types.agents import AgentResponse +async def _jwt_team_selection_identity( + monkeypatch: pytest.MonkeyPatch, + teams: tuple[LiteLLM_TeamTable, ...], + claimed_teams: tuple[str, ...], + membership_teams: tuple[str, ...], + requested_model: str | None, + *, + authorize: bool = False, + fallback_to_db_teams: bool | None = None, + matching_rule: bool | None = None, + scope_models: tuple[str, ...] | None = None, + role_models: tuple[str, ...] | None = None, +) -> UserAPIKeyAuth: + cache: Final = UserApiKeyCache() + database: Final = MagicMock() + handler: Final = JWTHandler() + handler.update_environment( + prisma_client=database, + user_api_key_cache=cache, + litellm_jwtauth=LiteLLM_JWTAuth( + user_id_jwt_field="sub", + team_ids_jwt_field="groups", + enforce_team_based_model_access=True, + enforce_rbac=role_models is not None, + fallback_to_db_teams=not claimed_teams if fallback_to_db_teams is None else fallback_to_db_teams, + enforce_scope_based_access=scope_models is not None, + scope_mappings=[ScopeMapping(scope="inference", models=list(scope_models))] + if scope_models is not None + else None, + ), + ) + private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + jwks_url: Final = "https://identity.invalid/selection/jwks" + cache.set_cache( + f"litellm_jwt_auth_keys_{jwks_url}", + [json.loads(jwt.algorithms.RSAAlgorithm.to_jwk(private_key.public_key()))], + ) + cache.set_cache( + "selection-user", + LiteLLM_UserTable(user_id="selection-user", user_role="internal_user", teams=list(membership_teams)), + model_type=LiteLLM_UserTable, + ) + for team in teams: + cache.set_cache(f"team_id:{team.team_id}", team, model_type=LiteLLM_TeamTable) + cache.set_cache( + f"team_membership:selection-user:{team.team_id}", + LiteLLM_TeamMembership(user_id="selection-user", team_id=team.team_id), + model_type=LiteLLM_TeamMembership, + ) + monkeypatch.setenv("JWT_PUBLIC_KEY_URL", jwks_url) + monkeypatch.delenv("JWT_ISSUER", raising=False) + monkeypatch.delenv("JWT_AUDIENCE", raising=False) + monkeypatch.setattr(proxy_server, "llm_router", None) + monkeypatch.setattr(proxy_server, "jwt_handler", handler) + monkeypatch.setattr(proxy_server, "prisma_client", database) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + monkeypatch.setattr(proxy_server, "premium_user", True) + monkeypatch.setattr(proxy_server, "master_key", "selection-master") + monkeypatch.setattr(proxy_server, "user_custom_auth", None) + monkeypatch.setattr( + proxy_server, + "general_settings", + { + "enable_jwt_auth": True, + **({"role_permissions": [RoleBasedPermissions( + role=LitellmUserRoles.INTERNAL_USER, models=list(role_models), routes=["/chat/completions", "/models"] + )]} if role_models is not None else {}), + **( + { + "access_group_rules": [ + { + "attribute": "jwt_claims.sub", + "equals": "selection-user" if matching_rule else "another-user", + "access_group": "selection-group", + } + ] + } + if matching_rule is not None + else {} + ), + }, + ) + group: Final = LiteLLM_AccessGroupTable( + access_group_id="selection-group-id", + access_group_name="selection-group", + access_model_names=["selection-model"], + ) + cache.set_cache("access_group_rule_name:selection-group", group, model_type=LiteLLM_AccessGroupTable) + cache.set_cache("access_group_id:selection-group-id", group, model_type=LiteLLM_AccessGroupTable) + token: Final = jwt.encode( + {"sub": "selection-user", "groups": list(claimed_teams), "scope": "inference"}, private_key, algorithm="RS256" + ) + request_data: Final = {"model": requested_model} if requested_model is not None else {} + route: Final = "/chat/completions" if requested_model is not None else "/models" + try: + if authorize: + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": route, + "headers": [(b"content-type", b"application/json")], + "query_string": b"", + } + ) + request._body = json.dumps(request_data).encode() + return await user_api_key_auth(request=request, api_key=f"Bearer {token}") + result: Final = await JWTAuthManager.auth_builder( + api_key=token, + jwt_handler=handler, + request_data=request_data, + general_settings=proxy_server.general_settings, + route=route, + prisma_client=database, + user_api_key_cache=cache, + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + identity: Final = JWTAuthManager.user_api_key_auth_from_result(result) + return identity + finally: + await handler.close() + + +@pytest.mark.asyncio +async def test_unresolved_claim_keeps_model_denial_after_single_membership_is_selected( + monkeypatch: pytest.MonkeyPatch, +) -> None: + membership_team: Final = LiteLLM_TeamTable( + team_id="membership-team", + models=["selection-model"], + members_with_roles=[Member(user_id="selection-user", role="user")], + ) + with pytest.raises(ProxyException) as denied: + await _jwt_team_selection_identity( + monkeypatch, + (membership_team,), + ("unresolved-claim",), + ("membership-team",), + "selection-model", + authorize=True, + ) + assert int(denied.value.code) == 403 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("claimed_teams", (("team-c", "team-a", "team-b"), ("team-b", "team-c", "team-a"))) +@pytest.mark.parametrize("requested_model,expected_team", (("selection-model", "team-b"), (None, "team-a"))) +async def test_auth_builder_claimed_teams_attribute_to_first_sorted_eligible_team( + monkeypatch: pytest.MonkeyPatch, + claimed_teams: tuple[str, ...], + requested_model: str | None, + expected_team: str, +) -> None: + teams: Final = tuple( + LiteLLM_TeamTable( + team_id=team_id, + models=["other-model"] if team_id == "team-a" else ["selection-model"], + members_with_roles=[Member(user_id="selection-user", role="user")], + ) + for team_id in ("team-a", "team-b", "team-c") + ) + identity: Final = await _jwt_team_selection_identity( + monkeypatch, teams, claimed_teams, ("team-c", "team-a", "team-b"), requested_model + ) + assert identity.team_id == expected_team + assert identity.user_id == "selection-user" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("membership_teams", (("team-a", "team-c", "team-b"), ("team-a", "team-b", "team-c"))) +async def test_auth_builder_db_teams_attribute_to_first_model_allowed_membership( + monkeypatch: pytest.MonkeyPatch, membership_teams: tuple[str, ...] +) -> None: + teams: Final = tuple( + LiteLLM_TeamTable( + team_id=team_id, + models=["other-model"] if team_id == "team-a" else ["selection-model"], + members_with_roles=[Member(user_id="selection-user", role="user")], + ) + for team_id in membership_teams + ) + identity: Final = await _jwt_team_selection_identity(monkeypatch, teams, (), membership_teams, "selection-model") + assert identity.team_id == membership_teams[1] + assert identity.user_id == "selection-user" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("matching", (True, False)) +async def test_jwt_rule_grants_a_claimed_teams_missing_model_without_losing_attribution( + monkeypatch: pytest.MonkeyPatch, + matching: bool, +) -> None: + team: Final = LiteLLM_TeamTable( + team_id="claimed-team", + models=["other-model"], + members_with_roles=[Member(user_id="selection-user", role="user")], + ) + if matching: + identity: Final = await _jwt_team_selection_identity( + monkeypatch, + (team,), + ("claimed-team",), + ("claimed-team",), + "selection-model", + authorize=True, + matching_rule=matching, + ) + assert identity.team_id == "claimed-team" + assert identity.user_id == "selection-user" + else: + with pytest.raises(ProxyException) as denied: + await _jwt_team_selection_identity( + monkeypatch, + (team,), + ("claimed-team",), + ("claimed-team",), + "selection-model", + authorize=True, + matching_rule=matching, + ) + assert denied.value.code == "403" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("requested_model", (None, "selection-model")) +async def test_jwt_native_team_access_with_rules_configured( + monkeypatch: pytest.MonkeyPatch, + requested_model: str | None, +) -> None: + team: Final = LiteLLM_TeamTable( + team_id="native-team", + models=["selection-model"], + members_with_roles=[Member(user_id="selection-user", role="user")], + ) + identity: Final = await _jwt_team_selection_identity( + monkeypatch, + (team,), + ("native-team",), + ("native-team",), + requested_model, + authorize=True, + matching_rule=True, + ) + assert identity.team_id == "native-team" + assert identity.user_id == "selection-user" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("granted", (True, False)) +@pytest.mark.parametrize("policy_kind", ("scope", "role")) +async def test_jwt_role_and_scope_model_restrictions_remain_in_builder( + monkeypatch: pytest.MonkeyPatch, + granted: bool, + policy_kind: str, +) -> None: + team: Final = LiteLLM_TeamTable( + team_id="scope-team", + models=["selection-model"], + members_with_roles=[Member(user_id="selection-user", role="user")], + ) + scope_models: Final = ("selection-model",) if granted else () + if granted: + identity: Final = await _jwt_team_selection_identity( + monkeypatch, + (team,), + ("scope-team",), + ("scope-team",), + "selection-model", + authorize=True, + scope_models=scope_models if policy_kind == "scope" else None, + role_models=scope_models if policy_kind == "role" else None, + ) + assert identity.team_id == "scope-team" + else: + with pytest.raises(ProxyException) as denied: + await _jwt_team_selection_identity( + monkeypatch, + (team,), + ("scope-team",), + ("scope-team",), + "selection-model", + authorize=True, + scope_models=scope_models if policy_kind == "scope" else None, + role_models=scope_models if policy_kind == "role" else None, + ) + assert denied.value.code == "403" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("matching", (True, False)) +async def test_teamless_jwt_is_rejected_before_rule_grants( + monkeypatch: pytest.MonkeyPatch, matching: bool, +) -> None: + with pytest.raises(ProxyException, match="No teams found in token") as denied: + await _jwt_team_selection_identity( + monkeypatch, (), (), (), "selection-model", authorize=True, + fallback_to_db_teams=False, matching_rule=matching, + ) + assert denied.value.code == "403" + + @pytest.mark.asyncio async def test_map_user_to_teams_user_already_in_team(): """Test that no action is taken when user is already in team""" @@ -213,7 +525,7 @@ async def test_find_team_with_model_access_reports_passthrough_allowlist_denial( assert exc_info.value.status_code == 403 assert "allowed_passthrough_routes" in exc_info.value.detail assert "requested model" not in exc_info.value.detail - mock_is_auth_enforced_pass_through_route.assert_called_once_with( + mock_is_auth_enforced_pass_through_route.assert_called_with( route="/my-pass-through", method="POST" ) @@ -3846,57 +4158,22 @@ async def test_find_team_with_model_access_enforce_team_based_access_still_raise @pytest.mark.asyncio async def test_find_team_with_model_access_resolved_team_without_model_still_raises_403( - monkeypatch, -): - """Regression guard: when the JWT group claim DOES resolve to a real - LiteLLM team but that team does not grant the requested model, keep the - original 403. Only the unresolved-claim case is softened.""" - import sys - import types - - from fastapi import HTTPException - - from litellm.router import Router - - router = Router( - model_list=[ - {"model_name": "gpt-4o-mini", "litellm_params": {"model": "gpt-4o-mini"}}, - { - "model_name": "gpt-3.5-turbo", - "litellm_params": {"model": "gpt-3.5-turbo"}, - }, - ] + monkeypatch: pytest.MonkeyPatch, +) -> None: + team: Final = LiteLLM_TeamTable( + team_id="real-team", + models=["other-model"], + members_with_roles=[Member(user_id="selection-user", role="user")], ) - proxy_server_module = types.ModuleType("proxy_server") - proxy_server_module.llm_router = router - monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module) - - team = LiteLLM_TeamTable(team_id="real-team", models=["gpt-3.5-turbo"]) - - async def mock_get_team_object(*_args, **_kwargs): - return team - - monkeypatch.setattr( - "litellm.proxy.auth.handle_jwt.get_team_object", mock_get_team_object + identity: Final = await _jwt_team_selection_identity( + monkeypatch, (team,), ("real-team",), ("real-team",), "selection-model" ) - - jwt_handler = JWTHandler() - jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() - - with pytest.raises(HTTPException) as exc_info: - await JWTAuthManager.find_team_with_model_access( - team_ids={"real-team"}, - requested_model="gpt-4o-mini", - route="/chat/completions", - jwt_handler=jwt_handler, - prisma_client=None, - user_api_key_cache=None, - parent_otel_span=None, - proxy_logging_obj=None, + assert identity.team_id == "real-team" + with pytest.raises(ProxyException) as denied: + await _jwt_team_selection_identity( + monkeypatch, (team,), ("real-team",), ("real-team",), "selection-model", authorize=True ) - - assert exc_info.value.status_code == 403 - assert "No team has access to the requested model" in str(exc_info.value.detail) + assert denied.value.code == "403" @pytest.mark.asyncio @@ -5441,44 +5718,22 @@ async def test_resolve_team_from_header_accepts_db_teams_provisionally_under_fal @pytest.mark.asyncio -async def test_find_team_with_model_access_defers_no_team_403_under_db_fallback(): - """find_team_with_model_access raises the early "no teams in token" 403 when - enforcement is on, but defers (returns no team) so auth_builder's DB fallback - can run when fallback_to_db_teams is enabled.""" - jwt_handler = JWTHandler() - jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( - enforce_team_based_model_access=True, - fallback_to_db_teams=False, - ) - with pytest.raises(HTTPException) as exc_info: - await JWTAuthManager.find_team_with_model_access( - team_ids=set(), - requested_model="gpt-4", - route="/chat/completions", - jwt_handler=jwt_handler, - prisma_client=None, - user_api_key_cache=MagicMock(), - parent_otel_span=None, - proxy_logging_obj=MagicMock(), +@pytest.mark.parametrize("fallback", (False, True)) +async def test_auth_builder_preserves_teamless_enforcement_denial( + monkeypatch: pytest.MonkeyPatch, + fallback: bool, +) -> None: + with pytest.raises(HTTPException) as builder_denied: + await _jwt_team_selection_identity( + monkeypatch, (), (), (), "selection-model", fallback_to_db_teams=fallback ) - assert exc_info.value.status_code == 403 - - jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( - enforce_team_based_model_access=True, - fallback_to_db_teams=True, - ) - team_id, team_object = await JWTAuthManager.find_team_with_model_access( - team_ids=set(), - requested_model="gpt-4", - route="/chat/completions", - jwt_handler=jwt_handler, - prisma_client=None, - user_api_key_cache=MagicMock(), - parent_otel_span=None, - proxy_logging_obj=MagicMock(), - ) - assert team_id is None - assert team_object is None + assert builder_denied.value.status_code == 403 + with pytest.raises(ProxyException) as denied: + await _jwt_team_selection_identity( + monkeypatch, (), (), (), "selection-model", authorize=True, fallback_to_db_teams=fallback + ) + assert denied.value.code == "403" + assert ("not a member of any team" if fallback else "No teams found in token") in denied.value.message def _db_fallback_handler(litellm_jwtauth: LiteLLM_JWTAuth | None = None) -> JWTHandler: @@ -6265,82 +6520,23 @@ async def test_auth_builder_db_fallback_does_not_validate_rbac_team_against_db_m @pytest.mark.asyncio -async def test_resolve_db_team_fallback_distinguishes_no_membership_vs_model_denied(): - """When enforce_team_based_model_access is on, a user with no DB memberships - and a user with memberships that all fail the model-access check must surface - different 403s; collapsing both into the no-membership message hides the real - cause and diverges from find_team_with_model_access's claim-based message.""" - membership_user = LiteLLM_UserTable( - user_id="u_no_model", - user_role=LitellmUserRoles.INTERNAL_USER, - teams=["only_team"], +async def test_resolve_db_team_fallback_distinguishes_no_membership_vs_model_denied( + monkeypatch: pytest.MonkeyPatch, +) -> None: + team: Final = LiteLLM_TeamTable( + team_id="only-team", + models=["other-model"], + members_with_roles=[Member(user_id="selection-user", role="user")], ) - no_membership_user = LiteLLM_UserTable( - user_id="u_empty", - user_role=LitellmUserRoles.INTERNAL_USER, - teams=[], - ) - - async def fake_get_team(team_id, **kwargs): - return LiteLLM_TeamTable(team_id=team_id, models=["other"]) - - async def fake_can_access(model, team_object, llm_router, team_model_aliases=None): - raise ProxyException( - message="team not allowed to access model", - type=ProxyErrorTypes.team_model_access_denied, - param="model", - code=403, - ) - - with ( - patch( - "litellm.proxy.auth.handle_jwt.get_team_object", - new_callable=AsyncMock, - side_effect=fake_get_team, - ), - patch( - "litellm.proxy.auth.handle_jwt.can_team_access_model", - new_callable=AsyncMock, - side_effect=fake_can_access, - ), - ): - with pytest.raises(HTTPException) as model_denied: - await JWTAuthManager._resolve_db_team_fallback( - user_object=membership_user, - user_id=None, - requested_model="gpt-4", - route="/chat/completions", - jwt_handler=_db_fallback_handler(), - enforce_team_based_model_access=True, - team_id_upsert=False, - prisma_client=None, - user_api_key_cache=MagicMock(), - parent_otel_span=None, - proxy_logging_obj=MagicMock(), - ) - - with pytest.raises(HTTPException) as no_member: - await JWTAuthManager._resolve_db_team_fallback( - user_object=no_membership_user, - user_id=None, - requested_model="gpt-4", - route="/chat/completions", - jwt_handler=_db_fallback_handler(), - enforce_team_based_model_access=True, - team_id_upsert=False, - prisma_client=None, - user_api_key_cache=MagicMock(), - parent_otel_span=None, - proxy_logging_obj=MagicMock(), - ) - - assert model_denied.value.status_code == 403 - assert "requested model" in model_denied.value.detail - assert "gpt-4" in model_denied.value.detail - assert "only_team" not in model_denied.value.detail - - assert no_member.value.status_code == 403 - assert "not a member of any team" in no_member.value.detail + identity: Final = await _jwt_team_selection_identity(monkeypatch, (team,), (), ("only-team",), "selection-model") + assert identity.team_id == "only-team" + with pytest.raises(ProxyException) as model_denied: + await _jwt_team_selection_identity(monkeypatch, (team,), (), ("only-team",), "selection-model", authorize=True) + assert model_denied.value.code == "403" + assert "not a member of any team" not in model_denied.value.message + with pytest.raises(ProxyException, match="not a member of any team") as no_member: + await _jwt_team_selection_identity(monkeypatch, (), (), (), "selection-model", authorize=True) + assert no_member.value.code == "403" @pytest.mark.asyncio diff --git a/tests/unit/proxy/auth/test_user_api_key_auth.py b/tests/unit/proxy/auth/test_user_api_key_auth.py index 19f9d0dcc73..ce4a844550b 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth.py @@ -2,27 +2,29 @@ ## Unit tests for user_api_key_auth helper functions -import litellm.proxy -import litellm.proxy.proxy_server - +import json +import logging from typing import Dict, Final, List, Mapping, Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fastapi import HTTPException, WebSocket, status from starlette.datastructures import URL +from starlette.requests import Request as StarletteRequest from starlette.types import Message -from litellm._logging import verbose_proxy_logger -import json -import logging -import litellm -from litellm.proxy.auth.user_api_key_auth import ( - user_api_key_auth, - UserAPIKeyAuth, - get_api_key_from_custom_header, -) -from fastapi import WebSocket, HTTPException, status -from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles +import litellm +import litellm.proxy +import litellm.proxy.proxy_server +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import LiteLLM_AccessGroupTable, LiteLLM_UserTable, LitellmUserRoles, ProxyException +from litellm.proxy.auth.user_api_key_auth import ( + UserAPIKeyAuth, + _authorize_authenticated_request, + get_api_key_from_custom_header, + user_api_key_auth, +) +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from tests._master_key import MASTER_KEY @@ -33,6 +35,146 @@ class Request: self.headers: Dict[str, str] = {} +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model,fallbacks,allowed", (("model-a", [], True), ("model-b", [], False), ("model-a", ["model-b"], False)) +) +async def test_virtual_key_model_policy_on_first_lookup_and_cache_hit( + monkeypatch: pytest.MonkeyPatch, model: str, fallbacks: list[str], allowed: bool +) -> None: + server: Final = litellm.proxy.proxy_server + cache: Final = UserApiKeyCache() + key: Final = "sk-central-model-check" + stored: Final = UserAPIKeyAuth(token=server.hash_token(key), models=["model-a"]) + database: Final = MagicMock(get_data=AsyncMock(side_effect=(stored, AssertionError("unexpected cache miss")))) + monkeypatch.setattr(server, "master_key", MASTER_KEY) + monkeypatch.setattr(server, "general_settings", {}) + monkeypatch.setattr(server, "user_api_key_cache", cache) + monkeypatch.setattr(server, "prisma_client", database) + monkeypatch.setattr(server, "llm_router", None) + monkeypatch.setattr(server, "user_custom_auth", None) + for _ in range(2): + request: Final = StarletteRequest( + {"type": "http", "method": "POST", "path": "/chat/completions", "headers": []}, + receive=AsyncMock( + return_value={ + "type": "http.request", + "body": json.dumps({"model": model, "fallbacks": fallbacks}).encode(), + "more_body": False, + } + ), + ) + if allowed: + identity: Final = await user_api_key_auth(request=request, api_key=f"Bearer {key}") + assert identity.models == ["model-a"] + else: + with pytest.raises(ProxyException) as denied: + await user_api_key_auth(request=request, api_key=f"Bearer {key}") + assert denied.value.code == "403" + assert "model-b" in denied.value.message + + +@pytest.mark.asyncio +async def test_scim_deactivation_precedes_central_key_model_denial(monkeypatch: pytest.MonkeyPatch) -> None: + server: Final = litellm.proxy.proxy_server + key: Final = "sk-deactivated-owner" + cache: Final = UserApiKeyCache() + stored: Final = UserAPIKeyAuth(token=server.hash_token(key), models=["model-a"], user_id="inactive-owner") + cache.set_cache( + "inactive-owner", + LiteLLM_UserTable(user_id="inactive-owner", metadata={"scim_active": False}), + model_type=LiteLLM_UserTable, + ) + monkeypatch.setattr(server, "master_key", MASTER_KEY) + monkeypatch.setattr(server, "general_settings", {}) + monkeypatch.setattr(server, "user_api_key_cache", cache) + monkeypatch.setattr(server, "prisma_client", MagicMock(get_data=AsyncMock(return_value=stored))) + monkeypatch.setattr(server, "llm_router", None) + monkeypatch.setattr(server, "user_custom_auth", None) + request: Final = StarletteRequest( + {"type": "http", "method": "POST", "path": "/chat/completions", "headers": []}, + receive=AsyncMock(return_value={"type": "http.request", "body": b'{"model":"model-b"}'}), + ) + with pytest.raises(ProxyException) as denied: + await user_api_key_auth(request=request, api_key=f"Bearer {key}") + assert denied.value.code == "401" + assert "deactivated via SCIM" in denied.value.message + + +@pytest.mark.asyncio +async def test_restricted_key_model_list_denial_names_the_denied_model( + monkeypatch: pytest.MonkeyPatch, +) -> None: + server: Final = litellm.proxy.proxy_server + cache: Final = UserApiKeyCache() + group: Final = LiteLLM_AccessGroupTable( + access_group_id="science-id", access_group_name="science", access_model_names=["model-a"] + ) + cache.set_cache("access_group_rule_name:science", group, model_type=LiteLLM_AccessGroupTable) + cache.set_cache("access_group_id:science-id", group, model_type=LiteLLM_AccessGroupTable) + monkeypatch.setattr(server, "user_api_key_cache", cache) + monkeypatch.setattr(server, "prisma_client", MagicMock()) + monkeypatch.setattr( + server, + "general_settings", + {"access_group_rules": [{"attribute": "user_id", "equals": "science", "access_group": "science"}]}, + ) + monkeypatch.setattr(server, "llm_router", None) + identity: Final = UserAPIKeyAuth(user_id="science", models=["model-a"]) + request: Final = StarletteRequest({"type": "http", "method": "POST", "path": "/chat/completions", "headers": []}) + with pytest.raises(ProxyException) as denied: + await _authorize_authenticated_request( + identity, request, {"model": ["model-a", "model-b"]}, "/chat/completions", "synthetic", + force_virtual_key_checks=True, + ) + assert denied.value.code == "403" + assert "model-b" in denied.value.message + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model,enabled,allowed", + (("group-model", True, True), ("outside", True, False), ("group-model", False, False), ("explicit", False, True)), +) +async def test_user_access_group_rules_are_enforced_at_central_authorization( + monkeypatch: pytest.MonkeyPatch, model: str, enabled: bool, allowed: bool +) -> None: + server: Final = litellm.proxy.proxy_server + cache: Final = UserApiKeyCache() + group: Final = LiteLLM_AccessGroupTable( + access_group_id="science-id", access_group_name="science", access_model_names=["group-model"] + ) + cache.set_cache("access_group_rule_name:science", group, model_type=LiteLLM_AccessGroupTable) + cache.set_cache("access_group_id:science-id", group, model_type=LiteLLM_AccessGroupTable) + cache.set_cache("science", LiteLLM_UserTable(user_id="science"), model_type=LiteLLM_UserTable) + monkeypatch.setattr( + server, + "general_settings", + {"access_group_rules": [{"attribute": "user_id", "equals": "science", "access_group": "science"}]} + if enabled + else {}, + ) + monkeypatch.setattr(server, "user_api_key_cache", cache) + monkeypatch.setattr(server, "prisma_client", MagicMock()) + monkeypatch.setattr(server, "llm_router", None) + monkeypatch.setattr(server, "user_custom_auth", None) + principal: Final = UserAPIKeyAuth(models=["explicit"], user_id="science") + request: Final = StarletteRequest({"type": "http", "method": "POST", "path": "/chat/completions", "headers": []}) + if allowed: + assert ( + await _authorize_authenticated_request( + principal, request, {"model": model}, "/chat/completions", "synthetic", force_virtual_key_checks=True + ) + is None + ) + assert principal.access_group_ids == (["science-id"] if enabled else None) + else: + with pytest.raises(ProxyException, match="model"): + await _authorize_authenticated_request( + principal, request, {"model": model}, "/chat/completions", "synthetic", force_virtual_key_checks=True + ) + + @pytest.mark.parametrize( "allowed_ips, client_ip, expected_result", [ diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index 8440d3d2ab2..1489afc02af 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -87,6 +87,7 @@ from litellm.proxy.auth.user_api_key_auth import ( _reserve_budget_after_common_checks, _route_requires_auth_despite_public, _routing_selector_matches_claim, + _authorize_authenticated_request, run_centralized_common_checks, _run_post_custom_auth_checks, user_api_key_auth_builder, @@ -446,29 +447,47 @@ async def test_should_not_reuse_cached_key_object_for_request_state(): assert second_request_key.request_route is None -@pytest.mark.asyncio -async def test_custom_auth_does_not_enforce_key_model_access_by_default(): - valid_token = UserAPIKeyAuth(token="test_token", models=["gpt-4o-mini"]) - request_data = {"model": "gpt-4o"} - - with ( - patch( - "litellm.proxy.auth.user_api_key_auth.can_key_call_model", - new_callable=AsyncMock, - ) as mock_can_key, - patch( - "litellm.proxy.proxy_server.general_settings", - {}, - ), - ): - await _run_post_custom_auth_checks( - valid_token=valid_token, - request=None, - request_data=request_data, - route="/v1/chat/completions", - parent_otel_span=None, +async def _custom_auth_model_request( + monkeypatch: pytest.MonkeyPatch, requested_model: str, *, opt_in: bool, from_header: bool = False +) -> UserAPIKeyAuth: + async def custom_auth(request: Request, api_key: str) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + models=["allowed-model"], + user_role=LitellmUserRoles.INTERNAL_USER, + request_route=request.url.path, + api_key=api_key, ) - mock_can_key.assert_not_awaited() + + server: Final = litellm.proxy.proxy_server + monkeypatch.setattr(litellm, "enable_post_custom_auth_checks", True, raising=False) + monkeypatch.setattr(server, "user_custom_auth", custom_auth) + monkeypatch.setattr(server, "general_settings", {"custom_auth_run_common_checks": opt_in}) + monkeypatch.setattr(server, "prisma_client", None) + monkeypatch.setattr(server, "llm_router", None) + monkeypatch.setattr(server, "llm_model_list", None) + monkeypatch.setattr(server, "user_api_key_cache", UserApiKeyCache()) + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": "/v1/files" if from_header else "/v1/chat/completions", + "headers": [(b"x-litellm-model", requested_model.encode())] if from_header else [], + }, + receive=AsyncMock( + return_value={ + "type": "http.request", + "more_body": False, + "body": json.dumps({} if from_header else {"model": requested_model}).encode(), + } + ), + ) + return await user_api_key_auth(request=request, api_key="Bearer external-token") + + +@pytest.mark.asyncio +async def test_custom_auth_does_not_enforce_key_model_access_by_default(monkeypatch: pytest.MonkeyPatch): + identity: Final = await _custom_auth_model_request(monkeypatch, "restricted-model", opt_in=False) + assert identity.models == ["allowed-model"] @pytest.mark.asyncio @@ -492,99 +511,30 @@ async def test_post_custom_auth_expired_key_returns_unauthorized(): @pytest.mark.asyncio -async def test_custom_auth_honors_key_level_model_access_restriction_allowed_with_opt_in(): - valid_token = UserAPIKeyAuth(token="test_token", models=["gpt-4o-mini"]) - request_data = {"model": "gpt-4o-mini"} - - with ( - patch( - "litellm.proxy.auth.user_api_key_auth.can_key_call_model", - new_callable=AsyncMock, - ) as mock_can_key, - patch("litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock), - patch( - "litellm.proxy.proxy_server.general_settings", - {"custom_auth_run_common_checks": True}, - ), - ): - await _run_post_custom_auth_checks( - valid_token=valid_token, - request=None, - request_data=request_data, - route="/v1/chat/completions", - parent_otel_span=None, - ) - mock_can_key.assert_awaited_once_with( - model="gpt-4o-mini", - llm_model_list=ANY, - valid_token=valid_token, - llm_router=ANY, - ) +async def test_custom_auth_honors_key_level_model_access_restriction_allowed_with_opt_in( + monkeypatch: pytest.MonkeyPatch, +): + identity: Final = await _custom_auth_model_request(monkeypatch, "allowed-model", opt_in=True) + assert identity.models == ["allowed-model"] @pytest.mark.asyncio -async def test_custom_auth_enforces_key_model_access_from_file_route_header_with_opt_in(): - valid_token = UserAPIKeyAuth(token="test_token", models=["allowed-model"]) - request = _RoutingRequest(headers={"x-litellm-model": "restricted-model"}) - - with ( - patch( - "litellm.proxy.auth.user_api_key_auth.can_key_call_model", - new_callable=AsyncMock, - ) as mock_can_key, - patch("litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock), - patch( - "litellm.proxy.proxy_server.general_settings", - {"custom_auth_run_common_checks": True}, - ), - ): - await _run_post_custom_auth_checks( - valid_token=valid_token, - request=request, - request_data={}, - route="/v1/files", - parent_otel_span=None, - ) - mock_can_key.assert_awaited_once_with( - model="restricted-model", - llm_model_list=ANY, - valid_token=valid_token, - llm_router=ANY, - ) +async def test_custom_auth_enforces_key_model_access_from_file_route_header_with_opt_in( + monkeypatch: pytest.MonkeyPatch, +): + with pytest.raises(ProxyException) as denied: + await _custom_auth_model_request(monkeypatch, "restricted-model", opt_in=True, from_header=True) + assert denied.value.code == "403" @pytest.mark.asyncio -async def test_custom_auth_honors_key_level_model_access_restriction_denied_with_opt_in(): - valid_token = UserAPIKeyAuth(token="test_token", models=["gpt-4o-mini"]) - request_data = {"model": "gpt-4o"} - - with ( - patch( - "litellm.proxy.auth.user_api_key_auth.can_key_call_model", - new_callable=AsyncMock, - ) as mock_can_key, - patch("litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock), - patch( - "litellm.proxy.proxy_server.general_settings", - {"custom_auth_run_common_checks": True}, - ), - ): - mock_can_key.side_effect = ProxyException( - message="Key not allowed to access model", - type=ProxyErrorTypes.key_model_access_denied, - param="model", - code=401, - ) - with pytest.raises(ProxyException) as exc: - await _run_post_custom_auth_checks( - valid_token=valid_token, - request=None, - request_data=request_data, - route="/v1/chat/completions", - parent_otel_span=None, - ) - - assert exc.value.type == ProxyErrorTypes.key_model_access_denied +async def test_custom_auth_honors_key_level_model_access_restriction_denied_with_opt_in( + monkeypatch: pytest.MonkeyPatch, +): + with pytest.raises(ProxyException) as denied: + await _custom_auth_model_request(monkeypatch, "restricted-model", opt_in=True) + assert denied.value.type == ProxyErrorTypes.key_model_access_denied + assert denied.value.code == "403" def _proxy_server_attrs_for_custom_auth(*, user_custom_auth): @@ -1830,8 +1780,7 @@ async def test_db_virtual_key_auth_sets_via_virtual_key_marker(): @pytest.mark.asyncio @pytest.mark.parametrize("model_allowed", [True, False]) -async def test_auth_prefetches_referenced_objects_only_after_the_key_may_call_the_model(model_allowed): - """A request denied by the key's model list must not pay for the team/user/org MGET or DB join.""" +async def test_auth_prefetches_referenced_objects_before_central_model_authorization(model_allowed): from fastapi import Request from starlette.datastructures import URL @@ -1841,9 +1790,9 @@ async def test_auth_prefetches_referenced_objects_only_after_the_key_may_call_th api_key = "sk-prefetch-order-test" valid_token = UserAPIKeyAuth(api_key=api_key, token=hash_token(api_key), user_id="u1", team_id="t1") - mock_cache = AsyncMock() - mock_cache.async_get_cache = AsyncMock(return_value=None) - mock_cache.delete_cache = MagicMock() + mock_cache = UserApiKeyCache() + mock_cache.set_cache("team_id:t1", LiteLLM_TeamTable(team_id="t1")) + mock_cache.set_cache("team_membership:u1:t1", LiteLLM_TeamMembership(user_id="u1", team_id="t1")) mock_proxy_logging_obj = MagicMock() mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) @@ -1873,7 +1822,7 @@ async def test_auth_prefetches_referenced_objects_only_after_the_key_may_call_th try: for attr, val in _attrs_to_set.items(): setattr(_proxy_server_mod, attr, val) - request = Request(scope={"type": "http"}) + request = Request(scope={"type": "http", "method": "POST", "path": "/chat/completions"}) request._url = URL(url="/chat/completions") with ( @@ -1894,7 +1843,7 @@ async def test_auth_prefetches_referenced_objects_only_after_the_key_may_call_th "litellm.proxy.auth.user_api_key_auth.get_user_object", new_callable=AsyncMock, return_value=None ), ): - call = user_api_key_auth_builder( + resolved: Final = await user_api_key_auth_builder( request=request, api_key=f"Bearer {api_key}", azure_api_key_header="", @@ -1903,15 +1852,18 @@ async def test_auth_prefetches_referenced_objects_only_after_the_key_may_call_th azure_apim_header=None, request_data={"model": "gpt-4o"}, ) + assert isinstance(resolved, UserAPIKeyAuth) + mock_prefetch.assert_awaited_once() + assert mock_prefetch.await_args.kwargs["refs"].team_id == "t1" + call: Final = _authorize_authenticated_request( + resolved, request, {"model": "gpt-4o"}, "/chat/completions", api_key + ) if model_allowed: - assert isinstance(await call, UserAPIKeyAuth) - mock_prefetch.assert_awaited_once() - assert mock_prefetch.await_args.kwargs["refs"].team_id == "t1" + await call else: with pytest.raises(ProxyException) as exc: await call assert exc.value.type == ProxyErrorTypes.key_model_access_denied - mock_prefetch.assert_not_awaited() finally: for attr, val in _original_values.items(): setattr(_proxy_server_mod, attr, val) @@ -2601,6 +2553,48 @@ async def test_auto_register_map_existing_key_uses_the_issuers_own_user_field_ov assert generate_key.await_count == (0 if expect_reuse else 1) +def _validated_jwt_admission_cache(*, org_id: str | None = None, end_user_id: str | None = None) -> UserApiKeyCache: + cache: Final = UserApiKeyCache() + cache.set_cache( + "validated-user", + LiteLLM_UserTable( + user_id="validated-user", + user_email="validated@example.com", + user_role="internal_user", + teams=["validated-team"], + ), + model_type=LiteLLM_UserTable, + ) + cache.set_cache( + "team_id:validated-team", + LiteLLM_TeamTableCachedObj( + team_id="validated-team", + models=["gpt-4o-mini", "gpt-5.6"], + members_with_roles=[{"user_id": "validated-user", "role": "user"}], + organization_id=org_id, + ), + model_type=LiteLLM_TeamTableCachedObj, + ) + cache.set_cache( + team_membership_reservation_cache_key(user_id="validated-user", team_id="validated-team"), + LiteLLM_TeamMembership(user_id="validated-user", team_id="validated-team"), + model_type=LiteLLM_TeamMembership, + ) + if org_id is not None: + org: Final = LiteLLM_OrganizationTable( + organization_id=org_id, budget_id="org-budget", created_by="admin", updated_by="admin" + ) + cache.set_cache(f"org_id:{org_id}", org, model_type=LiteLLM_OrganizationTable) + cache.set_cache(f"org_id:{org_id}:with_budget", org, model_type=LiteLLM_OrganizationTable) + if end_user_id is not None: + cache.set_cache( + f"end_user_id:{end_user_id}", + LiteLLM_EndUserTable(user_id=end_user_id, blocked=False), + model_type=LiteLLM_EndUserTable, + ) + return cache + + @pytest.mark.asyncio @pytest.mark.parametrize( ("map_existing_key", "master_key", "reused_key_models", "expect_denied"), @@ -2615,7 +2609,7 @@ async def test_auto_register_map_existing_key_first_request_runs_key_checks( map_existing_key: bool, master_key: str | None, reused_key_models: list[str], expect_denied: bool ) -> None: jwt_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature" - user_api_key_cache = DualCache() + user_api_key_cache: Final = _validated_jwt_admission_cache() prisma_client = MagicMock() jwt_handler = MagicMock() jwt_handler.is_jwt.return_value = True @@ -2648,12 +2642,21 @@ async def test_auto_register_map_existing_key_first_request_runs_key_checks( "jwt_claims": {"sub": "user1"}, } - mock_request = MagicMock() - mock_request.url.path = "/v1/chat/completions" - mock_request.method = "POST" - mock_request.headers = {"authorization": f"Bearer {jwt_token}"} - mock_request.query_params = {} - mock_request.state = SimpleNamespace() + mock_request: Final = Request( + { + "type": "http", + "method": "POST", + "path": "/v1/chat/completions", + "headers": [(b"authorization", f"Bearer {jwt_token}".encode())], + }, + receive=AsyncMock( + return_value={ + "type": "http.request", + "more_body": False, + "body": json.dumps({"model": "gpt-4o-mini"}).encode(), + } + ), + ) with ( patch("litellm.proxy.proxy_server.general_settings", {"enable_jwt_auth": True}), @@ -2686,14 +2689,13 @@ async def test_auto_register_map_existing_key_first_request_runs_key_checks( return_value=reused_key, ), ): - call = user_api_key_auth_builder( + call = user_api_key_auth( request=mock_request, api_key=jwt_token, azure_api_key_header="", anthropic_api_key_header=None, google_ai_studio_api_key_header=None, azure_apim_header=None, - request_data={"model": "gpt-4o-mini"}, ) if expect_denied: with pytest.raises(ProxyException, match="not available for this API key"): @@ -2719,7 +2721,7 @@ async def test_auto_register_first_request_propagates_user_email(active: bool) - """ jwt_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature" general_settings = {"enable_jwt_auth": True} - user_api_key_cache = DualCache() + user_api_key_cache: Final = _validated_jwt_admission_cache() prisma_client = MagicMock() jwt_handler = MagicMock() jwt_handler.is_jwt.return_value = True @@ -2952,7 +2954,10 @@ async def test_jwt_auto_register_forwards_bound_agent_id(): id auth_builder resolved must reach the key creation, not be dropped when valid_token is swapped for the freshly registered key.""" jwt_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature" - user_api_key_cache = DualCache() + user_api_key_cache: Final = _validated_jwt_admission_cache() + prisma_client: Final = MagicMock() + prisma_client.writer_db.litellm_agentstable.find_unique = AsyncMock(return_value=None) + prisma_client.writer_db.litellm_retiredagent.find_unique = AsyncMock(return_value=None) jwt_handler = MagicMock() jwt_handler.is_jwt.return_value = True jwt_handler.auth_jwt = AsyncMock(return_value={"sub": "user1", "appid": "2f5c9b1e-6a4d-4c8e-9f0b-7d1a3e5c9b21"}) @@ -3001,7 +3006,7 @@ async def test_jwt_auto_register_forwards_bound_agent_id(): general_settings={"enable_jwt_auth": True}, premium_user=True, master_key="sk-master", - prisma_client=MagicMock(), + prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=MagicMock(), jwt_handler=jwt_handler, @@ -3274,7 +3279,7 @@ class TestJWTOAuth2Coexistence: async def test_auto_register_passes_validated_org_context_to_generated_key(self): jwt_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature" general_settings = {"enable_jwt_auth": True} - user_api_key_cache = DualCache() + user_api_key_cache: Final = _validated_jwt_admission_cache(org_id="validated-org", end_user_id="validated-end-user") prisma_client = MagicMock() jwt_handler = MagicMock() jwt_handler.is_jwt.return_value = True @@ -5094,6 +5099,70 @@ async def test_centralized_common_checks_carries_team_and_user_budget_state_on_t } +@pytest.mark.asyncio +@pytest.mark.parametrize("prometheus_registered", [True, False], ids=["prometheus_logger", "no_prometheus_logger"]) +async def test_centralized_common_checks_carry_the_committed_spend_beside_the_stale_cached_spend( + monkeypatch: pytest.MonkeyPatch, prometheus_registered: bool +): + """The enforcement counters also hold budget holds, the committed twins only charged spend.""" + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as _proxy_server_mod + from litellm.integrations.prometheus import PrometheusLogger + from litellm.proxy.spend_tracking.committed_spend import committed_spend_key + from litellm.proxy.utils import ProxyLogging + from litellm.redis.stores.connection import connect + from litellm.redis.stores.spend_counters import SpendCounterStore + from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys + + for callback_list in ("success_callback", "failure_callback", "_async_success_callback", "_async_failure_callback"): + monkeypatch.setattr(litellm, callback_list, []) + monkeypatch.setattr(litellm, "callbacks", [MagicMock(spec=PrometheusLogger)] if prometheus_registered else []) + token: Final = UserAPIKeyAuth(api_key="sk-live-spend", max_budget=10.0, spend=0.0, team_id="t-live", team_spend=0.0) + user_api_key_cache: Final = UserApiKeyCache() + user_api_key_cache.set_cache( + key="team_id:t-live", value=LiteLLM_TeamTableCachedObj(team_id="t-live", max_budget=20.0, spend=0.0) + ) + spend_counter_store: Final = SpendCounterStore(connect(None).spend_counters()) + for counter_key, held, committed in ((f"spend:key:{token.token}", 0.67, 0.37), ("spend:team:t-live", 1.55, 1.25)): + spend_counter_store.set_local(counter_key, held) + spend_counter_store.set_local(committed_spend_key(counter_key), committed) + attrs: Final = { + **_proxy_attrs_for_centralized_checks(), + "llm_router": None, + "prisma_client": MagicMock(), + "user_api_key_cache": user_api_key_cache, + "spend_counter_store": spend_counter_store, + "proxy_logging_obj": ProxyLogging(user_api_key_cache=user_api_key_cache), + } + originals: Final = {attribute: getattr(_proxy_server_mod, attribute, None) for attribute in attrs} + request: Final = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + try: + for attribute, value in attrs.items(): + setattr(_proxy_server_mod, attribute, value) + await run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 1000}, + route="/chat/completions", + ) + finally: + for attribute, value in originals.items(): + setattr(_proxy_server_mod, attribute, value) + + assert token.committed_budget_spend == ( + CommittedBudgetSpend( + counter_keys=EntitySpendCounterKeys(key=f"spend:key:{token.token}", team=f"spend:team:{token.team_id}"), + key=0.37, + team=1.25, + ) + if prometheus_registered + else None + ) + assert (token.spend, token.team_spend) == (0.0, 0.0) + + def _end_user_budget_row(budget_id: str, max_budget: float) -> MagicMock: row = MagicMock() row.dict = lambda: {"budget_id": budget_id, "max_budget": max_budget} diff --git a/tests/unit/proxy/common_utils/test_callback_config_validation.py b/tests/unit/proxy/common_utils/test_callback_config_validation.py index f42387d3e76..ac1fe2d785e 100644 --- a/tests/unit/proxy/common_utils/test_callback_config_validation.py +++ b/tests/unit/proxy/common_utils/test_callback_config_validation.py @@ -7,6 +7,7 @@ from litellm.proxy.common_utils.callback_config_validation import ( callback_config_error, conflicting_callback_var_error, conflicting_span_scope_error, + credential_reference_rejection, cross_entry_family_error, logging_metadata_config_error, ) @@ -499,3 +500,78 @@ def test_key_logging_entries_compare_excluded_services_by_parsed_value( def test_a_failure_only_entry_rejects_excluded_services() -> None: error: Final = callback_config_error("langfuse_otel", {"excluded_services": "redis"}, "failure") assert error is not None and "excluded_services" in error and "success_and_failure" in error + + +_REFERENCE: Final = {"callback_name": "arize", "callback_type": "success", "litellm_credential_name": "space-a"} +_INLINE: Final = {"callback_name": "langfuse", "callback_vars": {"langfuse_public_key": "pk"}} +_SIBLING: Final = {"callback_name": "arize", "callback_vars": {"arize_endpoint": "http://attacker.example/v1/traces"}} + + +@pytest.mark.parametrize( + ("requested", "stored_logging", "is_proxy_admin", "expected_status"), + [ + pytest.param({"logging": [_REFERENCE]}, None, False, 403, id="non-admin-adds-reference"), + pytest.param({"logging": [_INLINE, _REFERENCE]}, None, False, 403, id="non-admin-adds-reference-second"), + pytest.param({"team": "x"}, [_REFERENCE], False, 403, id="non-admin-removes-reference"), + pytest.param({"logging": [_REFERENCE, _SIBLING]}, [_REFERENCE], False, 403, id="non-admin-adds-sibling"), + pytest.param({"logging": [_REFERENCE]}, [_REFERENCE], False, None, id="non-admin-resends-same"), + pytest.param(None, [_REFERENCE], False, None, id="non-admin-leaves-metadata-alone"), + pytest.param({"logging": [_INLINE]}, [], False, None, id="non-admin-without-references"), + pytest.param({"logging": [_REFERENCE]}, [], True, None, id="admin-adds-defined-reference"), + pytest.param({"team": "x"}, [_REFERENCE], True, None, id="admin-removes-reference"), + ], +) +def test_credential_reference_rejection_lets_only_proxy_admins_change_referencing_logging( + requested: dict[str, object] | None, + stored_logging: list[object] | None, + is_proxy_admin: bool, + expected_status: int | None, +) -> None: + rejection: Final = credential_reference_rejection( + requested, stored_logging, is_proxy_admin=is_proxy_admin, is_defined=lambda name: name == "space-a" + ) + assert (None if rejection is None else rejection.status_code) == expected_status + + +def test_credential_reference_rejection_names_undefined_credentials_only_to_admins() -> None: + requested: Final = {"logging": [_INLINE, {**_REFERENCE, "litellm_credential_name": "missing"}]} + + admin: Final = credential_reference_rejection(requested, None, is_proxy_admin=True, is_defined=lambda _: False) + non_admin: Final = credential_reference_rejection(requested, None, is_proxy_admin=False, is_defined=lambda _: False) + + assert admin is not None and admin.status_code == 400 and "missing" in admin.message + assert non_admin is not None and non_admin.status_code == 403 and "missing" not in non_admin.message + + +@pytest.mark.parametrize("is_proxy_admin", [True, False], ids=["admin", "non-admin"]) +def test_credential_reference_rejection_lets_a_stored_reference_to_a_deleted_credential_stay( + is_proxy_admin: bool, +) -> None: + stale: Final = {**_REFERENCE, "litellm_credential_name": "deleted"} + + kept: Final = credential_reference_rejection( + {"logging": [stale], "note": "unrelated edit"}, + [stale], + is_proxy_admin=is_proxy_admin, + is_defined=lambda _: False, + ) + added: Final = credential_reference_rejection( + {"logging": [stale, {**_REFERENCE, "litellm_credential_name": "also-missing"}]}, + [stale], + is_proxy_admin=True, + is_defined=lambda _: False, + ) + + assert kept is None + assert added is not None and added.status_code == 400 + assert "also-missing" in added.message and "deleted" not in added.message + + +@pytest.mark.parametrize("name", [5, ["space-a"], None], ids=["int", "list", "null"]) +def test_a_non_string_credential_name_is_not_a_reference(name: object) -> None: + entry: Final = {**_REFERENCE, "litellm_credential_name": name} + + assert ( + credential_reference_rejection({"logging": [entry]}, None, is_proxy_admin=False, is_defined=lambda _: False) + is None + ) diff --git a/tests/unit/proxy/common_utils/test_callback_utils.py b/tests/unit/proxy/common_utils/test_callback_utils.py index c8d909ac2eb..c7739b9cc2f 100644 --- a/tests/unit/proxy/common_utils/test_callback_utils.py +++ b/tests/unit/proxy/common_utils/test_callback_utils.py @@ -1,5 +1,6 @@ import copy import json +import logging import sys from types import ModuleType, SimpleNamespace from typing import Final @@ -21,6 +22,7 @@ from litellm.proxy.common_utils.callback_utils import ( get_logging_caching_headers, get_remaining_tokens_and_requests_from_request_data, initialize_callbacks_on_proxy, + key_logging_entries, normalize_callback_names, process_callback, sanitize_openai_provider_metadata, @@ -28,6 +30,7 @@ from litellm.proxy.common_utils.callback_utils import ( ) from litellm.proxy.utils import ProxyLogging from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import CredentialItem def test_get_remaining_tokens_and_requests_from_request_data(): @@ -775,3 +778,114 @@ def test_initialize_callbacks_on_proxy_accepts_instance_non_list_value(probe_con assert len(litellm.callbacks) == 1 assert isinstance(litellm.callbacks[0], CustomLogger) + + +def _credential(name: str, values: dict[str, object]) -> CredentialItem: + return CredentialItem(credential_name=name, credential_values=values, credential_info={}) + + +def test_key_logging_entries_merges_each_referenced_credential_under_its_entry(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + litellm, + "credential_list", + [ + _credential("space-a", {"arize_api_key": "key-a", "arize_space_id": "space-a", "api_base": "x"}), + _credential("space-b", {"arize_api_key": "key-b", "arize_space_id": "space-b"}), + ], + ) + entries: Final = [ + {"callback_name": "langfuse", "callback_vars": {"langfuse_public_key": "pk"}}, + {"callback_name": "arize", "litellm_credential_name": "space-a", "callback_vars": {"arize_space_id": "own"}}, + {"callback_name": "arize", "litellm_credential_name": "space-b", "callback_vars": {}}, + ] + + resolved: Final = key_logging_entries({"logging": entries}) + + assert resolved is not None + assert [entry["callback_vars"] for entry in resolved] == [ + {"langfuse_public_key": "pk"}, + {"arize_api_key": "key-a", "arize_space_id": "own"}, + {"arize_api_key": "key-b", "arize_space_id": "space-b"}, + ] + + +@pytest.mark.parametrize( + "credential_values", + [{"api_key": "sk-model-key"}, {"arize_api_key": None}], + ids=["model-credential", "null-values"], +) +def test_key_logging_entries_drops_an_entry_whose_credential_supplies_no_logging_values( + monkeypatch: pytest.MonkeyPatch, credential_values: dict[str, object] +) -> None: + monkeypatch.setattr(litellm, "credential_list", [_credential("model-only", credential_values)]) + kept: Final = {"callback_name": "langfuse", "callback_vars": {"langfuse_public_key": "pk"}} + entries: Final = [ + kept, + {"callback_name": "arize", "litellm_credential_name": "model-only", "callback_vars": {}}, + ] + + assert key_logging_entries({"logging": entries}) == (kept,) + + +def test_a_dropped_credential_reference_is_logged_as_an_error_so_key_health_reports_it( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + monkeypatch.setattr(litellm, "credential_list", []) + entry: Final = {"callback_name": "arize", "litellm_credential_name": "missing", "callback_vars": {}} + + with caplog.at_level(logging.WARNING): + key_logging_entries({"logging": [entry]}) + + assert [record.levelno for record in caplog.records if "missing" in record.getMessage()] == [logging.ERROR] + + +def test_key_logging_entries_drops_an_entry_whose_credential_is_not_defined(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "credential_list", [_credential("space-a", {"arize_api_key": "key-a"})]) + kept: Final = {"callback_name": "langfuse", "callback_vars": {"langfuse_public_key": "pk"}} + entries: Final = [ + kept, + {"callback_name": "arize", "litellm_credential_name": "missing", "callback_vars": {}}, + ] + + assert key_logging_entries({"logging": entries}) == (kept,) + + +@pytest.mark.parametrize( + "logging_value", + [ + pytest.param("abc", id="string"), + pytest.param({"callback_name": "arize"}, id="object"), + pytest.param(5, id="number"), + pytest.param(["x"], id="list-with-string"), + pytest.param([{"callback_name": "langfuse", "callback_vars": {}}, None], id="list-with-null"), + ], +) +def test_key_logging_entries_fails_on_non_object_logging_like_the_request_path_always_has( + logging_value: object, +) -> None: + with pytest.raises(TypeError): + key_logging_entries({"logging": logging_value}) + + +@pytest.mark.parametrize( + ("logging_value", "expected"), + [ + pytest.param("", (), id="empty-string"), + pytest.param({}, (), id="empty-object"), + pytest.param([], (), id="empty-list"), + pytest.param([{"bad": 1}], ({"bad": 1},), id="object-with-unknown-fields"), + ], +) +def test_key_logging_entries_passes_empty_and_object_entries_through( + logging_value: object, expected: tuple[dict[str, object], ...] +) -> None: + assert key_logging_entries({"logging": logging_value}) == expected + + +def test_key_logging_entries_leaves_an_entry_with_a_non_string_credential_name_as_written( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "credential_list", [_credential("5", {"arize_api_key": "k"})]) + entry: Final = {"callback_name": "arize", "litellm_credential_name": 5, "callback_vars": {}} + + assert key_logging_entries({"logging": [entry]}) == (entry,) diff --git a/tests/unit/proxy/common_utils/test_reset_budget_job.py b/tests/unit/proxy/common_utils/test_reset_budget_job.py index 4a4490afc33..aaa686f2a98 100644 --- a/tests/unit/proxy/common_utils/test_reset_budget_job.py +++ b/tests/unit/proxy/common_utils/test_reset_budget_job.py @@ -1438,7 +1438,47 @@ def test_reset_budget_for_keys_invalidates_redis_counter(reset_budget_job, mock_ asyncio.run(reset_budget_job.reset_budget_for_litellm_keys()) - assert "spend:key:sk-abc" in counter_cache.deleted() + assert {"spend:key:sk-abc", "spend:committed:key:sk-abc"} <= counter_cache.deleted() + + +def _expired_key_row() -> object: + return type( + "Key", + (), + { + "spend": 100.0, + "budget_duration": "30d", + "budget_reset_at": datetime.now(timezone.utc), + "id": "key-1", + "token": "sk-abc", + }, + ) + + +def test_reset_budget_for_keys_without_redis_clears_the_counter_and_its_committed_twin( + reset_budget_job, mock_prisma_client, monkeypatch +): + counter_cache: Final = _make_counter_invalidation_job(monkeypatch) + counter_cache.redis = None + counter_cache.memory.update({"spend:key:sk-abc": 100.0, "spend:committed:key:sk-abc": 100.0}) + mock_prisma_client.data["key"] = [_expired_key_row()] + + asyncio.run(reset_budget_job.reset_budget_for_litellm_keys()) + + assert counter_cache.memory == {} + + +def test_a_redis_failure_on_the_key_reset_still_drops_its_committed_twin_locally( + reset_budget_job, mock_prisma_client, monkeypatch +): + counter_cache: Final = _make_counter_invalidation_job(monkeypatch) + counter_cache.failing = frozenset({"delete"}) + counter_cache.memory.update({"spend:key:sk-abc": 100.0, "spend:committed:key:sk-abc": 100.0}) + mock_prisma_client.data["key"] = [_expired_key_row()] + + asyncio.run(reset_budget_job.reset_budget_for_litellm_keys()) + + assert counter_cache.memory == {} def test_reset_budget_for_users_invalidates_redis_counter(reset_budget_job, mock_prisma_client, monkeypatch): @@ -1790,7 +1830,9 @@ def test_enduser_invalidation_is_paged_and_batched(reset_budget_job, mock_prisma assert counter_cache.user_api_key_cache.async_delete_cache_keys.await_count == 6 invalidated: Final = counter_cache.deleted() - assert invalidated == {f"spend:end_user:cust-{i:06d}" for i in range(population)} + assert invalidated == {f"spend:end_user:cust-{i:06d}" for i in range(population)} | { + f"spend:committed:end_user:cust-{i:06d}" for i in range(population) + } evicted: Final = { key for call in counter_cache.user_api_key_cache.async_delete_cache_keys.await_args_list for key in call.args[0] } diff --git a/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py index e9e0a585949..d2161fddeb0 100644 --- a/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py @@ -416,38 +416,6 @@ async def test_deferred_anthropic_messages_bridged_to_the_responses_api_logs_the # --------------------------------------------------------------------------- -@pytest.mark.parametrize("call_type", ["ocr", "aocr", "completion", "acompletion", "embedding", "responses"]) -@pytest.mark.parametrize("exception_raised", [False, True]) -def test_native_pending_logging_is_released_only_for_ocr(call_type: str, exception_raised: bool) -> None: - pending: Final = MagicMock() - enqueue: Final = MagicMock() - logger: Final = MagicMock( - call_type=call_type, - _native_pending_logging=pending, - enqueue_deferred_logging=enqueue, - ) - - ProxyBaseLLMRequestProcessing._flush_deferred_async_logging( - logging_obj=logger, - exception_raised=exception_raised, - ) - ProxyBaseLLMRequestProcessing._flush_deferred_async_logging( - logging_obj=logger, - exception_raised=exception_raised, - ) - - if call_type in ("ocr", "aocr"): - pending.release.assert_called_once_with(not exception_raised) - assert logger._native_pending_logging is None - else: - pending.release.assert_not_called() - assert logger._native_pending_logging is pending - if exception_raised: - enqueue.assert_not_called() - else: - enqueue.assert_called_once_with() - - def test_flush_deferred_async_logging_fires_on_success(): """ Happy path: with no exception, the production flush helper invokes the diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index d06f95111aa..43f529faa6b 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -13,7 +13,7 @@ import pytest from litellm.redis.stores.connection import connect from litellm.redis.stores.spend_counters import SpendCounterStore -from fastapi import HTTPException +from fastapi import FastAPI, HTTPException import litellm from litellm._logging import verbose_proxy_logger @@ -35,6 +35,7 @@ from litellm.proxy.hooks.proxy_track_cost_callback import ( run_spend_event, ) from litellm.proxy.route_llm_request import ProxyModelNotFoundError +from litellm.proxy.pass_through_endpoints.pass_through_endpoints import InitPassThroughEndpointHelpers from litellm.proxy.spend_tracking.spend_event import SpendEventDecodeError, build_spend_event, decode_spend_event from litellm.proxy.spend_tracking.spend_event_producer import SpendEventProducer, UnixAddress from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload @@ -44,6 +45,7 @@ from litellm.proxy.utils import ( BilledFailedAttempt, ProxyUpdateSpend, ) +from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys from litellm.types.router import Deployment, LiteLLM_Params from litellm.types.utils import CallTypes, LiteLLMBatch, ModelResponse, PromptTokensDetailsWrapper, Usage from litellm.types.videos.main import VideoObject @@ -906,6 +908,74 @@ async def test_async_post_call_failure_hook_non_llm_route(): mock_update_database.assert_not_called() +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("route", "call_type", "written"), + [ + ("/azure-search", CallTypes.pass_through.value, True), + ("/unregistered-search", CallTypes.acompletion.value, False), + ("/key/generate", CallTypes.pass_through.value, False), + ("/team/new", CallTypes.pass_through.value, False), + ("/unregistered-search", CallTypes.pass_through.value, False), + ("/azure-search/", CallTypes.pass_through.value, True), + ("/configured-search/indexes/docs", CallTypes.pass_through.value, True), + (None, CallTypes.acompletion.value, True), + ], +) +async def test_async_post_call_failure_hook_tracks_config_defined_pass_through_route( + route: str | None, call_type: str, written: bool +) -> None: + api_base: Final = "http://upstream.example/indexes/x/docs/search" + request_data: Final = { + "model": "azure-search", + "call_type": call_type, + "litellm_params": {"api_base": api_base, "metadata": {}}, + } + user_api_key_dict: Final = UserAPIKeyAuth( + api_key="test_api_key", + request_route=route, + ) + + spend_writer: Final = create_autospec(DBSpendUpdateWriter, instance=True) + with patch.dict( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes", clear=True + ): + InitPassThroughEndpointHelpers.add_exact_path_route( + app=FastAPI(), + path="/azure-search", + target=api_base, + custom_headers={}, + forward_headers=False, + merge_query_params=False, + dependencies=[], + cost_per_request=None, + endpoint_id="configured-exact", + ) + InitPassThroughEndpointHelpers.add_subpath_route( + app=FastAPI(), + path="/configured-search", + target=api_base, + custom_headers={}, + forward_headers=False, + merge_query_params=False, + dependencies=[], + cost_per_request=None, + endpoint_id="configured-subpath", + ) + await ProxyDBLogger(spend_writer=lambda: spend_writer).async_post_call_failure_hook( + request_data=request_data, + original_exception=Exception("Provider error"), + user_api_key_dict=user_api_key_dict, + ) + + assert spend_writer.update_database.await_count == (1 if written else 0) + if not written: + return + call_args: Final = spend_writer.update_database.call_args + assert call_args.kwargs["kwargs"]["litellm_params"]["api_base"] == api_base + assert call_args.kwargs["kwargs"]["litellm_params"]["metadata"]["status"] == "failure" + + @pytest.mark.asyncio @pytest.mark.parametrize( ("team_id", "route", "written"), @@ -1129,6 +1199,95 @@ def test_get_budget_reservation_from_metadata_handles_dict_auth_object(): ) +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("metadata", "expected_scope"), + [ + ( + { + "user_api_key_auth": UserAPIKeyAuth( + committed_budget_spend=CommittedBudgetSpend( + counter_keys=EntitySpendCounterKeys(key="spend:key:test_api_key") + ) + ) + }, + EntitySpendCounterKeys(key="spend:key:test_api_key"), + ), + ({}, EntitySpendCounterKeys()), + ], +) +async def test_charge_moves_only_the_committed_counters_auth_named_as_budgeted(metadata, expected_scope): + proxy_logging_obj = MagicMock() + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock() + increment_spend_counters = AsyncMock() + + await _update_database_and_spend_counters( + proxy_logging_obj=proxy_logging_obj, + increment_spend_counters=increment_spend_counters, + user_api_key="test_api_key", + user_id="test_user_id", + end_user_id=None, + team_id="test_team_id", + org_id="test_org_id", + kwargs={"litellm_params": {"metadata": metadata}}, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + response_cost=0.2, + budget_reservation=None, + ) + + assert increment_spend_counters.await_args.kwargs["committed_spend_scope"] == expected_scope + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("tracked", "expected_key", "expected_org"), + [(True, pytest.approx(1.2), pytest.approx(2.2)), (False, 1.0, 2.0)], +) +async def test_charge_with_no_auth_on_its_row_moves_the_committed_counters_that_exist( + tracked, expected_key, expected_org +): + """A background response is billed from its creator's stored attribution, after the request's auth is gone.""" + from litellm.proxy.proxy_server import increment_spend_counters + + proxy_logging_obj = MagicMock() + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock() + spend_counter_store: Final = SpendCounterStore(connect(None).spend_counters()) + spend_counter_store.set_local("spend:committed:key:test_api_key", 1.0) + spend_counter_store.set_local("spend:committed:org:test_org_id", 2.0) + + with ( + patch("litellm.proxy.hooks.proxy_track_cost_callback.committed_spend_tracked", return_value=tracked), + patch("litellm.proxy.proxy_server.spend_counter_store", spend_counter_store), + patch( + "litellm.proxy.spend_tracking.committed_spend.SpendCounterReseed.from_db", + AsyncMock(return_value=50.0), + ), + ): + await _update_database_and_spend_counters( + proxy_logging_obj=proxy_logging_obj, + increment_spend_counters=increment_spend_counters, + user_api_key="test_api_key", + user_id="test_user_id", + end_user_id=None, + team_id="test_team_id", + org_id="test_org_id", + kwargs={"litellm_params": {"metadata": {"user_api_key_hash": "test_api_key"}}}, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + response_cost=0.2, + budget_reservation=None, + ) + + cached: Final = spend_counter_store.get_local + assert cached("spend:committed:key:test_api_key") == expected_key + assert cached("spend:committed:org:test_org_id") == expected_org + assert cached("spend:committed:team:test_team_id") is None + assert cached("spend:committed:user:test_user_id") is None + + @pytest.mark.asyncio async def test_update_database_and_spend_counters_releases_reservation_when_db_update_fails(): proxy_logging_obj = MagicMock() @@ -1257,6 +1416,7 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda request_started_at=start_time, model_access_groups=("premium",), project_id=None, + committed_spend_scope=EntitySpendCounterKeys(), ) @@ -2776,7 +2936,9 @@ _LIVE_COUNTERS: Final = { } -async def _fail_request(logger: ProxyDBLogger, failure: Mapping[str, object]) -> None: +async def _fail_request( + logger: ProxyDBLogger, failure: Mapping[str, object], committed: CommittedBudgetSpend | None = None +) -> None: await logger.async_post_call_failure_hook( request_data={ "model": "gpt-4.1", @@ -2793,6 +2955,7 @@ async def _fail_request(logger: ProxyDBLogger, failure: Mapping[str, object]) -> end_user_id="end-user-1", org_id="org-1", project_id="project-1", + committed_budget_spend=committed, ), ) @@ -2831,6 +2994,43 @@ async def test_failure_spend_rows_raise_the_live_spend_counters_by_their_cost( ) +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("failure", "charged"), + [ + ({"combined_usage_object": Usage(prompt_tokens=30), "response_cost": 0.125}, 0.125), + (_failure_with_billed_attempts(), 1.25), + ({"standard_logging_object": {"cost_finalized": True, "response_cost": 0.125}}, 0.125), + ], + ids=["recovered_stream_cost", "billed_retries", "completed_cost"], +) +async def test_a_failed_charge_moves_only_the_committed_counters_auth_named_as_budgeted( + live_spend_counters: SpendCounterStore, failure: Mapping[str, object], charged: float +): + """The remaining budget gauges read the committed counters, so a charge missing from them shows as budget the + request never left behind, until the next reseed from the database.""" + budgeted: Final = EntitySpendCounterKeys.of( + token="hashed-key", team_id="team-1", user_id=None, org_id=None, end_user_id=None + ) + starting: Final = { + "spend:committed:key:hashed-key": 0.2, + "spend:committed:team:team-1": 0.3, + "spend:committed:user:user-1": 0.4, + } + for key, spend in starting.items(): + live_spend_counters.set_local(key, spend) + + await _fail_request(ProxyDBLogger(), failure, CommittedBudgetSpend(counter_keys=budgeted)) + + assert {key: live_spend_counters.get_local(key) for key in starting} == pytest.approx( + { + "spend:committed:key:hashed-key": 0.2 + charged, + "spend:committed:team:team-1": 0.3 + charged, + "spend:committed:user:user-1": 0.4, + } + ) + + @pytest.mark.asyncio async def test_a_completed_cost_failure_raises_the_live_spend_counters_once(live_spend_counters: SpendCounterStore): """A failure whose cost was already finalized settles through the success path's write, which raises the diff --git a/tests/unit/proxy/management_endpoints/test_customer_endpoints.py b/tests/unit/proxy/management_endpoints/test_customer_endpoints.py index e9f10e8c0db..c42690050da 100644 --- a/tests/unit/proxy/management_endpoints/test_customer_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_customer_endpoints.py @@ -249,6 +249,29 @@ def test_delete_customer_not_found(mock_prisma_client, mock_user_api_key_auth): assert response_json["error"]["code"] == "404" +def test_delete_customer_drops_its_committed_spend_so_a_recreated_customer_starts_from_its_own_row( + monkeypatch, mock_prisma_client, mock_user_api_key_auth +): + from types import SimpleNamespace + + from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters + + spend: Final = FakeSpendCounters( + redis=None, + memory={"spend:committed:end_user:cust-gone": 7.0, "spend:committed:end_user:cust-stays": 3.0}, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_store", spend.store()) + mock_prisma_client.db.litellm_endusertable.find_many = AsyncMock(return_value=[SimpleNamespace(user_id="cust-gone")]) + mock_prisma_client.db.litellm_endusertable.delete_many = AsyncMock(return_value=1) + + response: Final = client.post( + "/customer/delete", json={"user_ids": ["cust-gone"]}, headers={"Authorization": "Bearer test-key"} + ) + + assert response.status_code == 200 + assert spend.memory == {"spend:committed:end_user:cust-stays": 3.0} + + def test_error_schema_consistency(mock_prisma_client, mock_user_api_key_auth): """ Test that all customer endpoints return the same error schema format. diff --git a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py index 120cf8a2b84..b7be80da31a 100644 --- a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py @@ -4222,7 +4222,10 @@ async def test_admin_user_update_spend_invalidates_counter(mocker): admin_caller = UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) await _update_single_user_helper(user_request=user_request, user_api_key_dict=admin_caller) - mock_invalidate.assert_awaited_once_with(counter_key="spend:user:target-user") + assert [call.kwargs["counter_key"] for call in mock_invalidate.await_args_list] == [ + "spend:user:target-user", + "spend:committed:user:target-user", + ] @pytest.mark.asyncio @@ -5223,6 +5226,39 @@ async def test_delete_user_evicts_cached_user_rows(mocker: MockerFixture) -> Non broadcast.assert_awaited_once_with(cache_key=deleted.user_id) +@pytest.mark.asyncio +async def test_delete_user_drops_the_committed_spend_of_the_users_it_deleted(mocker: MockerFixture) -> None: + from litellm.proxy._types import DeleteUserRequest, LiteLLM_UserTable + from litellm.proxy.management_endpoints.internal_user_endpoints import delete_user + from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters + + spend: Final = FakeSpendCounters( + redis=None, memory={"spend:committed:user:user-gone": 7.0, "spend:committed:user:user-stays": 3.0} + ) + mocker.patch("litellm.proxy.proxy_server.spend_counter_store", spend.store()) # test-quality-ok: substitute the spend counter dependency + deleted: Final = LiteLLM_UserTable(user_id="user-gone", user_email="gone@example.test", teams=[]) + prisma_client: Final = mocker.MagicMock() + prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(return_value=deleted) + prisma_client.db.litellm_teamtable.find_many = mocker.AsyncMock(return_value=[]) + prisma_client.db.litellm_jwtkeymapping.find_many = mocker.AsyncMock(return_value=[]) + prisma_client.db.litellm_verificationtoken.find_many = mocker.AsyncMock(return_value=[]) + prisma_client.db.litellm_verificationtoken.delete_many = mocker.AsyncMock(return_value=0) + prisma_client.db.litellm_invitationlink.delete_many = mocker.AsyncMock(return_value=0) + prisma_client.db.litellm_organizationmembership.delete_many = mocker.AsyncMock(return_value=0) + prisma_client.db.litellm_teammembership.delete_many = mocker.AsyncMock(return_value=0) + prisma_client.db.litellm_usertable.delete_many = mocker.AsyncMock(return_value=1) + _serve_locked_user_delete(prisma_client, deleted) + mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency + mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj", None) # test-quality-ok: delete_user reads it off proxy_server at call time + + await delete_user( + data=DeleteUserRequest(user_ids=[deleted.user_id]), + user_api_key_dict=UserAPIKeyAuth(user_id="proxy-admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert spend.memory == {"spend:committed:user:user-stays": 3.0} + + _DB_OUTAGE_503_BODY: Final = { "error": { "message": "Service Unavailable, the authentication database is temporarily unreachable. Please retry shortly.", diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index f83d53ac2e5..2e18a0c1099 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -51,6 +51,7 @@ from litellm.proxy.auth.auth_checks import ( ) from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import AuthCacheInvalidationSubscriber +from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, budget_windows_absent_cache_key, @@ -101,6 +102,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( persist_deleted_verification_tokens, prepare_key_update_data, prepare_metadata_fields, + raise_on_logging_credential_change, reset_key_spend_fn, update_key_fn, validate_key_list_check, @@ -112,6 +114,7 @@ from litellm.types.proxy.management_endpoints.key_management_endpoints import ( BulkUpdateKeyResponse, CustomKeyPolicyRequest, ) +from litellm.types.utils import CredentialItem from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters client = TestClient(app) @@ -8693,6 +8696,7 @@ async def test_reset_key_spend_success(monkeypatch): mock_spend_counter_cache.in_memory_cache.set_cache.assert_any_call( key=f"spend_db_floor:spend:key:{hashed_key}", value=50.0, ttl=5 ) + assert spend.deleted() == {f"spend:committed:key:{hashed_key}"} @pytest.mark.asyncio @@ -9013,7 +9017,10 @@ async def test_update_key_spend_updates_counter(monkeypatch): mock_delete_cache.assert_awaited_once() assert spend.memory == {f"spend:key:{hashed_key}": 0.0} - assert spend.commands == [f"SET spend:key:{hashed_key} 0.0 60"] + assert spend.commands == [ + f"SET spend:key:{hashed_key} 0.0 60", + f"DEL spend:committed:key:{hashed_key}", + ] @pytest.mark.asyncio @@ -19793,6 +19800,150 @@ async def test_bulk_key_update_batch_enqueued_token_limit_rejected_for_non_admin assert "Only proxy admins can set" in str(exc.value.detail) + +_INDEX_GRANT = [{"index_name": "docs", "index_permissions": ["read", "write"]}] + + +@pytest.mark.parametrize( + "label, top_level, sent, stored, allowed", + [ + ("granted top-level on a key with none stored", _INDEX_GRANT, _UNTOUCHED, None, False), + ("granted in metadata on a key with none stored", None, _INDEX_GRANT, None, False), + ("widened from read to read and write", None, _INDEX_GRANT, [{"index_name": "docs", "index_permissions": ["read"]}], False), + ("cleared by replacing the blob", None, None, _INDEX_GRANT, False), + ("resent unchanged in metadata", None, _INDEX_GRANT, _INDEX_GRANT, True), + ("resent unchanged top-level", _INDEX_GRANT, _UNTOUCHED, _INDEX_GRANT, True), + ("left untouched", None, _UNTOUCHED, _INDEX_GRANT, True), + ("cleared top-level with an empty list", [], _UNTOUCHED, _INDEX_GRANT, False), + ("empty top-level list on a key with none stored", [], _UNTOUCHED, None, True), + ], +) +def test_vector_store_index_grants_admin_gate_matrix( + label: str, + top_level: list[dict[str, object]] | None, + sent: object, + stored: list[dict[str, object]] | None, + allowed: bool, +) -> None: + from litellm.proxy.auth.auth_utils import ( + enforce_vector_store_index_grants_are_admin_only, + ) + + request_body = {} if sent is _UNTOUCHED else {"metadata": {} if sent is None else {"allowed_vector_store_indexes": sent}} + if top_level is not None: + request_body["allowed_vector_store_indexes"] = top_level + data = UpdateKeyRequest(key="sk-1", **request_body) + existing_metadata = None if stored is None else {"allowed_vector_store_indexes": stored} + + def _call(caller: UserAPIKeyAuth) -> None: + enforce_vector_store_index_grants_are_admin_only( + data=data, + top_level_grants=data.allowed_vector_store_indexes, + existing_metadata=existing_metadata, + user_api_key_dict=caller, + entity="key", + ) + + non_admin = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-non-admin", user_id="alice") + if allowed: + _call(non_admin) + else: + with pytest.raises(HTTPException) as exc: + _call(non_admin) + assert exc.value.status_code == 403 + assert "Only proxy admins can set allowed_vector_store_indexes" in str(exc.value.detail) + + _call(UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin")) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "request_fields", + [{"allowed_vector_store_indexes": _INDEX_GRANT}, {"metadata": {"allowed_vector_store_indexes": _INDEX_GRANT}}], + ids=["top-level", "metadata"], +) +async def test_generate_key_vector_store_index_grants_rejected_for_non_admin(request_fields: dict[str, object]) -> None: + """An index grant on a key opens document reads and writes on that index, so a + non-admin minting their own key cannot pick its grants.""" + with patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()): + with pytest.raises(HTTPException) as exc: + await _common_key_generation_helper( + data=GenerateKeyRequest(**request_fields), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-alice", + user_id="alice", + ), + litellm_changed_by=None, + team_table=None, + ) + assert int(getattr(exc.value, "status_code", 0)) == 403 + assert "Only proxy admins can set allowed_vector_store_indexes" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_update_key_vector_store_index_grants_unchanged_allows_non_admin_edit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A key owner renaming a key that carries admin-set index grants resends them + with the rest of the form, which must not count as setting them.""" + from litellm.proxy.management_endpoints.key_management_endpoints import ( + update_key_fn, + ) + + token = "b3b2c3d4e5f6789012345678901234567890123456789012345678901234abcd" + _wire_update_key_fn(monkeypatch, _estimate_key_row(token, {"allowed_vector_store_indexes": _INDEX_GRANT})) + + mock_request = MagicMock() + mock_request.query_params = {} + + result = await update_key_fn( + request=mock_request, + data=UpdateKeyRequest(key=token, key_alias="my-alias", metadata={"allowed_vector_store_indexes": _INDEX_GRANT}), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-internal", + user_id="internal_user", + ), + litellm_changed_by=None, + ) + + assert result is not None + + +@pytest.mark.asyncio +async def test_update_key_vector_store_index_grants_widened_rejected_for_non_admin( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy.management_endpoints.key_management_endpoints import ( + update_key_fn, + ) + + token = "c3b2c3d4e5f6789012345678901234567890123456789012345678901234abcd" + _wire_update_key_fn( + monkeypatch, + _estimate_key_row(token, {"allowed_vector_store_indexes": [{"index_name": "docs", "index_permissions": ["read"]}]}), + ) + + mock_request = MagicMock() + mock_request.query_params = {} + + with pytest.raises(ProxyException) as exc: + await update_key_fn( + request=mock_request, + data=UpdateKeyRequest(key=token, allowed_vector_store_indexes=_INDEX_GRANT), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-internal", + user_id="internal_user", + ), + litellm_changed_by=None, + ) + + assert str(exc.value.code) == "403" + assert "Only proxy admins can set allowed_vector_store_indexes" in str(exc.value.message) + + @pytest.mark.asyncio async def test_execute_virtual_key_regeneration_stamps_settings_updated_at(): """Regenerate rewrites the key's config, so it must move settings_updated_at.""" @@ -24259,3 +24410,29 @@ async def test_key_aliases_answers_503_when_the_database_is_unreachable(monkeypa ) assert (exc_info.value.code, exc_info.value.type) == ("503", "no_db_connection") + + +def test_non_admin_resaving_referencing_key_metadata_read_back_encrypted_is_not_a_logging_change(monkeypatch): + monkeypatch.setenv("LITELLM_SALT_KEY", "test-salt-32-bytes-aaaaaaaaaaaaaa") + monkeypatch.setattr( + litellm, + "credential_list", + [CredentialItem(credential_name="space-a", credential_values={"arize_api_key": "k"}, credential_info={})], + ) + stored = encrypt_callback_vars( + { + "logging": [ + {"callback_name": "arize", "litellm_credential_name": "space-a", "callback_vars": {}}, + {"callback_name": "langfuse", "callback_vars": {"langfuse_secret_key": "sk-lf-secret"}}, + ] + } + ) + assert stored["logging"][1]["callback_vars"]["langfuse_secret_key"] != "sk-lf-secret" + key_row = LiteLLM_VerificationToken(token="hashed", metadata=stored) + internal_user = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER) + + raise_on_logging_credential_change({**stored, "note": "unrelated edit"}, key_row, internal_user) + with pytest.raises(HTTPException) as removed: + raise_on_logging_credential_change({"note": "unrelated edit"}, key_row, internal_user) + + assert removed.value.status_code == 403 diff --git a/tests/unit/proxy/management_endpoints/test_organization_endpoints.py b/tests/unit/proxy/management_endpoints/test_organization_endpoints.py index 8d2c6c0054c..70cfaf4140d 100644 --- a/tests/unit/proxy/management_endpoints/test_organization_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_organization_endpoints.py @@ -1188,6 +1188,27 @@ async def _run_legacy_update_organization( return mock_prisma_client +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("body", "twin_after"), + [ + ({"organization_id": "org-1", "spend": 0.0}, None), + ({"organization_id": "org-1", "organization_alias": "renamed"}, 7.0), + ], +) +async def test_legacy_update_drops_the_committed_org_spend_only_when_it_writes_the_stored_spend( + monkeypatch, body, twin_after +): + from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters + + spend: Final = FakeSpendCounters(redis=None, memory={"spend:committed:org:org-1": 7.0}) + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_store", spend.store()) + + await _run_legacy_update_organization(monkeypatch, body=body, existing_budget_id="budget-1") + + assert spend.memory.get("spend:committed:org:org-1") == twin_after + + @pytest.mark.asyncio @pytest.mark.parametrize( "body", @@ -1582,6 +1603,49 @@ async def test_delete_organization_evicts_the_cache_of_the_keys_it_deletes(monke assert jwt_table.rows == (kept_row,) +@pytest.mark.asyncio +async def test_delete_organization_drops_the_committed_spend_of_the_org_and_its_teams(monkeypatch): + from litellm.proxy._types import DeleteOrganizationRequest, LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.management_endpoints.organization_endpoints import delete_organization + from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters + + spend: Final = FakeSpendCounters( + redis=None, + memory={ + "spend:committed:org:org-doomed": 7.0, + "spend:committed:team:team-in-org": 5.0, + "spend:committed:org:org-kept": 3.0, + "spend:committed:team:team-elsewhere": 2.0, + }, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_store", spend.store()) + prisma_client: Final = AsyncMock() + prisma_client.db.litellm_teamtable.find_many = AsyncMock( + return_value=[SimpleNamespace(team_id="team-in-org", organization_id="org-doomed")] + ) + prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + prisma_client.db.litellm_jwtkeymapping.find_many = AsyncMock(return_value=[]) + prisma_client.db.litellm_organizationtable.delete = AsyncMock(return_value=MagicMock()) + + @asynccontextmanager + async def transaction(): + yield prisma_client.db + + prisma_client.tx = transaction + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True, raising=False) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", None) + + await delete_organization( + data=DeleteOrganizationRequest(organization_ids=["org-doomed"]), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert spend.memory == {"spend:committed:org:org-kept": 3.0, "spend:committed:team:team-elsewhere": 2.0} + + @pytest.mark.asyncio async def test_delete_organization_is_refused_while_one_of_its_teams_holds_a_ptu_share(monkeypatch): """Deleting an organization deletes its teams, so a team still holding a PTU share stops it the diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index 94a8022fe96..0881c902a9d 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -9966,6 +9966,37 @@ async def test_delete_validated_team_rows_deletes_only_the_keys_its_rostered_mem assert deleted_keys.kwargs["where"] == {"token": {"in": ["hashed-member"]}} +@pytest.mark.asyncio +async def test_delete_validated_team_rows_drops_the_committed_spend_of_the_teams_it_deleted( + monkeypatch, + disable_audit_logging_for_mocked_team, +): + from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters + + spend: Final = FakeSpendCounters( + redis=None, memory={"spend:committed:team:team-doomed": 7.0, "spend:committed:team:team-kept": 3.0} + ) + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_store", spend.store()) + team = LiteLLM_TeamTable(team_id="team-doomed", members_with_roles=[], metadata={}, model_max_budget={}, model_spend={}) + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_teammembership.delete_many = AsyncMock() + mock_access_group_tx_cm = MagicMock() + mock_access_group_tx_cm.__aenter__ = AsyncMock(return_value=AsyncMock()) + mock_access_group_tx_cm.__aexit__ = AsyncMock(return_value=False) + mock_prisma_client.db.tx = MagicMock(return_value=mock_access_group_tx_cm) + _wire_team_delete_tx(mock_prisma_client, locked_teams=(team,), team_keys=()) + monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") + + await delete_validated_team_rows( + teams=(team,), + prisma_client=mock_prisma_client, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="hashed-admin"), + litellm_changed_by=None, + ) + + assert spend.memory == {"spend:committed:team:team-kept": 3.0} + + @pytest.mark.asyncio async def test_delete_validated_teams_writes_the_deleted_audit_row_only_for_a_team_still_there_under_the_lock( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unit/proxy/management_helpers/test_bulk_user_creation.py b/tests/unit/proxy/management_helpers/test_bulk_user_creation.py index 2b966a8ebeb..96c84578111 100644 --- a/tests/unit/proxy/management_helpers/test_bulk_user_creation.py +++ b/tests/unit/proxy/management_helpers/test_bulk_user_creation.py @@ -7,6 +7,7 @@ import pytest from prisma.errors import UniqueViolationError from pydantic import BaseModel, ConfigDict, ValidationError +import litellm from litellm.caching.caching import DualCache from litellm.proxy._types import LiteLLM_TeamTable, LitellmUserRoles, Member, UserAPIKeyAuth from litellm.proxy.list_api.common import ManagementProblem @@ -15,6 +16,7 @@ from litellm.types.proxy.management_endpoints.internal_user_endpoints import ( BulkNewUserItem, BulkNewUserRequest, ) +from litellm.types.utils import CredentialItem ADMIN: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) INTERNAL: Final = UserAPIKeyAuth(user_id="someone", user_role=LitellmUserRoles.INTERNAL_USER) @@ -434,6 +436,39 @@ async def test_only_proxy_admin_sets_batch_limits_on_a_created_key(): assert calls[-1]["metadata"] == limits +@pytest.mark.asyncio +async def test_only_proxy_admin_creates_a_key_whose_logging_references_a_credential( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + litellm, + "credential_list", + [CredentialItem(credential_name="space-a", credential_values={"arize_api_key": "k"}, credential_info={})], + ) + reference: Final = {"logging": [{"callback_name": "arize", "litellm_credential_name": "space-a", "callback_vars": {}}]} + calls: list[dict[str, object]] = [] + + async def generate_key(**kwargs: object) -> dict[str, object]: + calls.append(kwargs) + return {"token": f"sk-{kwargs['user_id']}"} + + rows: Final = [ + {"user_id": "u1", "auto_create_key": True, "metadata": reference}, + {"user_id": "u2", "auto_create_key": True}, + ] + + response = await _run(_FakePrisma(), rows, caller=INTERNAL, generate_key=generate_key) + + assert [r.success for r in response.data] == [False, True] + assert "litellm_credential_name" in (response.data[0].error or "") + assert [call["user_id"] for call in calls] == ["u2"] + + admin_response = await _run(_FakePrisma(), rows[:1], caller=ADMIN, generate_key=generate_key) + + assert [r.key for r in admin_response.data] == ["sk-u1"] + assert calls[-1]["metadata"] == reference + + @pytest.mark.asyncio async def test_license_is_checked_once_against_the_whole_batch(): prisma = _FakePrisma() diff --git a/tests/unit/proxy/management_helpers/test_bulk_user_deletion.py b/tests/unit/proxy/management_helpers/test_bulk_user_deletion.py index e212a433393..880e5746150 100644 --- a/tests/unit/proxy/management_helpers/test_bulk_user_deletion.py +++ b/tests/unit/proxy/management_helpers/test_bulk_user_deletion.py @@ -2,11 +2,12 @@ import asyncio import copy import itertools import json -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence from contextlib import asynccontextmanager from typing import Final, cast import pytest +from pytest_mock import MockerFixture from pydantic import BaseModel, ConfigDict, ValidationError import litellm @@ -22,6 +23,7 @@ from litellm.proxy.management_helpers.bulk_user_deletion import ( ) from litellm.types.proxy.management_endpoints.internal_user_endpoints import BulkDeleteUserRequest from litellm.types.proxy.management_endpoints.team_endpoints import BulkTeamMemberDeleteRequest, TeamMemberRef +from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters ADMIN: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin") INTERNAL: Final = UserAPIKeyAuth(user_id="someone", user_role=LitellmUserRoles.INTERNAL_USER) @@ -697,6 +699,34 @@ async def test_bulk_delete_evicts_jwt_key_mappings_of_the_deleted_users_keys(): assert cache.get_cache(key=kept) is not None +async def _delete_through_scim_or_user_delete(prisma: _FakePrisma, user_ids: Sequence[str]) -> object: + return await delete_users_from_current_teams( + prisma_client=prisma, # pyright: ignore[reportArgumentType] # fake stands in for PrismaClient + user_ids=user_ids, + user_api_key_dict=ADMIN, + user_api_key_cache=UserApiKeyCache(), + proxy_logging_obj=None, + litellm_changed_by=None, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("delete", [_delete, _delete_through_scim_or_user_delete]) +async def test_user_deletes_drop_the_committed_spend_of_the_users_they_deleted( + delete: Callable[[_FakePrisma, Sequence[str]], Awaitable[object]], mocker: MockerFixture +) -> None: + spend: Final = FakeSpendCounters( + redis=None, memory={"spend:committed:user:u1": 7.0, "spend:committed:user:keep": 3.0} + ) + mocker.patch("litellm.proxy.proxy_server.spend_counter_store", spend.store()) # test-quality-ok: substitute the spend counter dependency + prisma: Final = _FakePrisma(users=[_user("u1"), _user("keep")]) + + _ = await delete(prisma, ["u1"]) + + assert sorted(prisma.db.litellm_usertable.rows) == ["keep"] + assert spend.memory == {"spend:committed:user:keep": 3.0} + + @pytest.mark.asyncio async def test_bulk_delete_rejects_non_admin_callers_before_touching_the_db(): prisma = _FakePrisma(users=[_user("u1")]) diff --git a/tests/unit/proxy/mcp_server/test_idp_token_exchange.py b/tests/unit/proxy/mcp_server/test_idp_token_exchange.py index 036ee4a332b..75b0a5cd308 100644 --- a/tests/unit/proxy/mcp_server/test_idp_token_exchange.py +++ b/tests/unit/proxy/mcp_server/test_idp_token_exchange.py @@ -1,20 +1,34 @@ +import json import logging from typing import Final +import jwt import pytest -from fastapi import HTTPException +from cryptography.hazmat.primitives.asymmetric import rsa +from fastapi import HTTPException, Request from prisma.engine.errors import BinaryNotFoundError from prisma.errors import DataError from litellm.caching.caching import DualCache -from litellm.proxy._types import JWTIssuerConfig, LiteLLM_JWTAuth, ProxyException +from litellm.proxy import proxy_server +from litellm.proxy._types import ( + JWTIssuerConfig, + LiteLLM_JWTAuth, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + LiteLLM_UserTable, + Member, + ProxyException, +) from litellm.proxy.auth.handle_jwt import JWKSUnreachableError, JWTHandler, jwks_unavailable_exception +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.mcp_server.gateway_dcr_flow import SubjectIdentity, SubjectTokenRefusal from litellm.proxy.mcp_server.idp_token_exchange import ( REJECTED_SUBJECT_TOKEN, SUBJECT_TOKEN_CHECK_FAULTED, SUBJECT_TOKEN_CHECK_UNAVAILABLE, TokenExchangePrerequisites, + exchange_idp_subject_token, identity_from_subject_token, token_exchange_available, ) @@ -246,3 +260,71 @@ async def test_agent_token_cannot_be_exchanged_for_a_user_identity(user_id: str assert isinstance(result, SubjectTokenRefusal) assert result.error == "invalid_request" assert "direct JWT authentication" in result.description + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("enforce_team", "team_id", "expected"), + ( + (True, "exchange-team", SubjectIdentity(user_id="exchange-user", team_id="exchange-team")), + (False, None, SubjectIdentity(user_id="exchange-user", team_id=None)), + (True, None, SubjectTokenRefusal(error="invalid_request", description=REJECTED_SUBJECT_TOKEN)), + ), + ids=("native-team", "teamless-without-enforcement", "teamless-under-enforcement"), +) +async def test_subject_exchange_applies_team_admission_before_returning_an_identity( + monkeypatch: pytest.MonkeyPatch, + enforce_team: bool, + team_id: str | None, + expected: SubjectIdentity | SubjectTokenRefusal, +) -> None: + cache: Final = UserApiKeyCache() + database: Final = object() + handler: Final = JWTHandler() + handler.update_environment( + prisma_client=database, + user_api_key_cache=cache, + litellm_jwtauth=LiteLLM_JWTAuth( + user_id_jwt_field="sub", + team_ids_jwt_field="groups", + enforce_team_based_model_access=enforce_team, + team_allowed_routes=["/token"], + ), + ) + signing_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + jwks_url: Final = "https://identity.invalid/exchange/jwks" + cache.set_cache( + f"litellm_jwt_auth_keys_{jwks_url}", + [json.loads(jwt.algorithms.RSAAlgorithm.to_jwk(signing_key.public_key()))], + ) + cache.set_cache( + "exchange-user", + LiteLLM_UserTable(user_id="exchange-user", user_role="internal_user", teams=[team_id] if team_id else []), + model_type=LiteLLM_UserTable, + ) + if team_id is not None: + cache.set_cache( + f"team_id:{team_id}", + LiteLLM_TeamTable(team_id=team_id, members_with_roles=[Member(user_id="exchange-user", role="user")]), + model_type=LiteLLM_TeamTable, + ) + cache.set_cache( + f"team_membership:exchange-user:{team_id}", + LiteLLM_TeamMembership(user_id="exchange-user", team_id=team_id), + model_type=LiteLLM_TeamMembership, + ) + monkeypatch.setenv("JWT_PUBLIC_KEY_URL", jwks_url) + monkeypatch.delenv("JWT_ISSUER", raising=False) + monkeypatch.delenv("JWT_AUDIENCE", raising=False) + monkeypatch.setattr(proxy_server, "jwt_handler", handler) + monkeypatch.setattr(proxy_server, "prisma_client", database) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + monkeypatch.setattr(proxy_server, "general_settings", {"enable_jwt_auth": True}) + monkeypatch.setattr(proxy_server, "premium_user", True) + monkeypatch.setattr(proxy_server, "llm_router", None) + monkeypatch.setattr(proxy_server, "llm_model_list", None) + token: Final = jwt.encode( + {"sub": "exchange-user", "groups": [team_id] if team_id else []}, signing_key, algorithm="RS256" + ) + request: Final = Request({"type": "http", "method": "POST", "path": "/token", "headers": []}) + assert await exchange_idp_subject_token(token, request) == expected diff --git a/tests/unit/proxy/middleware/test_billable_request_metrics_middleware.py b/tests/unit/proxy/middleware/test_billable_request_metrics_middleware.py index d2e961ad158..2488af7952c 100644 --- a/tests/unit/proxy/middleware/test_billable_request_metrics_middleware.py +++ b/tests/unit/proxy/middleware/test_billable_request_metrics_middleware.py @@ -105,6 +105,7 @@ def test_is_pure_asgi_not_base_http_middleware(): # SpendLogs-producing routes surfaced by the route-inventory audit ("/v1/search", (BillableCategory.LLM, "/search")), ("/v1/vector_stores/vs_1/search", (BillableCategory.LLM, "/search")), + ("/opensearch/docs/_search", (BillableCategory.LLM, "/opensearch")), ("/v1/rag/query", (BillableCategory.LLM, "/rag/query")), ("/rag/ingest", (BillableCategory.LLM, "/rag/ingest")), # Provider passthrough carries real inference and writes SpendLogs @@ -164,6 +165,18 @@ def test_classify_billable(path: str, expected: Tuple[BillableCategory, str]): "/v1beta/interactions/int_123/cancel", # observability passthrough writes no SpendLogs row "/langfuse/api/public/ingestion", + "/opensearch/docs/_doc", + "/opensearch/docs/_doc/_search", + "/opensearch/docs/_doc/1", + "/opensearch/docs/_bulk", + "/opensearch/docs/_update/1", + "/opensearch/docs/_count", + "/opensearch/docs/_mget", + "/opensearch/docs/_doc/search", + "/opensearch/docs/_update/search", + "/opensearch/docs/_create/embeddings", + "/opensearch/docs/_doc/completions", + "/opensearch/docs/_explain/rerank", # a bare provider prefix is not an inference call "/bedrock", "/eu_jev/v1/systemone/extra", diff --git a/tests/unit/proxy/pass_through_endpoints/test_codex_passthrough_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_codex_passthrough_endpoints.py index 1a4751b4598..322752782fd 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_codex_passthrough_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_codex_passthrough_endpoints.py @@ -11,6 +11,7 @@ import jwt import pytest import respx from cryptography.hazmat.primitives.asymmetric import rsa +from fastapi import Request from fastapi.testclient import TestClient from jwt.algorithms import RSAAlgorithm @@ -704,10 +705,23 @@ def _key_with_alias_to_budgeted_target(monkeypatch: pytest.MonkeyPatch, target_b return key +@pytest.mark.parametrize("custom_auth", (False, True), ids=("virtual-key", "custom-auth")) def test_native_route_refuses_a_saved_alias_when_its_target_is_over_budget( - client: TestClient, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch + client: TestClient, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, custom_auth: bool ) -> None: key: Final = _key_with_alias_to_budgeted_target(monkeypatch, target_budget=0.0) + if custom_auth: + monkeypatch.setattr(litellm, "enable_post_custom_auth_checks", True, raising=False) + identity: Final = proxy_server.user_api_key_cache.get_cache(key=hash_token(key), model_type=UserAPIKeyAuth) + assert identity is not None + + async def authenticate(request: Request, api_key: str) -> UserAPIKeyAuth: + return identity + + monkeypatch.setattr(proxy_server, "user_custom_auth", authenticate) + monkeypatch.setattr( + proxy_server, "general_settings", {**proxy_server.general_settings, "custom_auth_run_common_checks": True} + ) _account_catalog(respx_mock) provider: Final = _open_weights_provider(respx_mock) response: Final = client.post( diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 73eb1088d20..70f5e855e16 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -1,13 +1,16 @@ import asyncio import base64 import contextlib +import hashlib import itertools import json import logging import os +import re import traceback from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping from datetime import datetime, timezone +from pathlib import Path from types import MappingProxyType, SimpleNamespace from typing import Final, Literal from unittest import mock @@ -17,6 +20,9 @@ from urllib.parse import parse_qs import httpx import pytest import respx +from botocore.auth import SigV4Auth +from botocore.awsrequest import AWSRequest +from botocore.credentials import Credentials from fastapi import FastAPI, HTTPException, Request, Response from fastapi.routing import APIRoute from fastapi.responses import StreamingResponse @@ -35,7 +41,7 @@ from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessin from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException, ModifyResponseException from litellm.proxy._types import AwsErrorTypeProxyException, ProxyException from litellm.proxy.spend_tracking.spend_tracking_utils import get_spend_logs_id -from litellm.proxy.utils import ProxyLogging +from litellm.proxy.utils import ProxyLogging, hash_token from litellm.passthrough.main import raise_for_provider_status from litellm.llms.bedrock.common_utils import BedrockError from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe @@ -7653,8 +7659,16 @@ _BEDROCK_COUNT_ALIAS: Final = "count-log-bedrock" _BEDROCK_COUNT_MODEL: Final = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" +_BEDROCK_COUNT_BODY: Final = MappingProxyType( + {"model": _BEDROCK_COUNT_ALIAS, "messages": [{"role": "user", "content": "hi"}]} +) + + async def _bedrock_native_count( - monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, aws_reply: httpx.Response + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + aws_reply: httpx.Response, + request_body: Mapping[str, object] = _BEDROCK_COUNT_BODY, ) -> dict[str, object]: from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( handle_bedrock_count_tokens, @@ -7694,7 +7708,7 @@ async def _bedrock_native_count( request=request, fastapi_response=Response(), user_api_key_dict=MagicMock(), - request_body={"model": _BEDROCK_COUNT_ALIAS, "messages": [{"role": "user", "content": "hi"}]}, + request_body=dict(request_body), ) @@ -7748,6 +7762,81 @@ async def test_bedrock_native_count_tokens_logs_the_provider_status_when_the_cou ) +@pytest.mark.parametrize( + ("request_body", "problem"), + [ + ({"model": _BEDROCK_COUNT_ALIAS}, "Request must contain content to count tokens"), + ({"model": _BEDROCK_COUNT_ALIAS, "messages": []}, "messages parameter is required for Converse input"), + ({"model": _BEDROCK_COUNT_ALIAS, "messages": [{"role": "user"}]}, "Message 0 must have a 'content' field"), + ({"model": _BEDROCK_COUNT_ALIAS, "messages": [{"content": "hi"}]}, "Message 0 must have a 'role' field"), + ], +) +@pytest.mark.asyncio +async def test_bedrock_native_count_tokens_answers_400_for_a_malformed_body_without_calling_bedrock( + monkeypatch: pytest.MonkeyPatch, + httpx_transport: None, + respx_mock: respx.MockRouter, + request_body: Mapping[str, object], + problem: str, +) -> None: + with pytest.raises(HTTPException) as exc_info: + await _bedrock_native_count( + monkeypatch, respx_mock, httpx.Response(200, json={"inputTokens": 12}), request_body + ) + + assert (exc_info.value.status_code, exc_info.value.detail) == (400, {"error": problem}) + assert len(respx_mock.calls) == 0 + + +@pytest.mark.parametrize( + ("model", "model_id"), + [ + ("us", "us"), + ("global", "global"), + ("apac", "apac"), + ("bedrock/us", "us"), + ("converse/us", "us"), + ("invoke/us", "us"), + ("bedrock/converse/global", "global"), + ], +) +@pytest.mark.asyncio +async def test_bedrock_native_count_tokens_for_a_bare_region_prefix_answers_bedrocks_own_400( + monkeypatch: pytest.MonkeyPatch, + httpx_transport: None, + respx_mock: respx.MockRouter, + model: str, + model_id: str, +) -> None: + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + handle_bedrock_count_tokens, + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "AKIATESTACCESSKEY") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "test-secret") + monkeypatch.setenv("AWS_REGION_NAME", "us-east-1") + invalid_model: Final = {"message": "The provided model identifier is invalid."} + respx_mock.post(host="bedrock-runtime.us-east-1.amazonaws.com", path=f"/model/{model_id}/count-tokens").mock( + return_value=httpx.Response(400, json=invalid_model) + ) + request: Final = Request({"type": "http", "method": "POST", "path": f"/bedrock/model/{model}/count-tokens", "headers": []}) + + with pytest.raises(HTTPException) as exc_info: + await handle_bedrock_count_tokens( + endpoint=f"model/{model}/count-tokens", + request=request, + fastapi_response=Response(), + user_api_key_dict=MagicMock(), + request_body={"model": model, "messages": [{"role": "user", "content": "hi"}]}, + ) + + assert exc_info.value.status_code == 400 + assert json.loads(exc_info.value.detail["error"]) == invalid_model + assert len(respx_mock.calls) == 1 + + @pytest.mark.asyncio async def test_bedrock_native_count_tokens_without_a_model_logs_the_400_it_answers( token_count_log_lines: Callable[[], tuple[dict[str, object], ...]], @@ -7814,9 +7903,7 @@ async def test_bedrock_native_count_tokens_with_a_region_only_model_answers_the_ request_body={"model": "us"}, ) - assert exc_info.value.detail == { - "error": "CountTokens processing error: Request must contain content to count tokens" - } + assert exc_info.value.detail == {"error": "Request must contain content to count tokens"} assert token_count_log_lines() == ( { "litellm_call_id": "bedrock-native-region-only", @@ -10169,3 +10256,914 @@ class TestTinyFishRouteTimeout: monkeypatch.setattr(proxy_server, "general_settings", {"pass_through_request_timeout": 30}, raising=False) assert _tinyfish_route_timeout() is None + + +OPENSEARCH_HOST: Final = "abc123.us-east-1.aoss.amazonaws.com" +OPENSEARCH_ACCESS_KEY: Final = "AKIDOPENSEARCHPASSTHROUGH" +OPENSEARCH_SECRET_KEY: Final = "opensearch-passthrough-secret" +DOCS_UPSTREAM: Final = f"https://{OPENSEARCH_HOST}/prod-docs-v2" +JSON_TYPE: Final = "application/json" +NDJSON_TYPE: Final = "application/x-ndjson" +_SIGV4_AUTHORIZATION: Final = re.compile( + r"\AAWS4-HMAC-SHA256 Credential=(?P[^/]+)/\d{8}/(?P[^/]+)/(?P[^/]+)/aws4_request, " + r"SignedHeaders=(?P[^,]+), Signature=(?P[0-9a-f]{64})\Z" +) + + +def _verified_sigv4_scope(sent: httpx.Request) -> tuple[str, str]: + authorization: Final = _SIGV4_AUTHORIZATION.match(sent.headers["authorization"]) + assert authorization is not None, sent.headers["authorization"] + signed_names: Final = authorization["signed"].split(";") + assert {"content-type", "host", "x-amz-content-sha256", "x-amz-date"} <= set(signed_names), signed_names + assert sent.headers["x-amz-content-sha256"] == hashlib.sha256(sent.content).hexdigest() + received: Final = AWSRequest( + method=sent.method, + url=str(sent.url), + data=sent.content, + headers={name: sent.headers[name] for name in signed_names}, + ) + received.context["timestamp"] = sent.headers["x-amz-date"] + signer: Final = SigV4Auth( + Credentials(OPENSEARCH_ACCESS_KEY, OPENSEARCH_SECRET_KEY), authorization["service"], authorization["region"] + ) + expected: Final = signer.signature(signer.string_to_sign(received, signer.canonical_request(received)), received) + assert (authorization["key"], authorization["signature"]) == (OPENSEARCH_ACCESS_KEY, expected) + return authorization["service"], authorization["region"] + + +def _opensearch_caller(*permissions: str, index_name: str = "docs") -> UserAPIKeyAuth: + grant: Final = { + "allowed_vector_store_indexes": [{"index_name": index_name, "index_permissions": list(permissions)}] + } + return UserAPIKeyAuth(api_key="sk-virtual", metadata=grant) + + +def _act_as(monkeypatch: pytest.MonkeyPatch, caller: UserAPIKeyAuth) -> None: + from litellm.proxy.proxy_server import app + + monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, lambda: caller) + + +@pytest.fixture +def opensearch_client(monkeypatch: pytest.MonkeyPatch) -> Iterator[TestClient]: + from litellm.proxy.proxy_server import app + from litellm.types.vector_stores import IndexCreateLiteLLMParams, LiteLLM_ManagedVectorStoreIndex + from litellm.vector_stores.vector_store_registry import VectorStoreIndexRegistry, VectorStoreRegistry + + monkeypatch.delenv("SERVER_ROOT_PATH", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + indexes: Final = ( + ("docs", "prod-docs-v2", "opensearch-store"), + ("hr", "hr-private", "opensearch-store"), + ("no-credentials", "prod-docs-v3", "opensearch-store-without-credentials"), + ("azure-docs", "azure-docs", "azure-store"), + ) + monkeypatch.setattr( + litellm, + "vector_store_index_registry", + VectorStoreIndexRegistry( + [ + LiteLLM_ManagedVectorStoreIndex( + id=f"index-{name}", + index_name=name, + litellm_params=IndexCreateLiteLLMParams(vector_store_index=provider_index, vector_store_name=store), + ) + for name, provider_index, store in indexes + ] + ), + ) + monkeypatch.setattr( + litellm, + "vector_store_registry", + VectorStoreRegistry( + [ + { + "vector_store_id": "opensearch-store", + "vector_store_name": "opensearch-store", + "custom_llm_provider": "opensearch", + "litellm_params": { + "api_base": f"https://{OPENSEARCH_HOST}", + "aws_access_key_id": OPENSEARCH_ACCESS_KEY, + "aws_secret_access_key": OPENSEARCH_SECRET_KEY, + }, + }, + { + "vector_store_id": "opensearch-store-without-credentials", + "vector_store_name": "opensearch-store-without-credentials", + "custom_llm_provider": "opensearch", + "litellm_params": { + "api_base": f"https://{OPENSEARCH_HOST}", + "aws_profile_name": "opensearch-missing-profile", + }, + }, + { + "vector_store_id": "azure-store", + "vector_store_name": "azure-store", + "custom_llm_provider": "azure_ai", + "litellm_params": {"api_base": "https://search.example.search.windows.net"}, + }, + ] + ), + ) + yield TestClient(app) + + +class TestOpenSearchProxyRoute: + SEARCH_BODY: Final = b'{"query":{"match":{"text":"remote work"}}}' + HITS: Final = MappingProxyType( + { + "took": 3, + "timed_out": False, + "hits": { + "total": {"value": 1, "relation": "eq"}, + "hits": [{"_index": "prod-docs-v2", "_id": "1", "_score": 1.2, "_source": {"text": "Remote work"}}], + }, + } + ) + + @pytest.mark.parametrize("content_type", [JSON_TYPE, "Application/JSON; charset=UTF-8"]) + def test_a_read_key_search_is_signed_for_the_mapped_index_and_returns_the_hits( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, content_type: str + ) -> None: + _act_as(monkeypatch, _opensearch_caller("read")) + with respx.mock(assert_all_called=True) as upstream: + route = upstream.post(f"{DOCS_UPSTREAM}/_search").mock( + return_value=httpx.Response(200, json=dict(self.HITS)) + ) + response = opensearch_client.post( + "/opensearch/docs/_search", + content=self.SEARCH_BODY, + headers={"Content-Type": content_type, "Authorization": "Bearer sk-virtual"}, + ) + + assert (response.status_code, response.json()) == (200, dict(self.HITS)) + sent: Final = route.calls.last.request + assert (sent.content, sent.headers["content-type"]) == (self.SEARCH_BODY, JSON_TYPE) + assert _verified_sigv4_scope(sent) == ("aoss", "us-east-1") + + def test_a_team_grant_lets_the_key_search( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch + ) -> None: + team_grant: Final = {"allowed_vector_store_indexes": [{"index_name": "docs", "index_permissions": ["read"]}]} + _act_as(monkeypatch, UserAPIKeyAuth(api_key="sk-virtual", team_id="team-a", team_metadata=team_grant)) + with respx.mock(assert_all_called=True) as upstream: + upstream.post(f"{DOCS_UPSTREAM}/_search").mock(return_value=httpx.Response(200, json=dict(self.HITS))) + response = opensearch_client.post( + "/opensearch/docs/_search", content=self.SEARCH_BODY, headers={"Content-Type": JSON_TYPE} + ) + + assert (response.status_code, response.json()) == (200, dict(self.HITS)) + + @pytest.mark.parametrize( + "method, api", + [ + ("POST", "_doc"), + ("PUT", "_doc/1"), + ("POST", "_doc/1"), + ("DELETE", "_doc/1"), + ("PUT", "_create/1"), + ("POST", "_update/1"), + ("POST", "_bulk"), + ("POST", "_delete_by_query"), + ("POST", "_update_by_query"), + ], + ) + def test_a_read_only_key_cannot_write( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, method: str, api: str + ) -> None: + _act_as(monkeypatch, _opensearch_caller("read")) + with respx.mock(assert_all_called=False) as upstream: + route = upstream.route(host=OPENSEARCH_HOST).mock(return_value=httpx.Response(201, json={})) + response = opensearch_client.request( + method, + f"/opensearch/docs/{api}", + content=b'{"title":"Remote work"}', + headers={"Content-Type": JSON_TYPE}, + ) + + assert (response.status_code, route.called) == (403, False) + + @pytest.mark.parametrize( + "document", + [b'{"title":"Remote work","index":"hr-private"}', b'{"title":"Remote work","stream":true}'], + ids=["field-named-index", "field-named-stream"], + ) + def test_a_read_write_key_writes_the_document_as_sent( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, document: bytes + ) -> None: + _act_as(monkeypatch, _opensearch_caller("read", "write")) + with respx.mock(assert_all_called=True) as upstream: + route = upstream.post(f"{DOCS_UPSTREAM}/_doc").mock( + return_value=httpx.Response(201, json={"_id": "abc", "result": "created"}) + ) + response = opensearch_client.post( + "/opensearch/docs/_doc", content=document, headers={"Content-Type": JSON_TYPE} + ) + + assert (response.status_code, response.json()["result"]) == (201, "created") + assert route.calls.last.request.content == document + assert _verified_sigv4_scope(route.calls.last.request) == ("aoss", "us-east-1") + + @pytest.mark.parametrize( + "path, body, status_code", + [ + pytest.param("_update/1", b'{"doc":{},"_source":true}', 403, id="update-body-source"), + pytest.param("_update/1?_source=text", b'{"doc":{"title":"b"}}', 403, id="update-source-param"), + pytest.param( + "_update/1", + b'{"script":{"source":"throw new IllegalArgumentException(ctx._source.text)"}}', + 403, + id="update-script", + ), + pytest.param("_update/1", b'{"doc":{} /* c */ ,"_source":true}', 400, id="update-body-json-with-comment"), + pytest.param("_bulk", b'{"update":{"_id":"1","_source":true}}\n{"doc":{}}\n', 403, id="bulk-action-source"), + pytest.param( + "_bulk", + b'{"index":{"_id":"2"}}\n{"title":"a"}\n{"update":{"_id":"1"}}\n{"doc":{},"_source":true}\n', + 403, + id="bulk-update-payload-source", + ), + pytest.param("_bulk?_source_includes=text", b'{"update":{"_id":"1"}}\n{"doc":{}}\n', 403, id="bulk-param"), + pytest.param("_delete_by_query", b'{"query":{"prefix":{"text":"a"}}}', 403, id="delete-by-query"), + pytest.param("_update_by_query", b'{"query":{"match_all":{}}}', 403, id="update-by-query"), + ], + ) + def test_a_write_only_key_cannot_read_documents_back_through_a_write( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, path: str, body: bytes, status_code: int + ) -> None: + _act_as(monkeypatch, _opensearch_caller("write")) + with respx.mock(assert_all_called=False) as upstream: + route = upstream.route(host=OPENSEARCH_HOST).mock( + return_value=httpx.Response(200, json={"get": {"_source": {"text": "Remote work"}}}) + ) + response = opensearch_client.post( + f"/opensearch/docs/{path}", content=body, headers={"Content-Type": NDJSON_TYPE} + ) + + assert (response.status_code, route.called) == (status_code, False) + assert "Remote work" not in response.text + + @pytest.mark.parametrize( + "caller, path, body", + [ + pytest.param( + _opensearch_caller("write"), "_update/1", b'{"doc":{"script":"a field","_source":"b"}}', id="write-only" + ), + pytest.param( + _opensearch_caller("read", "write"), "_update/1?_source=text", b'{"doc":{},"_source":true}', id="update" + ), + pytest.param( + _opensearch_caller("read", "write"), + "_bulk", + b'{"update":{"_id":"1","_source":true}}\n{"doc":{},"script":{"source":"ctx._source.n += 1"}}\n', + id="bulk", + ), + pytest.param( + _opensearch_caller("read", "write"), "_delete_by_query", b'{"query":{"term":{"a":1}}}', id="by-query" + ), + ], + ) + def test_a_write_that_reads_no_document_or_has_the_read_grant_is_sent_as_is( + self, + opensearch_client: TestClient, + monkeypatch: pytest.MonkeyPatch, + caller: UserAPIKeyAuth, + path: str, + body: bytes, + ) -> None: + _act_as(monkeypatch, caller) + with respx.mock(assert_all_called=True) as upstream: + route = upstream.post(url__startswith=f"{DOCS_UPSTREAM}/{path.partition('?')[0]}").mock( + return_value=httpx.Response(200, json={"result": "updated"}) + ) + response = opensearch_client.post( + f"/opensearch/docs/{path}", content=body, headers={"Content-Type": NDJSON_TYPE} + ) + + assert (response.status_code, route.calls.last.request.content) == (200, body) + + @pytest.mark.parametrize( + "caller, index_name", + [ + (_opensearch_caller("read", "write", index_name="hr"), "docs"), + (UserAPIKeyAuth(api_key="sk-virtual"), "docs"), + (UserAPIKeyAuth(api_key="sk-virtual"), "unregistered"), + ], + ids=["grant-on-another-index", "no-grant", "no-grant-unregistered-index"], + ) + def test_a_key_without_a_grant_on_the_index_is_refused( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, caller: UserAPIKeyAuth, index_name: str + ) -> None: + _act_as(monkeypatch, caller) + with respx.mock(assert_all_called=False) as upstream: + route = upstream.route(host=OPENSEARCH_HOST).mock(return_value=httpx.Response(200, json={})) + response = opensearch_client.post( + f"/opensearch/{index_name}/_search", content=self.SEARCH_BODY, headers={"Content-Type": JSON_TYPE} + ) + + assert (response.status_code, route.called) == (403, False) + assert index_name in response.json()["detail"] + + @pytest.mark.parametrize( + "method, path", + [ + ("POST", "/opensearch/docs/_search/template"), + ("POST", "/opensearch/docs/_msearch"), + ("GET", "/opensearch/docs/_msearch"), + ("PUT", "/opensearch/docs/_mapping"), + ("GET", "/opensearch/docs/_settings"), + ("POST", "/opensearch/docs/_close"), + ("POST", "/opensearch/docs/_refresh"), + ("GET", "/opensearch/docs/_doc/1/x"), + ("PATCH", "/opensearch/docs/_doc/1"), + ("PUT", "/opensearch/docs"), + ("DELETE", "/opensearch/docs"), + ], + ) + def test_an_unlisted_api_is_for_proxy_admins_only( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, method: str, path: str + ) -> None: + _act_as(monkeypatch, _opensearch_caller("read", "write")) + with respx.mock(assert_all_called=False) as upstream: + route = upstream.route(host=OPENSEARCH_HOST).mock(return_value=httpx.Response(200, json={})) + response = opensearch_client.request( + method, path, content=b'{"source":"{}"}', headers={"Content-Type": JSON_TYPE} + ) + + assert (response.status_code, route.called) == (403, False) + + @pytest.mark.parametrize( + "method, path, upstream_path, status", + [ + pytest.param("GET", "/opensearch/docs", "/prod-docs-v2", 200, id="get-index"), + pytest.param("HEAD", "/opensearch/docs", "/prod-docs-v2", 200, id="index-exists"), + pytest.param("HEAD", "/opensearch/docs/_doc/1", "/prod-docs-v2/_doc/1", 404, id="doc-exists"), + pytest.param("HEAD", "/opensearch/docs/_source/1", "/prod-docs-v2/_source/1", 200, id="source-exists"), + pytest.param( + "POST", "/opensearch/docs/_explain/_bulk1", "/prod-docs-v2/_explain/_bulk1", 200, id="id-named-bulk" + ), + ], + ) + def test_a_read_key_reaches_the_index_and_its_existence_checks( + self, + opensearch_client: TestClient, + monkeypatch: pytest.MonkeyPatch, + method: str, + path: str, + upstream_path: str, + status: int, + ) -> None: + _act_as(monkeypatch, _opensearch_caller("read")) + with respx.mock(assert_all_called=True) as upstream: + route = upstream.route(method=method, host=OPENSEARCH_HOST, path=upstream_path).mock( + return_value=httpx.Response(status, content=b"" if method == "HEAD" else b"{}") + ) + response = opensearch_client.request( + method, path, content=b"" if method == "HEAD" else self.SEARCH_BODY, headers={"Content-Type": JSON_TYPE} + ) + + assert response.status_code == status + assert _verified_sigv4_scope(route.calls.last.request) == ("aoss", "us-east-1") + + def test_a_proxy_admin_can_call_an_unlisted_api( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch + ) -> None: + settings: Final = {"prod-docs-v2": {"settings": {"index": {"number_of_shards": "2"}}}} + _act_as(monkeypatch, UserAPIKeyAuth(api_key="sk-admin", user_role=LitellmUserRoles.PROXY_ADMIN)) + with respx.mock(assert_all_called=True) as upstream: + route = upstream.get(f"{DOCS_UPSTREAM}/_settings").mock(return_value=httpx.Response(200, json=settings)) + response = opensearch_client.get("/opensearch/docs/_settings") + + assert (response.status_code, response.json()) == (200, settings) + assert _verified_sigv4_scope(route.calls.last.request) == ("aoss", "us-east-1") + + @pytest.mark.parametrize( + "index_name, detail", + [ + ("unknown", "Index unknown is not a litellm managed vector store index"), + ("azure-docs", "not an OpenSearch vector store"), + ], + ) + def test_an_index_that_is_not_a_managed_opensearch_index_is_a_400( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, index_name: str, detail: str + ) -> None: + _act_as(monkeypatch, _opensearch_caller("read", index_name=index_name)) + with respx.mock(assert_all_called=False) as upstream: + route = upstream.route().mock(return_value=httpx.Response(200, json={})) + response = opensearch_client.post( + f"/opensearch/{index_name}/_search", content=self.SEARCH_BODY, headers={"Content-Type": JSON_TYPE} + ) + + assert (response.status_code, route.called) == (400, False) + assert detail in response.json()["detail"] + + def test_an_index_whose_store_name_two_stores_share_is_a_400( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch + ) -> None: + from litellm.vector_stores.vector_store_registry import VectorStoreRegistry + + same_name: Final = ( + ("opensearch-store", f"https://{OPENSEARCH_HOST}"), + ("opensearch-store-copy", "https://other123.us-east-1.aoss.amazonaws.com"), + ) + monkeypatch.setattr( + litellm, + "vector_store_registry", + VectorStoreRegistry( + [ + { + "vector_store_id": store_id, + "vector_store_name": "opensearch-store", + "custom_llm_provider": "opensearch", + "litellm_params": {"api_base": api_base}, + } + for store_id, api_base in same_name + ] + ), + ) + _act_as(monkeypatch, _opensearch_caller("read")) + with respx.mock(assert_all_called=False) as upstream: + route = upstream.route().mock(return_value=httpx.Response(200, json={})) + response = opensearch_client.post( + "/opensearch/docs/_search", content=self.SEARCH_BODY, headers={"Content-Type": JSON_TYPE} + ) + + assert (response.status_code, route.called) == (400, False) + assert "2 vector stores are named opensearch-store" in response.json()["detail"] + + @pytest.mark.parametrize( + "path", + [ + "/opensearch/docs/", + "/opensearch/docs//_search", + "/opensearch/docs/%2E/_search", + "/opensearch/docs/%2E%2E/hr/_search", + "/opensearch/docs/_search/%2E%2E/%2E%2E/hr/_search", + ], + ) + def test_a_path_with_an_empty_or_dot_segment_is_a_400( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, path: str + ) -> None: + _act_as(monkeypatch, _opensearch_caller("read", "write")) + with respx.mock(assert_all_called=False) as upstream: + route = upstream.route(host=OPENSEARCH_HOST).mock(return_value=httpx.Response(200, json={})) + response = opensearch_client.post(path, content=self.SEARCH_BODY, headers={"Content-Type": JSON_TYPE}) + + assert (response.status_code, route.called) == (400, False) + + @pytest.mark.parametrize( + "api, content_type, body, detail", + [ + pytest.param( + "_search", + JSON_TYPE, + b'{"query":{"terms":{"user":{"index":"hr-private","id":"1","path":"followers"}}}}', + "'index'", + id="search-terms-lookup", + ), + pytest.param( + "_search", + JSON_TYPE, + b'{"query":{"more_like_this":{"like":[{"_index":"hr-private","_id":"1"}]}}}', + "'_index'", + id="search-more-like-this", + ), + pytest.param( + "_count", + JSON_TYPE, + b'{"query":{"terms":{"user":{"\\u0069ndex":"hr-private","id":"1","path":"followers"}}}}', + "'index'", + id="count-escaped-key", + ), + pytest.param( + "_search", + JSON_TYPE, + b'{"query":{"terms":{"user":{"index":"hr-private","id":"1","path":"f"}}},"query":{"match_all":{}}}', + "'index'", + id="search-duplicate-key", + ), + pytest.param( + "_delete_by_query", + JSON_TYPE, + b'{"query":{"terms":{"user":{"index":"hr-private","id":"1","path":"followers"}}}}', + "'index'", + id="delete-by-query-terms-lookup", + ), + pytest.param( + "_mget", + JSON_TYPE, + b'{"docs":[{"_id":"1"},{"_index":"hr-private","_id":"2"}]}', + "'_index'", + id="mget-docs", + ), + pytest.param( + "_bulk", + NDJSON_TYPE, + b'{"index":{"_id":"1"}}\n{"title":"a"}\n{"delete":{"_index":"hr-private","_id":"2"}}\n', + "'_index'", + id="bulk-action-index", + ), + pytest.param( + "_bulk", + NDJSON_TYPE, + b'{"index":{"_index":"prod-docs-v2","_id":"1"}}\n{"title":"a"}\n', + "'_index'", + id="bulk-action-names-the-provider-index", + ), + pytest.param( + "_bulk", + NDJSON_TYPE, + b'{"index":{"_index":"docs","_index":"hr-private","_id":"1"}}\n{"title":"a"}\n', + "'_index'", + id="bulk-action-duplicate-index", + ), + pytest.param( + "_bulk", + NDJSON_TYPE, + b'{"index":{"_id":"1"}}\n\n{"delete":{"_index":"hr-private","_id":"2"}}\n', + "'_index'", + id="bulk-blank-source-line", + ), + pytest.param( + "_bulk", + NDJSON_TYPE, + b'{"index":{"_id":"1","pipeline":"reroute"}}\n{"title":"a"}\n', + "'pipeline'", + id="bulk-action-pipeline", + ), + pytest.param("_bulk", NDJSON_TYPE, b"not json\n", "action line", id="bulk-unparseable-action"), + pytest.param( + "_search", + JSON_TYPE, + b'{"profile":true,"query":{"wrapper":{"query":"eyJ0ZXJtcyI6e319"}}}', + "'wrapper'", + id="search-wrapper-query", + ), + pytest.param( + "_search", + JSON_TYPE, + b'{"query":{"geo_shape":{"location":{"indexed_shape":{"id":"deu","path":"location"}}}}}', + "'indexed_shape'", + id="search-indexed-shape", + ), + pytest.param( + "_search", + JSON_TYPE, + b'{"query":{"match_all":{}},"search_pipeline":{"request_processors":[]}}', + "'search_pipeline'", + id="search-temporary-pipeline", + ), + pytest.param( + "_mget?source=%7B%22docs%22%3A%5B%7B%22_index%22%3A%22hr-private%22%2C%22_id%22%3A%221%22%7D%5D%7D" + "&source_content_type=application%2Fjson", + JSON_TYPE, + b"", + "source query parameter", + id="mget-source-query", + ), + pytest.param( + "_search", + "application/x-ndjson; multipart/form-data", + b'{"stream":true,"custom_body":{}}', + "application/json", + id="multipart-parameter", + ), + pytest.param( + "_doc", "application/json; boundary=x", b'{"title":"a"}', "application/json", id="content-type-parameter" + ), + pytest.param("_doc?pipeline=reroute", JSON_TYPE, b'{"title":"a"}', "pipeline", id="pipeline-query"), + pytest.param("_search", JSON_TYPE, b'{/*c*/"query":{"match_all":{}}}', "could not parse", id="comment"), + pytest.param("_search", "application/smile", b"\x3a\x29\x0a\x01", "application/json", id="smile"), + pytest.param("_search", "application/yaml", b"query:\n match_all: {}\n", "application/json", id="yaml"), + pytest.param( + "_search", JSON_TYPE, b'{"query":' + b"[" * 100_000 + b"]" * 100_000 + b"}", "nested", id="deep-nesting" + ), + ], + ) + def test_a_body_or_query_that_could_reach_another_index_is_a_400( + self, + opensearch_client: TestClient, + monkeypatch: pytest.MonkeyPatch, + api: str, + content_type: str, + body: bytes, + detail: str, + ) -> None: + _act_as(monkeypatch, _opensearch_caller("read", "write")) + with respx.mock(assert_all_called=False) as upstream: + route = upstream.route(host=OPENSEARCH_HOST).mock(return_value=httpx.Response(200, json={})) + response = opensearch_client.post( + f"/opensearch/docs/{api}", content=body, headers={"Content-Type": content_type} + ) + + assert (response.status_code, route.called) == (400, False) + assert detail in response.json()["detail"] + + @pytest.mark.parametrize( + "api, body", + [ + pytest.param( + "_bulk", + b'{"index":{"_id":"1"}}\n{"title":"a","index":"hr-private"}\n{"delete":{"_id":"2"}}\n\n' + b'{"update":{"_id":"3"}}\r\n{"doc":{"_index":"hr-private"}}\n', + id="bulk", + ), + ], + ) + @pytest.mark.parametrize("content_type", [NDJSON_TYPE, "application/json"]) + def test_an_ndjson_body_without_index_references_is_sent_byte_for_byte( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, api: str, body: bytes, content_type: str + ) -> None: + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", SHARED_MASTER_KEY) + with respx.mock(assert_all_called=True) as upstream: + route = upstream.post(f"{DOCS_UPSTREAM}/{api}").mock( + return_value=httpx.Response(200, json={"errors": False, "items": []}) + ) + response = opensearch_client.post( + f"/opensearch/docs/{api}", + content=body, + headers={"Content-Type": content_type, "Authorization": f"Bearer {SHARED_MASTER_KEY}"}, + ) + + assert response.status_code == 200 + sent: Final = route.calls.last.request + assert (sent.content, sent.headers["content-type"]) == (body, content_type) + assert _verified_sigv4_scope(sent) == ("aoss", "us-east-1") + + @pytest.mark.parametrize( + "document", + [ + b'{"title":"Remote work","model":"another-model"}', + b'{"title":"Remote work","api_base":"https://elsewhere.example"}', + b'{"title":"Remote work","base_url":"https://elsewhere.example"}', + ], + ids=["model", "api_base", "base_url"], + ) + def test_a_document_field_named_like_an_llm_param_is_written_for_a_model_limited_key( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, document: bytes + ) -> None: + from litellm.proxy import proxy_server + + key: Final = "sk-opensearch-model-limited-writer" + grant: Final = {"allowed_vector_store_indexes": [{"index_name": "docs", "index_permissions": ["write"]}]} + cache: Final = DualCache() + cache.set_cache( + key=hash_token(key), value=UserAPIKeyAuth(token=hash_token(key), models=["gpt-x"], metadata=grant) + ) + monkeypatch.setattr(proxy_server, "master_key", SHARED_MASTER_KEY) + monkeypatch.setattr(proxy_server, "prisma_client", "test-database") + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + with respx.mock(assert_all_called=True) as upstream: + route = upstream.post(f"{DOCS_UPSTREAM}/_doc").mock( + return_value=httpx.Response(201, json={"_id": "abc", "result": "created"}) + ) + response = opensearch_client.post( + "/opensearch/docs/_doc", + content=document, + headers={"Content-Type": JSON_TYPE, "Authorization": f"Bearer {key}"}, + ) + + assert (response.status_code, route.calls.last.request.content) == (201, document) + + def test_a_bulk_action_naming_the_path_index_is_sent_to_the_mapped_index( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch + ) -> None: + _act_as(monkeypatch, _opensearch_caller("write")) + with respx.mock(assert_all_called=True) as upstream: + route = upstream.post(f"{DOCS_UPSTREAM}/_bulk").mock( + return_value=httpx.Response(200, json={"errors": False, "items": []}) + ) + response = opensearch_client.post( + "/opensearch/docs/_bulk", + content=b'{"index":{"_index":"docs","_id":"1"}}\n{"title":"a","_index":"docs"}\n' + b'{"delete":{"_id":"2"}}\n{"update":{"_id":"3","_index":"docs"}}\n{"doc":{"title":"b"}}\n', + headers={"Content-Type": NDJSON_TYPE}, + ) + + assert response.status_code == 200 + sent: Final = route.calls.last.request + assert sent.content == ( + b'{"index":{"_index":"prod-docs-v2","_id":"1"}}\n{"title":"a","_index":"docs"}\n' + b'{"delete":{"_id":"2"}}\n{"update":{"_id":"3","_index":"prod-docs-v2"}}\n{"doc":{"title":"b"}}\n' + ) + assert _verified_sigv4_scope(sent) == ("aoss", "us-east-1") + + @pytest.mark.parametrize( + "method, path, body, sent_query", + [ + pytest.param( + "GET", + "/opensearch/docs/_search?q=title:remote+work&routing=a&routing=b", + b"", + b"q=title%3Aremote%20work&routing=a&routing=b", + id="space-and-repeated-key", + ), + pytest.param( + "POST", + "/opensearch/docs/_search", + b'{"size":1,"query_params":{"routing":"z"}}', + b"", + id="body-query-params-ignored", + ), + pytest.param( + "POST", + "/opensearch/docs/_search?search_pipeline=hybrid-min-max", + b'{"query":{"hybrid":{"queries":[{"match":{"title":"remote"}},' + b'{"knn":{"embedding":{"vector":[0.1,0.2],"k":2,"filter":{"term":{"lang":"en"}}}}}]}}}', + b"search_pipeline=hybrid-min-max", + id="hybrid-query-with-named-search-pipeline", + ), + ], + ) + def test_the_query_string_is_signed_as_sent( + self, + opensearch_client: TestClient, + monkeypatch: pytest.MonkeyPatch, + method: str, + path: str, + body: bytes, + sent_query: bytes, + ) -> None: + _act_as(monkeypatch, _opensearch_caller("read")) + with respx.mock(assert_all_called=True) as upstream: + route = upstream.route(method=method, host=OPENSEARCH_HOST, path="/prod-docs-v2/_search").mock( + return_value=httpx.Response(200, json=dict(self.HITS)) + ) + response = opensearch_client.request(method, path, content=body, headers={"Content-Type": JSON_TYPE}) + + assert response.status_code == 200 + sent: Final = route.calls.last.request + assert sent.url.query == sent_query + assert _verified_sigv4_scope(sent) == ("aoss", "us-east-1") + + def test_document_fields_named_like_request_params_leave_the_rate_limiter_alone( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch + ) -> None: + from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 + from litellm.proxy.utils import InternalUsageCache + + limiter: Final = PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(DualCache())) + + class RateLimit(CustomLogger): + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], + call_type: CallTypesLiteral, + ) -> dict[str, object]: + await limiter.async_pre_call_hook(user_api_key_dict, cache, data, call_type) + return data + + monkeypatch.setattr(litellm, "callbacks", [RateLimit()]) + grant: Final = {"index_name": "docs", "index_permissions": ["write"]} + _act_as( + monkeypatch, + UserAPIKeyAuth( + api_key="sk-opensearch-rate-limited-writer", + metadata={"allowed_vector_store_indexes": [grant], "model_rpm_limit": {"text-embedding-3-small": 1}}, + ), + ) + documents: Final = ( + b'{"text":"a","metadata":"handbook"}', + b'{"text":"b","metadata":["handbook"]}', + b'{"text":"c","model":"text-embedding-3-small"}', + b'{"text":"d","model":"text-embedding-3-small"}', + ) + with respx.mock(assert_all_called=True) as upstream: + route = upstream.post(f"{DOCS_UPSTREAM}/_doc").mock( + return_value=httpx.Response(201, json={"result": "created"}) + ) + statuses: Final = tuple( + opensearch_client.post( + "/opensearch/docs/_doc", content=document, headers={"Content-Type": JSON_TYPE} + ).status_code + for document in documents + ) + + assert statuses == (201, 201, 201, 201) + assert tuple(call.request.content for call in route.calls) == documents + + @pytest.mark.parametrize( + "document", [b'[{"text":"array body"}]', b'{"text":"\xff\xfe"}'], ids=["json-array", "invalid-utf-8"] + ) + def test_a_malformed_document_gets_the_opensearch_answer( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, document: bytes + ) -> None: + rejected: Final = {"error": {"type": "mapper_parsing_exception"}, "status": 400} + _act_as(monkeypatch, _opensearch_caller("write")) + with respx.mock(assert_all_called=True) as upstream: + route = upstream.post(f"{DOCS_UPSTREAM}/_doc").mock(return_value=httpx.Response(400, json=rejected)) + response = opensearch_client.post( + "/opensearch/docs/_doc", content=document, headers={"Content-Type": JSON_TYPE} + ) + + assert (response.status_code, response.json()) == (400, rejected) + assert route.calls.last.request.content == document + + def test_an_opensearch_error_keeps_its_status_and_body( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch + ) -> None: + not_found: Final = {"error": {"type": "index_not_found_exception", "index": "prod-docs-v2"}, "status": 404} + _act_as(monkeypatch, _opensearch_caller("read")) + with respx.mock(assert_all_called=True) as upstream: + upstream.post(f"{DOCS_UPSTREAM}/_search").mock(return_value=httpx.Response(404, json=not_found)) + response = opensearch_client.post( + "/opensearch/docs/_search", content=self.SEARCH_BODY, headers={"Content-Type": JSON_TYPE} + ) + + assert (response.status_code, response.json()) == (404, not_found) + + def test_missing_aws_credentials_are_a_401_before_anything_is_sent( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + for name in ("AWS_CONFIG_FILE", "AWS_SHARED_CREDENTIALS_FILE"): + empty_file: Path = tmp_path / name + empty_file.write_text("") + monkeypatch.setenv(name, str(empty_file)) + _act_as(monkeypatch, _opensearch_caller("read", index_name="no-credentials")) + with respx.mock(assert_all_called=False) as upstream: + route = upstream.route(host=OPENSEARCH_HOST).mock(return_value=httpx.Response(200, json={})) + response = opensearch_client.post( + "/opensearch/no-credentials/_search", content=self.SEARCH_BODY, headers={"Content-Type": JSON_TYPE} + ) + + assert (response.status_code, route.called) == (401, False) + + @pytest.mark.parametrize( + "method, api, action", + [ + ("GET", "_doc/1", "es:ESHttpGet"), + ("HEAD", "_doc/1", "es:ESHttpHead"), + ("PUT", "_doc/1", "es:ESHttpPut"), + ("DELETE", "_doc/1", "es:ESHttpDelete"), + ("POST", "_search", "es:ESHttpPost"), + ], + ) + def test_a_web_identity_session_allows_the_verb_sent_to_a_managed_domain( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, method: str, api: str, action: str + ) -> None: + from litellm.types.vector_stores import IndexCreateLiteLLMParams, LiteLLM_ManagedVectorStoreIndex + from litellm.vector_stores.vector_store_registry import VectorStoreIndexRegistry, VectorStoreRegistry + + domain: Final = "search-docs-abc123.us-east-1.es.amazonaws.com" + monkeypatch.setattr( + litellm, + "vector_store_index_registry", + VectorStoreIndexRegistry( + [ + LiteLLM_ManagedVectorStoreIndex( + id="index-domain-docs", + index_name="docs", + litellm_params=IndexCreateLiteLLMParams( + vector_store_index="prod-docs-v2", vector_store_name="opensearch-domain" + ), + ) + ] + ), + ) + monkeypatch.setattr( + litellm, + "vector_store_registry", + VectorStoreRegistry( + [ + { + "vector_store_id": "opensearch-domain", + "vector_store_name": "opensearch-domain", + "custom_llm_provider": "opensearch", + "litellm_params": { + "api_base": f"https://{domain}", + "aws_web_identity_token": "oidc/google/litellm", + "aws_role_name": "arn:aws:iam::123456789012:role/litellm-role", + "aws_session_name": f"opensearch-index-{method}", + "aws_region_name": "us-east-1", + }, + } + ] + ), + ) + sts: Final = MagicMock() + sts.assume_role_with_web_identity.return_value = { + "Credentials": { + "AccessKeyId": "k", + "SecretAccessKey": "s", + "SessionToken": "t", + "Expiration": datetime(2099, 1, 1, tzinfo=timezone.utc), + }, + "PackedPolicySize": 0, + } + _act_as(monkeypatch, _opensearch_caller("read", "write")) + with ( + patch("boto3.client", return_value=sts), + patch("litellm.llms.bedrock.base_aws_llm.get_secret", return_value="oidc-jwt-token"), + respx.mock(assert_all_called=True) as upstream, + ): + upstream.route(method=method, host=domain).mock(return_value=httpx.Response(200, json={})) + response = opensearch_client.request( + method, f"/opensearch/docs/{api}", content=b"{}", headers={"Content-Type": JSON_TYPE} + ) + + policy: Final = json.loads(sts.assume_role_with_web_identity.call_args.kwargs["Policy"]) + granted: Final = frozenset(action for statement in policy["Statement"] for action in statement["Action"]) + assert (response.status_code, action in granted) == (200, True) diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 45abb880f92..31b7b499845 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -44,6 +44,8 @@ from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitErro from litellm.proxy.hooks.parallel_request_limiter import PROXY_MaxParallelRequestsHandler from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup +from litellm.types.proxy.auth.user_api_key_auth import COMMITTED_BUDGET_SPEND_PARAM +from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import router as llm_passthrough_router from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler import ( VertexPassthroughLoggingHandler, @@ -108,6 +110,7 @@ PROTOCOL_CONSTRAINED_PASS_THROUGH_ROUTES: Final = MappingProxyType( "/laya/v1/systemone": frozenset({"POST"}), "/bespoke/v1/systemone": frozenset({"POST"}), "/eu_jev/v1/systemone": frozenset({"POST"}), + "/opensearch/{endpoint:path}": frozenset({"GET", "HEAD", "POST", "PUT", "DELETE", "PATCH"}), } ) @@ -446,6 +449,39 @@ async def test_non_streaming_http_request_handler_multipart_with_non_empty_parse assert call_args["files"] == [("file", ("test.txt", file_content, "text/plain"))] +@pytest.mark.parametrize("existing_api_base", [None, "http://existing.example/search"]) +def test_build_passthrough_failure_request_payload_carries_masked_api_base(existing_api_base: str | None) -> None: + logging_obj: Final = LiteLLMLoggingObj( + model="passthrough", + messages=[], + stream=False, + call_type=CallTypes.pass_through.value, + start_time=datetime(2026, 1, 1), + litellm_call_id="test-call-id", + function_id="test-function-id", + ) + logging_obj.pre_call( + input={}, + api_key="", + additional_args={ + "complete_input_dict": {}, + "api_base": "http://upstream.example/search?api-version=1", + "headers": {}, + }, + ) + params: Final = {"metadata": {}, **({"api_base": existing_api_base} if existing_api_base else {})} + payload: Final = _build_passthrough_failure_request_payload( + parsed_body={}, + kwargs={"litellm_params": params}, + logging_obj=logging_obj, + custom_llm_provider=None, + ) + assert payload["litellm_params"]["api_base"] == ( + existing_api_base or logging_obj.model_call_details["litellm_params"]["api_base"] + ) + assert params == {"metadata": {}, **({"api_base": existing_api_base} if existing_api_base else {})} + + @pytest.mark.asyncio async def test_pass_through_request_failure_handler(): """ @@ -465,6 +501,7 @@ async def test_pass_through_request_failure_handler(): assert "Request failed" in (failure.traceback_str or "") assert failure.request_data["contents"] == "not a list" assert failure.request_data["model"] == "" + assert failure.request_data["litellm_params"]["api_base"] == "http://test.com" @pytest.mark.asyncio @@ -9692,6 +9729,31 @@ async def _track_cost_for_passthrough_kwargs(kwargs: dict) -> AsyncMock: return increment_spend_counters +@pytest.mark.asyncio +async def test_passthrough_charge_moves_the_committed_spend_of_the_entities_auth_named_budgeted(): + """A forged body value stays out, and loggers that dump the row never get the auth object or the spend it read.""" + counter_keys: Final = EntitySpendCounterKeys.of( + token="hashed-token", team_id="t1", user_id="u1", org_id=None, end_user_id=None + ) + user_api_key_dict: Final = UserAPIKeyAuth( + api_key="hashed-token", + user_id="u1", + committed_budget_spend=CommittedBudgetSpend(counter_keys=counter_keys, key=0.123456), + metadata={"logging": [{"callback_name": "langfuse"}], "kept": "yes"}, + ) + + kwargs: Final = _passthrough_kwargs_for_reservation( + user_api_key_dict, + parsed_body={"litellm_metadata": {"user_api_key_auth": {"committed_budget_spend": None}}}, + ) + increment_spend_counters: Final = await _track_cost_for_passthrough_kwargs(kwargs) + + assert increment_spend_counters.await_args.kwargs["committed_spend_scope"] == counter_keys + assert "user_api_key_auth" not in kwargs["litellm_params"]["metadata"] + assert "committed_budget_spend" not in repr(kwargs["litellm_params"]["metadata"]) + assert "0.123456" not in json.dumps(kwargs, default=str) + + @pytest.mark.asyncio async def test_passthrough_success_reconciles_budget_reservation(): """ @@ -11134,6 +11196,9 @@ def _split_pass_through_body( GEMINI_BODY: Final = '{"contents": [{"parts": [{"text": "hi"}]}], "generationConfig": {"temperature": 0}}' +PROXY_SET_PARAMS: Final = ("metadata", "proxy_server_request", COMMITTED_BUDGET_SPEND_PARAM) + + def _metadata_of(split: _PassThroughSplit) -> Mapping[str, object]: return MappingProxyType(_LITELLM_PARAMS.validate_python(split.litellm_params["metadata"])) @@ -11145,7 +11210,7 @@ def test_passthrough_moves_every_litellm_owned_key_from_the_forwarded_body_into_ ) assert frozenset(split.litellm_params) == frozenset( - ("ttl", "num_retries", "litellm_trace_id", "metadata", "proxy_server_request") + ("ttl", "num_retries", "litellm_trace_id", *PROXY_SET_PARAMS) ) assert tuple(split.litellm_params[k] for k in ("ttl", "num_retries", "litellm_trace_id")) == (30, 2, "trace-a") assert split.forwarded_body == json.loads(GEMINI_BODY) @@ -11179,7 +11244,7 @@ PROXY_STAMPED_NAMES: Final = frozenset( def test_passthrough_keeps_each_registered_litellm_owned_name_out_of_the_forwarded_body(name: str) -> None: split: Final = _split_pass_through_body(json.dumps({name: "owned", **json.loads(GEMINI_BODY)})) - assert frozenset(split.litellm_params) == frozenset((name, "metadata", "proxy_server_request")) + assert frozenset(split.litellm_params) == frozenset((name, *PROXY_SET_PARAMS)) assert split.litellm_params[name] == "owned" assert split.forwarded_body == json.loads(GEMINI_BODY) @@ -11223,7 +11288,7 @@ def test_passthrough_keeps_a_client_model_info_object_for_keys_that_allow_client def test_passthrough_drops_a_client_supplied_proxy_stamped_name(name: str) -> None: split: Final = _split_pass_through_body(json.dumps({name: {"forged": "by-client"}, **json.loads(GEMINI_BODY)})) - assert frozenset(split.litellm_params) == frozenset(("metadata", "proxy_server_request")) + assert frozenset(split.litellm_params) == frozenset(PROXY_SET_PARAMS) assert split.litellm_params["proxy_server_request"] != {"forged": "by-client"}, split.litellm_params assert split.forwarded_body == json.loads(GEMINI_BODY) @@ -11234,7 +11299,7 @@ def test_passthrough_merges_both_metadata_carriers_from_the_body_into_one_metada ' "litellm_metadata": {"lm": "b"}, "generationConfig": {"temperature": 0}}' ) - assert frozenset(split.litellm_params) == frozenset(("ttl", "metadata", "proxy_server_request")) + assert frozenset(split.litellm_params) == frozenset(("ttl", *PROXY_SET_PARAMS)) assert _metadata_of(split) == {**_metadata_of(_split_pass_through_body(GEMINI_BODY)), "client_tag": "a", "lm": "b"} assert split.forwarded_body == json.loads(GEMINI_BODY) @@ -11273,7 +11338,7 @@ def test_passthrough_sees_a_name_appended_to_the_public_list_after_import() -> N finally: litellm.all_litellm_params.remove("registered_later") - assert frozenset(split.litellm_params) == frozenset(("registered_later", "metadata", "proxy_server_request")) + assert frozenset(split.litellm_params) == frozenset(("registered_later", *PROXY_SET_PARAMS)) assert split.forwarded_body == {"contents": [{"parts": [{"text": "hi"}]}]} @@ -11282,7 +11347,7 @@ def test_passthrough_sees_the_public_list_rebound_after_import(monkeypatch: pyte split: Final = _split_pass_through_body(LATE_REGISTERED_BODY) - assert frozenset(split.litellm_params) == frozenset(("registered_later", "metadata", "proxy_server_request")) + assert frozenset(split.litellm_params) == frozenset(("registered_later", *PROXY_SET_PARAMS)) assert split.forwarded_body == {"contents": [{"parts": [{"text": "hi"}]}]} diff --git a/tests/unit/proxy/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index 19db80b8930..d04021758ef 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/tests/unit/proxy/proxy_server/test_spend_counters.py @@ -38,10 +38,13 @@ from litellm.models.end_user import LiteLLM_EndUserTable from litellm.models.tag import LiteLLM_TagTable from litellm.models.team import LiteLLM_TeamTable, LiteLLM_TeamTableCachedObj from litellm.models.user import LiteLLM_UserTable +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_local from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, end_user_cache_key, tag_cache_key from litellm.proxy.management_endpoints.team_endpoints import refresh_cached_team +from litellm.proxy.spend_tracking.committed_spend import committed_spend_key +from litellm.types.proxy.carried_budget_state import EntitySpendCounterKeys from .conftest import normalize from litellm.rust_bridge._native import NativeRedisStoreError @@ -679,7 +682,7 @@ async def test_increment_spend_counters_pipelines_all_scopes_in_one_redis_call( "spend:org:org-1", } spend = _spend_counters(monkeypatch, redis=dict.fromkeys(expected_keys, 1.0)) - monkeypatch.setattr(ps, "user_api_key_cache", ps.DualCache()) + monkeypatch.setattr(ps, "user_api_key_cache", UserApiKeyCache()) monkeypatch.setattr(ps, "prisma_client", None) await ps.increment_spend_counters( @@ -697,6 +700,97 @@ async def test_increment_spend_counters_pipelines_all_scopes_in_one_redis_call( assert spend.memory == dict.fromkeys(expected_keys, 1.5) +class _VerificationTokenTable: + def __init__(self, spend: float) -> None: + self._spend: Final = spend + self.lookups: list[dict[str, str]] = [] + + async def find_unique(self, where: dict[str, str]) -> MagicMock: + self.lookups.append(where) + return MagicMock(token=where["token"], spend=self._spend) + + +def _prisma_with_key_spend(spend: float) -> MagicMock: + prisma: Final = MagicMock() + prisma.db.litellm_verificationtoken = _VerificationTokenTable(spend) + return prisma + + +@pytest.mark.asyncio +async def test_committed_twin_takes_the_charge_but_never_a_budget_hold(monkeypatch): + """The key counter holds the DB spend 10.0, this request's 2.0 reservation and another request's 1.0 hold.""" + spend: Final = _spend_counters(monkeypatch, redis=None, memory={"spend:key:hashed": 13.0}) + monkeypatch.setattr(ps, "user_api_key_cache", UserApiKeyCache()) + monkeypatch.setattr(ps, "prisma_client", _prisma_with_key_spend(10.0)) + reservation: Final = {"reserved_cost": 2.0, "entries": [{"counter_key": "spend:key:hashed"}], "finalized": False} + scope: Final = EntitySpendCounterKeys(key="spend:key:hashed") + + await ps.increment_spend_counters( + token="hashed", + team_id=None, + user_id=None, + response_cost=0.5, + budget_reservation=reservation, + committed_spend_scope=scope, + ) + await ps.increment_spend_counters( + token="hashed", team_id=None, user_id=None, response_cost=0.25, committed_spend_scope=scope + ) + + assert spend.memory == pytest.approx({"spend:key:hashed": 11.75, committed_spend_key("spend:key:hashed"): 10.75}) + + +@pytest.mark.asyncio +async def test_a_charge_outside_the_committed_scope_keeps_no_twin_and_reads_no_stored_spend(monkeypatch): + spend: Final = _spend_counters(monkeypatch, redis=None, memory={"spend:key:hashed": 10.0}) + prisma: Final = _prisma_with_key_spend(10.0) + key_cache: Final = UserApiKeyCache() + key_cache.set_cache(key="hashed", value=UserAPIKeyAuth(token="hashed")) + monkeypatch.setattr(ps, "user_api_key_cache", key_cache) + monkeypatch.setattr(ps, "prisma_client", prisma) + + await ps.increment_spend_counters(token="hashed", team_id=None, user_id=None, response_cost=0.5) + + assert spend.memory == pytest.approx({"spend:key:hashed": 10.5}) + assert prisma.db.litellm_verificationtoken.lookups == [] + + +@pytest.mark.asyncio +async def test_only_the_entities_in_the_committed_scope_take_the_charge_on_their_twin(monkeypatch): + spend: Final = _spend_counters(monkeypatch, redis=None, memory={"spend:key:hashed": 10.0, "spend:team:t1": 10.0}) + monkeypatch.setattr(ps, "user_api_key_cache", UserApiKeyCache()) + monkeypatch.setattr(ps, "prisma_client", _prisma_with_key_spend(10.0)) + + await ps.increment_spend_counters( + token="hashed", + team_id="t1", + user_id=None, + response_cost=0.5, + committed_spend_scope=EntitySpendCounterKeys(key="spend:key:hashed"), + ) + + assert spend.memory.get(committed_spend_key("spend:key:hashed")) == pytest.approx(10.5) + assert committed_spend_key("spend:team:t1") not in spend.memory + + +@pytest.mark.asyncio +async def test_a_budgeted_team_whose_id_names_a_window_keeps_its_charge_and_gets_no_twin(monkeypatch): + team_counter_key: Final = "spend:team:t1:window:1d" + spend: Final = _spend_counters(monkeypatch, redis=None, memory={team_counter_key: 1.0}) + monkeypatch.setattr(ps, "user_api_key_cache", UserApiKeyCache()) + monkeypatch.setattr(ps, "prisma_client", MagicMock()) + + await ps.increment_spend_counters( + token=None, + team_id="t1:window:1d", + user_id=None, + response_cost=0.5, + committed_spend_scope=EntitySpendCounterKeys(team=team_counter_key), + ) + + assert spend.memory == pytest.approx({team_counter_key: 1.5}) + + _PENDING_SPEND_COUNTER_KEYS: tuple[str, ...] = ( "spend:key:hashed", "spend:team:team-1", @@ -717,7 +811,7 @@ def _warm_counters_behind_a_failing_pipeline(monkeypatch, failure: str) -> FakeS fail_pipelines=True, failure=failure, ) - monkeypatch.setattr(ps, "user_api_key_cache", ps.DualCache()) + monkeypatch.setattr(ps, "user_api_key_cache", UserApiKeyCache()) monkeypatch.setattr(ps, "prisma_client", None) return spend diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 24ad068e6df..2c23305d5c9 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -3691,12 +3691,19 @@ class TestResponsesWSFirstFrameModelAuth: "standard_logging_guardrail_information": [{"guardrail_cost": 0.2}, {"guardrail_cost": 0.4}], } - @pytest.mark.asyncio @pytest.mark.parametrize( "relay_outcome", ["still_relaying", "already_returned", "returns_despite_the_cancel", "crashes_despite_the_cancel"], ) - async def test_session_limit_frees_the_slot_at_once_and_never_cuts_the_failure_row(self, relay_outcome: str): + def test_session_limit_frees_the_slot_at_once_and_never_cuts_the_failure_row(self, relay_outcome: str) -> None: + loop: Final = AutojumpClockLoop() + try: + loop.run_until_complete(self._assert_session_limit_releases_slot_before_failure_row(relay_outcome)) + finally: + loop.run_until_complete(loop.shutdown_asyncgens()) + loop.close() + + async def _assert_session_limit_releases_slot_before_failure_row(self, relay_outcome: str) -> None: from litellm.exceptions import RateLimitError from litellm.proxy.response_api_endpoints.endpoints import responses_websocket_endpoint diff --git a/tests/unit/proxy/spend_tracking/test_carried_budget_state.py b/tests/unit/proxy/spend_tracking/test_carried_budget_state.py index 0bdf43b396c..75b9e422330 100644 --- a/tests/unit/proxy/spend_tracking/test_carried_budget_state.py +++ b/tests/unit/proxy/spend_tracking/test_carried_budget_state.py @@ -2,6 +2,9 @@ from datetime import datetime, timezone +import pytest + +from litellm.models.end_user import LiteLLM_EndUserTable from litellm.proxy._types import ( LiteLLM_BudgetTable, LiteLLM_OrganizationTable, @@ -11,10 +14,13 @@ from litellm.proxy._types import ( ) from litellm.proxy.spend_tracking.carried_budget_state import ( carried_budget_metadata, + carry_committed_budget_spend, carry_organization_budget_state, carry_team_and_user_budget_state, ) from litellm.types.proxy.carried_budget_state import ( + CommittedBudgetSpend, + EntitySpendCounterKeys, KeyBudgetSnapshot, OrgBudgetSnapshot, TeamBudgetSnapshot, @@ -123,10 +129,105 @@ def test_snapshots_never_reach_the_serialized_token(): user_object=LiteLLM_UserTable(user_id="u1", user_alias="Alice"), ) token.org_budget_snapshot = OrgBudgetSnapshot(spend=1.0, max_budget=2.0) + metadata_without_committed_spend = dict(carried_budget_metadata(token)) + token.committed_budget_spend = CommittedBudgetSpend(key=1.0, team=2.0, user=3.0, org=4.0, end_user=5.0) dumped = token.model_dump() assert "team_budget_snapshot" not in dumped assert "user_budget_snapshot" not in dumped assert "org_budget_snapshot" not in dumped + assert "committed_budget_spend" not in dumped + assert dict(carried_budget_metadata(token)) == metadata_without_committed_spend assert UserAPIKeyAuth(**dumped).team_budget_snapshot is None + + +def _committed_reader(committed: CommittedBudgetSpend): + reads: list[EntitySpendCounterKeys] = [] + + async def read_committed(counter_keys: EntitySpendCounterKeys) -> CommittedBudgetSpend: + reads.append(counter_keys) + return committed + + return read_committed, reads + + +@pytest.mark.asyncio +async def test_committed_spend_is_read_once_for_every_budgeted_entity(): + token = UserAPIKeyAuth(token="hashed", team_id="t1", user_id="u1", org_id="o1", spend=1.0, max_budget=10.0) + token.org_budget_snapshot = OrgBudgetSnapshot(spend=40.0, max_budget=500.0) + committed = CommittedBudgetSpend(key=1.3, team=20.3, user=5.3, org=40.3, end_user=0.8) + read_committed, reads = _committed_reader(committed) + + await carry_committed_budget_spend( + valid_token=token, + team_object=LiteLLM_TeamTable(team_id="t1", max_budget=300.0, spend=20.0), + user_object=LiteLLM_UserTable(user_id="u1", max_budget=50.0, spend=5.0), + end_user_object=LiteLLM_EndUserTable( + user_id="cust-1", blocked=False, spend=0.5, litellm_budget_table=LiteLLM_BudgetTable(max_budget=3.0) + ), + read_committed=read_committed, + ) + + assert reads == [ + EntitySpendCounterKeys( + key="spend:key:hashed", + team="spend:team:t1", + user="spend:user:u1", + org="spend:org:o1", + end_user="spend:end_user:cust-1", + ) + ] + assert token.committed_budget_spend == committed + assert (token.spend, token.team_spend) == (1.0, None) + + +@pytest.mark.asyncio +async def test_route_without_a_budget_check_carries_the_counters_it_charges_and_reads_nothing(): + token = UserAPIKeyAuth(token="hashed", max_budget=10.0) + + await carry_committed_budget_spend( + valid_token=token, team_object=None, user_object=None, end_user_object=None, read_committed=None + ) + + assert token.committed_budget_spend == CommittedBudgetSpend( + counter_keys=EntitySpendCounterKeys(key="spend:key:hashed") + ) + + +@pytest.mark.asyncio +async def test_committed_spend_skips_entities_without_a_max_budget(): + token = UserAPIKeyAuth(token="hashed", team_id="t1", user_id="u1", org_id="o1", spend=1.0) + token.org_budget_snapshot = OrgBudgetSnapshot(spend=40.0, max_budget=None) + read_committed, reads = _committed_reader(CommittedBudgetSpend()) + + await carry_committed_budget_spend( + valid_token=token, + team_object=LiteLLM_TeamTable(team_id="t1", spend=20.0), + user_object=LiteLLM_UserTable(user_id="u1", spend=5.0), + end_user_object=LiteLLM_EndUserTable( + user_id="cust-1", blocked=False, spend=0.5, litellm_budget_table=LiteLLM_BudgetTable(max_budget=None) + ), + read_committed=read_committed, + ) + + assert reads == [EntitySpendCounterKeys()] + assert token.committed_budget_spend == CommittedBudgetSpend() + + +@pytest.mark.asyncio +async def test_committed_spend_reads_the_team_org_when_the_key_has_none(): + token = UserAPIKeyAuth(token="hashed", team_id="t1") + token.org_budget_snapshot = OrgBudgetSnapshot(spend=40.0, max_budget=500.0) + read_committed, reads = _committed_reader(CommittedBudgetSpend(org=41.0)) + + await carry_committed_budget_spend( + valid_token=token, + team_object=LiteLLM_TeamTable(team_id="t1", organization_id="team-org"), + user_object=None, + end_user_object=None, + read_committed=read_committed, + ) + + assert reads == [EntitySpendCounterKeys(org="spend:org:team-org")] + assert token.committed_budget_spend == CommittedBudgetSpend(org=41.0) diff --git a/tests/unit/proxy/spend_tracking/test_committed_spend.py b/tests/unit/proxy/spend_tracking/test_committed_spend.py new file mode 100644 index 00000000000..1ad5a2e7ed4 --- /dev/null +++ b/tests/unit/proxy/spend_tracking/test_committed_spend.py @@ -0,0 +1,274 @@ +"""Committed spend twins hold only charged cost, so the remaining budget gauges never count a budget hold.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Callable, Mapping +from types import MappingProxyType, SimpleNamespace +from typing import Final +from unittest.mock import MagicMock +from uuid import uuid4 + +import pytest + +import litellm +from litellm.integrations.prometheus import PrometheusLogger +from litellm.proxy.spend_tracking.committed_spend import ( + bind_committed_admission_counters, + drop_committed_spend, + read_admission_committed_spend, + read_committed_budget_spend, + seed_committed_spend, +) +from litellm.proxy.spend_tracking.spend_counter_batch import ( + active_spend_counter_batch, + read_batched_spend_counter, + spend_counter_batch_scope, +) +from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys +from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters + + +def _no_side_effect() -> None: + return None + + +class _FindUniqueTable: + def __init__(self, row: SimpleNamespace | None, on_read: Callable[[], None]) -> None: + self._row: Final = row + self._on_read: Final = on_read + self.reads = 0 + + async def find_unique(self, where: dict[str, str]) -> SimpleNamespace | None: + self.reads += 1 + await asyncio.sleep(0) + self._on_read() + return self._row + + +def _prisma_with_spend(spend: float | None, on_read: Callable[[], None] = _no_side_effect) -> SimpleNamespace: + row: Final = None if spend is None else SimpleNamespace(spend=spend) + return SimpleNamespace( + db=SimpleNamespace( + litellm_verificationtoken=_FindUniqueTable(row, on_read), + litellm_teamtable=_FindUniqueTable(row, on_read), + litellm_endusertable=_FindUniqueTable(row, on_read), + ) + ) + + +@pytest.fixture(params=[False, True], ids=["in_memory", "redis"]) +def has_redis(request: pytest.FixtureRequest) -> bool: + return request.param + + +def _spend_counters(has_redis: bool, counters: Mapping[str, float] = MappingProxyType({})) -> FakeSpendCounters: + if has_redis: + return FakeSpendCounters(redis=dict(counters)) + return FakeSpendCounters(redis=None, memory=dict(counters)) + + +@pytest.mark.asyncio +async def test_read_returns_each_committed_twin_and_none_for_one_never_charged(has_redis: bool): + spend: Final = _spend_counters( + has_redis, + {"spend:committed:key:hashed": 1.5, "spend:committed:end_user:cust-1": 0.25, "spend:team:t1": 99.0}, + ) + counter_keys: Final = EntitySpendCounterKeys.of( + token="hashed", team_id="t1", user_id=None, org_id=None, end_user_id="cust-1" + ) + + committed: Final = await read_committed_budget_spend(spend.store(), counter_keys) + + assert committed == CommittedBudgetSpend(counter_keys=counter_keys, key=1.5, end_user=0.25) + + +@pytest.mark.asyncio +async def test_a_twin_holding_garbage_reads_as_unknown_without_hiding_the_other_twins(): + spend: Final = FakeSpendCounters( + redis={"spend:committed:key:hashed": "not-a-number", "spend:committed:team:t1": 4.0} + ) + counter_keys: Final = EntitySpendCounterKeys.of( + token="hashed", team_id="t1", user_id=None, org_id=None, end_user_id=None + ) + + committed: Final = await read_committed_budget_spend(spend.store(), counter_keys) + + assert committed == CommittedBudgetSpend(counter_keys=counter_keys, team=4.0) + + +@pytest.mark.parametrize( + ("counter_keys", "expected_spend"), + [ + (EntitySpendCounterKeys(key="spend:key:hashed"), {"key": 12.0}), + (EntitySpendCounterKeys(end_user="spend:end_user:cust-1"), {"end_user": 12.0}), + ], +) +@pytest.mark.asyncio +async def test_seed_writes_the_db_spend_when_the_twin_is_absent( + has_redis: bool, counter_keys: EntitySpendCounterKeys, expected_spend: dict[str, float] +): + store: Final = _spend_counters(has_redis).store() + (counter_key,) = counter_keys.tracked() + + seeded: Final = await seed_committed_spend(store, _prisma_with_spend(12.0), counter_key) + + assert seeded is True + assert await read_committed_budget_spend(store, counter_keys) == CommittedBudgetSpend( + counter_keys=counter_keys, **expected_spend + ) + + +@pytest.mark.asyncio +async def test_concurrent_charges_on_a_cold_twin_read_the_db_once(has_redis: bool): + store: Final = _spend_counters(has_redis).store() + prisma_client: Final = _prisma_with_spend(12.0) + counter_key: Final = f"spend:key:{uuid4()}" + + seeded: Final = await asyncio.gather(*(seed_committed_spend(store, prisma_client, counter_key) for _ in range(5))) + + assert seeded == [True] * 5 + assert prisma_client.db.litellm_verificationtoken.reads == 1 + + +@pytest.mark.asyncio +async def test_seed_reports_no_twin_when_the_twin_cannot_be_read(): + async def unreachable(counter: str) -> None: + raise ConnectionError("redis unreachable") + + spend: Final = FakeSpendCounters(read_hook=unreachable) + + assert await seed_committed_spend(spend.store(), _prisma_with_spend(12.0), "spend:team:t1") is False + + +@pytest.mark.asyncio +async def test_seed_never_lowers_a_twin_another_replica_already_advanced(has_redis: bool): + store: Final = _spend_counters(has_redis, {"spend:committed:team:t1": 20.0}).store() + + seeded: Final = await seed_committed_spend(store, _prisma_with_spend(12.0), "spend:team:t1") + + committed: Final = await read_committed_budget_spend(store, EntitySpendCounterKeys(team="spend:team:t1")) + assert seeded is True + assert committed.team == 20.0 + + +@pytest.mark.asyncio +async def test_seed_keeps_the_twin_another_worker_wrote_while_this_one_read_the_db(): + spend: Final = FakeSpendCounters(redis={}) + + def another_worker_seeds_and_charges() -> None: + assert spend.redis is not None + spend.redis["spend:committed:team:t1"] = 20.0 + + seeded: Final = await seed_committed_spend( + spend.store(), _prisma_with_spend(12.0, on_read=another_worker_seeds_and_charges), "spend:team:t1" + ) + + committed: Final = await read_committed_budget_spend(spend.store(), EntitySpendCounterKeys(team="spend:team:t1")) + assert seeded is True + assert committed.team == 20.0 + + +@pytest.mark.asyncio +async def test_read_without_any_counter_key_sends_redis_nothing(): + spend: Final = FakeSpendCounters(redis={}) + + committed: Final = await read_committed_budget_spend(spend.store(), EntitySpendCounterKeys()) + + assert committed == CommittedBudgetSpend() + assert spend.commands == [] + + +@pytest.mark.asyncio +async def test_seed_reports_no_twin_when_redis_rejects_the_write(): + spend: Final = FakeSpendCounters(redis={}, fail=True) + + seeded: Final = await seed_committed_spend(spend.store(), _prisma_with_spend(12.0), "spend:team:t1") + + committed: Final = await read_committed_budget_spend(spend.store(), EntitySpendCounterKeys(team="spend:team:t1")) + assert seeded is False + assert committed.team is None + + +@pytest.mark.asyncio +async def test_seed_creates_no_twin_when_the_db_has_no_spend_for_the_entity(has_redis: bool): + store: Final = _spend_counters(has_redis).store() + + seeded: Final = await seed_committed_spend(store, _prisma_with_spend(None), "spend:key:hashed") + + committed: Final = await read_committed_budget_spend(store, EntitySpendCounterKeys(key="spend:key:hashed")) + assert seeded is False + assert committed.key is None + + +@pytest.mark.asyncio +async def test_admission_read_rides_the_admission_mget_with_the_enforcement_counters( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "callbacks", [MagicMock(spec=PrometheusLogger)]) + spend: Final = FakeSpendCounters(redis={"spend:committed:key:hashed": 1.5}) + store: Final = spend.store() + token: Final = UserAPIKeyAuth(token="hashed", team_id="t1") + counter_keys: Final = EntitySpendCounterKeys(key="spend:key:hashed", team="spend:team:t1") + + with spend_counter_batch_scope(store, counter_keys=frozenset(("spend:key:hashed",))): + bind_committed_admission_counters(token, end_user_id=None) + await read_batched_spend_counter("spend:key:hashed") + committed: Final = await read_admission_committed_spend(store, counter_keys) + + assert committed == CommittedBudgetSpend(counter_keys=counter_keys, key=1.5) + assert len(spend.commands) == 1 + assert {"spend:committed:key:hashed", "spend:committed:team:t1"} <= set(spend.commands[0].split()[1:]) + + +@pytest.mark.asyncio +async def test_admission_read_outside_a_batch_leaves_redis_alone_and_the_spend_unknown(): + spend: Final = FakeSpendCounters(redis={"spend:committed:key:hashed": 1.5}) + counter_keys: Final = EntitySpendCounterKeys(key="spend:key:hashed") + + committed: Final = await read_admission_committed_spend(spend.store(), counter_keys) + + assert committed == CommittedBudgetSpend(counter_keys=counter_keys) + assert spend.commands == [] + + +@pytest.mark.asyncio +async def test_admission_read_without_redis_reads_the_in_memory_twin(): + spend: Final = FakeSpendCounters(redis=None, memory={"spend:committed:key:hashed": 1.5}) + + committed: Final = await read_admission_committed_spend( + spend.store(), EntitySpendCounterKeys(key="spend:key:hashed") + ) + + assert committed.key == 1.5 + + +def test_no_committed_counter_is_bound_without_a_prometheus_logger(): + with spend_counter_batch_scope(FakeSpendCounters(redis={}).store()): + bind_committed_admission_counters(UserAPIKeyAuth(token="hashed"), end_user_id=None) + batch: Final = active_spend_counter_batch() + assert batch is not None + bound: Final = batch.counter_keys + + assert bound == frozenset() + + +@pytest.mark.asyncio +async def test_drop_removes_the_twin_and_leaves_the_enforcement_counter(has_redis: bool): + store: Final = _spend_counters(has_redis, {"spend:committed:org:o1": 5.0, "spend:org:o1": 5.0}).store() + counter_keys: Final = EntitySpendCounterKeys(org="spend:org:o1") + + await drop_committed_spend(store, ("spend:org:o1",)) + + assert (await read_committed_budget_spend(store, counter_keys)).org is None + assert (await store.read("spend:org:o1"))[0] == 5.0 + + +@pytest.mark.asyncio +async def test_a_failed_drop_never_fails_the_write_that_already_committed(): + spend: Final = FakeSpendCounters(redis={"spend:committed:org:o1": 5.0}, failing=frozenset({"delete"})) + + await drop_committed_spend(spend.store(), ("spend:org:o1",)) + + assert spend.deleted() == frozenset({"spend:committed:org:o1"}) diff --git a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py index c3e916d581d..05667629305 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -1,12 +1,14 @@ import asyncio import collections +import contextlib import datetime import hashlib import json import re import sqlite3 -from collections.abc import Mapping, Sequence +from collections.abc import AsyncIterator, Mapping, Sequence from datetime import timezone +from types import SimpleNamespace from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch @@ -17,14 +19,17 @@ from fastapi.testclient import TestClient import litellm import litellm.proxy.proxy_server as ps +from litellm.integrations.prometheus import PrometheusLogger from litellm.proxy.auth.authorization import OwnedRows from litellm.proxy.auth.authorization_dependencies import get_log_team_lookup, load_permitted_log_team_ids from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.spend_tracking.spend_management_endpoints import global_spend_reset from tests.unit.proxy.db.fake_prisma_engine import ( ReadTargetOf, statement_timeout_error, statement_timeout_reads, ) +from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters def _default_date_range(): @@ -8336,6 +8341,109 @@ async def test_global_spend_refresh_runs_the_refresh_on_the_writer_when_a_reader ] +class _ResetTableRows: + def __init__(self) -> None: + self.update_many: Final = AsyncMock() + + +class _ResetDb: + def __init__(self) -> None: + self.litellm_verificationtoken: Final = _ResetTableRows() + self.litellm_teamtable: Final = _ResetTableRows() + + async def query_raw(self, sql_query: str, after_id: str, limit: int) -> list[dict[str, str]]: + rows: Final = ["hashed"] if "VerificationToken" in sql_query else ["t1"] + return [{"id": row} for row in rows if row > after_id][:limit] + + +class _ResetTransaction: + def __init__(self, db: _ResetDb) -> None: + self.litellm_verificationtoken: Final = db.litellm_verificationtoken + self.litellm_teamtable: Final = db.litellm_teamtable + + async def query_raw(self, sql_query: str) -> list[dict[str, str]]: + return [] + + +def _reset_prisma(db: _ResetDb) -> SimpleNamespace: + @contextlib.asynccontextmanager + async def tx() -> AsyncIterator[_ResetTransaction]: + yield _ResetTransaction(db) + + return SimpleNamespace(db=db, tx=tx) + + +def _committed_spend(monkeypatch: pytest.MonkeyPatch, twins: Sequence[str]) -> FakeSpendCounters: + spend: Final = FakeSpendCounters(redis=None, memory=dict.fromkeys(twins, 7.0)) + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_store", spend.store()) + return spend + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("registered_callbacks", "twin_after"), + [ + ([MagicMock(spec=PrometheusLogger)], None), + ([], 7.0), + ], +) +async def test_global_spend_reset_drops_the_committed_key_and_team_spend_the_gauges_read( + monkeypatch, registered_callbacks, twin_after +): + twins: Final = ("spend:committed:key:hashed", "spend:committed:team:t1") + spend: Final = _committed_spend(monkeypatch, twins) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _reset_prisma(_ResetDb())) + monkeypatch.setattr(litellm, "callbacks", registered_callbacks) + + response: Final = await global_spend_reset() + + assert response["status"] == "success" + assert [spend.memory.get(twin) for twin in twins] == [twin_after, twin_after] + + +class _PagedResetDb(_ResetDb): + def __init__(self, token_count: int) -> None: + super().__init__() + self._tokens: Final = sorted(f"token-{index}" for index in range(token_count)) + self.largest_page = 0 + + async def query_raw(self, sql_query: str, after_id: str, limit: int) -> list[dict[str, str]]: + rows: Final = self._tokens if "VerificationToken" in sql_query else [] + page: Final = [{"id": row} for row in rows if row > after_id][:limit] + self.largest_page = max(self.largest_page, len(page)) + return page + + +@pytest.mark.asyncio +@pytest.mark.parametrize("token_count", [999, 1000, 1001, 2500]) +async def test_global_spend_reset_drops_every_key_twin_reading_the_keys_in_bounded_pages(monkeypatch, token_count): + db: Final = _PagedResetDb(token_count) + twins: Final = tuple(f"spend:committed:key:token-{index}" for index in range(token_count)) + spend: Final = _committed_spend(monkeypatch, twins) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _reset_prisma(db)) + monkeypatch.setattr(litellm, "callbacks", [MagicMock(spec=PrometheusLogger)]) + + response: Final = await global_spend_reset() + + assert response["status"] == "success" + assert spend.memory == {} + assert db.largest_page <= 1000 + + +@pytest.mark.asyncio +async def test_global_spend_reset_still_succeeds_when_listing_the_committed_counters_fails(monkeypatch): + db: Final = _ResetDb() + db.query_raw = AsyncMock(side_effect=ConnectionError("db unreachable")) + _committed_spend(monkeypatch, ("spend:committed:team:t1",)) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _reset_prisma(db)) + monkeypatch.setattr(litellm, "callbacks", [MagicMock(spec=PrometheusLogger)]) + + response: Final = await global_spend_reset() + + assert response["status"] == "success" + db.litellm_teamtable.update_many.assert_awaited_once() + + def test_spend_calculate_rejection_body_is_openai_shaped(client): app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user", api_key="hashed-k" diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index af463fa19d4..e4a560adca8 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -26,6 +26,7 @@ from litellm.constants import ( SESSION_ID_GENERATED_METADATA_KEY, SESSION_ID_OMITTED_METADATA_KEY, ) +from litellm.integrations.otel.model.config import is_otel_v2_enabled from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs from litellm.litellm_core_utils.get_provider_specific_headers import ( ProviderSpecificHeaderUtils, @@ -59,6 +60,7 @@ from litellm.proxy.litellm_pre_call_utils import ( get_dynamic_logging_metadata, get_metadata_variable_name, move_guardrails_to_metadata, + resolve_tenant_otel_destinations, ) from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload from litellm.types.litellm_params import AGENTIC_LOOP_KWARG_NAMES @@ -2779,7 +2781,7 @@ def test_key_dynamic_logging_settings(): team_metadata={}, ) result = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(key_with_arize) - assert result == [{"callback_name": "arize", "callback_type": "success"}] + assert result == ({"callback_name": "arize", "callback_type": "success"},) # Test with langfuse logging key_with_langfuse = UserAPIKeyAuth( @@ -2788,7 +2790,7 @@ def test_key_dynamic_logging_settings(): team_metadata={}, ) result = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(key_with_langfuse) - assert result == [{"callback_name": "langfuse", "callback_type": "success"}] + assert result == ({"callback_name": "langfuse", "callback_type": "success"},) # Test with no logging metadata key_without_logging = UserAPIKeyAuth(api_key="test-key", metadata={}, team_metadata={}) @@ -2912,6 +2914,22 @@ def test_get_dynamic_logging_metadata_with_arize_team_logging(): assert result.callback_vars["arize_space_id"] == "test_arize_space_id" +@pytest.mark.parametrize( + "logging_value", + [ + pytest.param({"callback_name": "langsmith"}, id="object"), + pytest.param("langsmith", id="string"), + pytest.param(["langsmith"], id="list-with-string"), + pytest.param([None], id="list-with-null"), + ], +) +def test_non_object_team_logging_metadata_fails_as_a_server_error_not_a_bad_request(logging_value: object) -> None: + user_api_key_dict: Final = UserAPIKeyAuth(api_key="test-key", metadata={}, team_metadata={"logging": logging_value}) + + with pytest.raises(TypeError): + get_dynamic_logging_metadata(user_api_key_dict=user_api_key_dict, proxy_config=MagicMock()) + + def test_add_team_callback_rejects_env_reference(): with pytest.raises(PydanticValidationError) as exc_info: AddTeamCallback( @@ -3895,6 +3913,41 @@ async def test_non_child_requests_cannot_forge_a_child_tier_identity(surface: st assert updated[get_metadata_variable_name(request)]["session_id"] == "child-01" +@pytest.mark.asyncio +@pytest.mark.parametrize( + "agent_id,session_id,child_id", + [ + ( + "a7256c04dfd99c9b9", + "11111111-2222-3333-4444-555555555555", + "11111111-2222-3333-4444-555555555555:a7256c04dfd99c9b9", + ), + (None, "11111111-2222-3333-4444-555555555555", None), + ("short", "11111111-2222-3333-4444-555555555555", None), + ("agent id with spaces", "11111111-2222-3333-4444-555555555555", None), + ("a7256c04dfd99c9b9", "bad", None), + ], +) +async def test_claude_code_subagents_get_their_own_tier_identity_under_the_parent_session( + agent_id: str | None, session_id: str, child_id: str | None +) -> None: + headers: Final = { + "x-claude-code-session-id": session_id, + **({"x-claude-code-agent-id": agent_id} if agent_id is not None else {}), + } + request: Final = _make_request_mock("/v1/messages", headers) + updated: Final = await add_litellm_data_to_request( + data={"model": "test-router", "messages": []}, + request=request, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={}, + ) + metadata: Final = updated[get_metadata_variable_name(request)] + assert metadata.get(COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY) == child_id + assert COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY not in metadata.get("requester_metadata", {}) + + def test_add_litellm_metadata_from_request_headers_x_litellm_session_id_sets_chain_id(): """x-litellm-session-id sets both metadata and top-level litellm_session_id/litellm_trace_id for call chaining.""" headers = {"x-litellm-session-id": "bar"} @@ -8083,6 +8136,7 @@ async def test_add_litellm_data_to_request_refuses_a_model_config_credential_hol ("API-Key", True), ("Cookie", True), ("Proxy-Authorization", True), + ("X-Amz-Security-Token", True), ("anthropic-version", False), ("user-agent", False), ], @@ -10231,6 +10285,48 @@ def test_arize_otlp_protocol_on_a_key_logging_entry_reaches_the_destination(monk is_otel_v2_enabled.cache_clear() +def test_a_key_referencing_a_credential_routes_to_that_credential_with_its_own_excluded_services(monkeypatch): + monkeypatch.setenv("LITELLM_OTEL_V2", "true") + monkeypatch.setenv("ARIZE_HTTP_ENDPOINT", "https://arize.internal.example/v1/traces") + monkeypatch.delenv("ARIZE_API_KEY", raising=False) + monkeypatch.delenv("ARIZE_SPACE_ID", raising=False) + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="space-a", + credential_values={"arize_api_key": "k", "arize_space_id": "s"}, + credential_info={}, + ) + ], + ) + is_otel_v2_enabled.cache_clear() + try: + auth = UserAPIKeyAuth( + api_key="hashed-key", + metadata={ + "logging": [ + { + "callback_name": "arize", + "callback_type": "success", + "litellm_credential_name": "space-a", + "callback_vars": {"excluded_services": "redis,postgres"}, + } + ] + }, + team_metadata={}, + ) + + destinations = resolve_tenant_otel_destinations(auth) + + assert [(d.headers, d.excluded_db_systems) for d in destinations] == [ + ({"arize-space-id": "s", "api_key": "k"}, frozenset({"redis", "postgresql"})) + ] + finally: + is_otel_v2_enabled.cache_clear() + + class TestResolveUserProviderCredentials: """``_resolve_user_provider_credentials_for_request`` injects the calling user's stored provider tokens into secret_fields for deployments carrying a per_user_oauth credential.""" diff --git a/tests/unit/proxy/test_plugin_routes.py b/tests/unit/proxy/test_plugin_routes.py index 7d3fbda14fd..23145f13156 100644 --- a/tests/unit/proxy/test_plugin_routes.py +++ b/tests/unit/proxy/test_plugin_routes.py @@ -12,10 +12,19 @@ Covers three bugs: """ import asyncio +import base64 +import hashlib +import hmac +import json +import os +import time +from typing import Final, NoReturn from unittest.mock import MagicMock +import _hashlib import pytest import respx +from cryptography.fernet import Fernet from starlette.requests import Request import litellm @@ -25,7 +34,12 @@ from litellm.proxy._types import ( PluginConfig, UserAPIKeyAuth, ) -from litellm.proxy.plugin_routes import list_plugins, register_plugins_from_config +from litellm.proxy.plugin_routes import ( + issue_plugin_session_claim, + list_plugins, + register_plugins_from_config, + verify_plugin_session_claim, +) def _admin() -> UserAPIKeyAuth: @@ -320,3 +334,50 @@ def test_plugin_proxy_refuses_non_admin_callers_without_contacting_the_plugin( assert response.status_code == 403 assert response.body == b"Plugin proxy access requires proxy_admin role." assert not upstream.called + + +def test_plugin_session_claim_opens_under_the_documented_hmac_plugin_key(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-salt") + plugin_key: Final = base64.urlsafe_b64encode(hmac.new(b"sk-salt", b"chat-ui", hashlib.sha256).digest()) + + claim: Final[dict[str, object]] = json.loads( + Fernet(plugin_key).decrypt(issue_plugin_session_claim("chat-ui", "user-1", "proxy_admin")) + ) + + assert (claim["plugin"], claim["user_id"], claim["user_role"]) == ("chat-ui", "user-1", "proxy_admin") + + +@pytest.mark.skipif( + not hasattr(_hashlib, "__file__"), reason="this Python links OpenSSL statically, out of ctypes reach" +) +@pytest.mark.usefixtures("fips_mode") +def test_fips_mode_plugin_session_claim_opens_under_the_kbkdf_key_without_os_urandom_or_hmac_new( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def outside_the_fips_provider(*_: object) -> NoReturn: + raise AssertionError("fips_mode must not draw on os.urandom or hmac.new") + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-" + bytes(range(32)).hex()) + monkeypatch.setattr(os, "urandom", outside_the_fips_provider) + monkeypatch.setattr(hmac, "new", outside_the_fips_provider) + monkeypatch.setattr(time, "time", lambda: 1_800_000_000.0) + fixed_input: Final = ( + (1).to_bytes(4, "big") + b"litellm-plugin:session-claim:" + b"\x00" + b"chat-ui" + (256).to_bytes(4, "big") + ) + plugin_key: Final = base64.urlsafe_b64encode(hmac.digest(bytes(range(32)), fixed_input, "sha256")) + token: Final = issue_plugin_session_claim("chat-ui", "user-1", "proxy_admin") + + claim: Final[dict[str, object]] = json.loads(Fernet(plugin_key).decrypt(token)) + + assert (claim["plugin"], claim["user_id"], claim["user_role"]) == ("chat-ui", "user-1", "proxy_admin") + assert verify_plugin_session_claim("chat-ui", token) == claim + + +@pytest.mark.usefixtures("fips_mode") +def test_fips_mode_refuses_a_plugin_session_claim_under_a_salt_key_that_is_not_hex( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-salt") + + with pytest.raises(ValueError, match="LITELLM_SALT_KEY must be at least 32 bytes of hex"): + issue_plugin_session_claim("chat-ui", "user-1", "proxy_admin") diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 1cfb14e6bf9..c62e73232d0 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -918,6 +918,49 @@ async def test_initialize_scheduled_jobs_loads_credentials_only_through_add_depl scheduler.shutdown(wait=False) +@pytest.mark.asyncio +@pytest.mark.parametrize("proxy_batch_write_at", [10, 60]) +async def test_committed_spend_twins_outlive_the_slowest_spend_write_to_the_db( + monkeypatch: pytest.MonkeyPatch, proxy_batch_write_at: int +) -> None: + """A twin that expires before its charges reach the DB seeds again from a total still missing them.""" + monkeypatch.delenv("DISABLE_PRISMA_SCHEMA_UPDATE", raising=False) + monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False) + from apscheduler.schedulers.asyncio import AsyncIOScheduler + + from litellm.constants import SPEND_LOG_QUEUE_MAX_BACKOFF_SECONDS + from litellm.proxy.proxy_server import _BATCH_WRITE_JITTER_MAX_SECONDS, ProxyStartupEvent + from litellm.proxy.utils import ProxyLogging + from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters + + spend: Final = FakeSpendCounters() + mock_proxy_logging: Final = MagicMock(spec=ProxyLogging) + mock_proxy_logging.slack_alerting_instance = MagicMock() + mock_proxy_logging.db_spend_update_writer = MagicMock() + scheduler: Final = AsyncIOScheduler() + try: + with ( + patch("litellm.proxy.proxy_server.proxy_config", _mock_scheduled_proxy_config()), + patch("litellm.proxy.proxy_server.store_model_in_db", False), + patch("litellm.proxy.proxy_server.AsyncIOScheduler", return_value=scheduler), + patch("litellm.proxy.proxy_server.spend_counter_store", spend.store()), + ): + await ProxyStartupEvent.initialize_scheduled_background_jobs( + general_settings={}, + prisma_client=MagicMock(), + proxy_budget_rescheduler_min_time=1, + proxy_budget_rescheduler_max_time=2, + proxy_batch_write_at=proxy_batch_write_at, + proxy_logging_obj=mock_proxy_logging, + ) + finally: + scheduler.shutdown(wait=False) + + slowest_db_write: Final = SPEND_LOG_QUEUE_MAX_BACKOFF_SECONDS + proxy_batch_write_at + _BATCH_WRITE_JITTER_MAX_SECONDS + assert spend.committed_ttl is not None + assert spend.committed_ttl >= slowest_db_write + + @pytest.mark.asyncio async def test_periodic_reload_job_scheduled_without_store_model_in_db(monkeypatch): """ diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 2f85d1c7516..694b257b1f9 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -1080,33 +1080,32 @@ async def test_bedrock_handler_httpx_error_status_code_propagation(): response=mock_response, ) - with patch.object(handler, "validate_count_tokens_request"): - with patch.object(handler, "_get_aws_region_name", return_value="us-west-2"): - with patch.object(handler, "transform_anthropic_to_bedrock_count_tokens", return_value={}): - with patch.object( - handler, - "get_bedrock_count_tokens_endpoint", - return_value="https://example.com", - ): - with patch.object(handler, "_sign_request", return_value=({}, "{}")): - with patch("litellm.llms.bedrock.count_tokens.handler.get_async_httpx_client") as mock_client: - mock_async_client = AsyncMock() - mock_async_client.post = AsyncMock(side_effect=http_error) - mock_client.return_value = mock_async_client + with patch.object(handler, "_get_aws_region_name", return_value="us-west-2"): + with patch.object(handler, "transform_anthropic_to_bedrock_count_tokens", return_value={}): + with patch.object( + handler, + "get_bedrock_count_tokens_endpoint", + return_value="https://example.com", + ): + with patch.object(handler, "_sign_request", return_value=({}, "{}")): + with patch("litellm.llms.bedrock.count_tokens.handler.get_async_httpx_client") as mock_client: + mock_async_client = AsyncMock() + mock_async_client.post = AsyncMock(side_effect=http_error) + mock_client.return_value = mock_async_client - with pytest.raises(BedrockError) as exc_info: - await handler.handle_count_tokens_request( - request_data={ - "model": "test", - "messages": [{"role": "user", "content": "hello"}], - }, - litellm_params={}, - resolved_model="anthropic.claude-3-sonnet", - ) + with pytest.raises(BedrockError) as exc_info: + await handler.handle_count_tokens_request( + request_data={ + "model": "test", + "messages": [{"role": "user", "content": "hello"}], + }, + litellm_params={}, + resolved_model="anthropic.claude-3-sonnet", + ) - assert exc_info.value.status_code == 403 - # Message should be the raw response text - assert exc_info.value.message == "Forbidden - Invalid credentials" + assert exc_info.value.status_code == 403 + # Message should be the raw response text + assert exc_info.value.message == "Forbidden - Invalid credentials" @pytest.mark.asyncio diff --git a/tests/unit/redis/stores/spend_counters_double.py b/tests/unit/redis/stores/spend_counters_double.py index 1f3c4b94d45..c50133f7076 100644 --- a/tests/unit/redis/stores/spend_counters_double.py +++ b/tests/unit/redis/stores/spend_counters_double.py @@ -35,9 +35,22 @@ def _is_number(value: object) -> bool: return value is None or (isinstance(value, (int, float)) and not isinstance(value, bool)) +_TWIN_SCOPES: Final = frozenset({"key", "team", "user", "org", "end_user"}) +_WINDOWED_SCOPES: Final = frozenset({"key", "team"}) + + +def _native_store_takes_twin(twin: str) -> bool: + scope, _, counter_id = twin.partition(":") + _, separator, window = counter_id.rpartition(":window:") + return scope in _TWIN_SCOPES and not (scope in _WINDOWED_SCOPES and separator and window) + + def _check(counter: str) -> None: if not counter.startswith("spend:"): raise ValueError(f"{counter!r} is not a spend counter") + twin: Final = counter.removeprefix("spend:committed:") + if twin != counter and not _native_store_takes_twin(twin): + raise ValueError(f"{counter!r} is not a spend counter") @dataclass @@ -53,6 +66,11 @@ class DeclaredRead: return self.owner.mget(self.counters) +@dataclass +class FakeSettlement: + settled: bool = False + + @dataclass class FakeSpendCounters: """``redis`` is None for a worker without Redis, and ``redis_default`` answers reads of a key it lacks. ``fail`` @@ -68,11 +86,12 @@ class FakeSpendCounters: failure: str = "transport" redis_default: float | None = None redis_ttl: float = 60 + committed_ttl: float | None = None batch_reads: bool = True read_hook: Callable[[str], Awaitable[None]] | None = None on_delete: Callable[[tuple[str, ...]], None] | None = None deletes: list[tuple[str, ...]] = field(default_factory=list) - post_call: list[tuple[object, tuple[SpendIncrement, ...]]] = field(default_factory=list) + post_call: list[tuple[object, tuple[SpendIncrement, ...], FakeSettlement]] = field(default_factory=list) declared: list[DeclaredRead] = field(default_factory=list) def store(self) -> SpendCounterStore: @@ -86,6 +105,9 @@ class FakeSpendCounters: def redis_ttl_seconds(self) -> float | None: return None if self.redis is None else self.redis_ttl + def set_committed_ttl(self, ttl_seconds: float, /) -> None: + self.committed_ttl = ttl_seconds + def attach(self, client: NativeRedisClient, /) -> bool: if self.redis is not None: return False @@ -139,17 +161,19 @@ class FakeSpendCounters: self.commands.append(f"PIPELINE {' '.join(counter for counter, _ in increments)}") return _done([self._incr(counter, amount) for counter, amount in increments]) - def increment_many_post_call(self, request: object, increments: Sequence[SpendIncrement]) -> bool: + def increment_many_post_call(self, request: object, increments: Sequence[SpendIncrement]) -> FakeSettlement | None: if self.redis is None: - return False - self.post_call.append((request, tuple(increments))) - return True + return None + settlement: Final = FakeSettlement() + self.post_call.append((request, tuple(increments), settlement)) + return settlement def settle_post_call(self) -> None: """What the native task does once the post-call batch goes out.""" pending: Final = tuple(self.post_call) self.post_call.clear() - for _, increments in pending: + for _, increments, settlement in pending: + settlement.settled = True if self.fail: for counter, _ in increments: self.memory.pop(counter, None) diff --git a/tests/unit/redis/stores/test_spend_counters.py b/tests/unit/redis/stores/test_spend_counters.py index c54dc6cc81e..1d39cffc933 100644 --- a/tests/unit/redis/stores/test_spend_counters.py +++ b/tests/unit/redis/stores/test_spend_counters.py @@ -10,6 +10,7 @@ from litellm.redis.stores.connection import RedisStores, connect from litellm.redis.stores.spend_counters import ( SPEND_COUNTERS_TARGET, SpendCounterStore, + committed_counter, end_user_counter, is_unavailable, key_counter, @@ -28,7 +29,7 @@ from litellm.redis.store_settings import RedisStoreSettings from litellm.rust_bridge._native import NativeRedisClient, NativeRedisStoreError, NativeSpendCounters, SpendRead from tests.unit.caching.test_redis_batch import FakeNatives, fake_send -from .spend_counters_double import FakeSpendCounters +from .spend_counters_double import FakeSettlement, FakeSpendCounters _KEYS: Final = ( ( @@ -46,6 +47,11 @@ _KEYS: Final = ( (project_counter("proj-1"), "spend:project:proj-1"), (model_access_group_counter("gpt-group"), "spend:model_access_group:gpt-group"), (shadow_eval_counter("job-7"), "spend:shadow_eval:job-7"), + (committed_counter(key_counter("88dc28d0")), "spend:committed:key:88dc28d0"), + (committed_counter(team_counter("team-a")), "spend:committed:team:team-a"), + (committed_counter(user_counter("user@example.com")), "spend:committed:user:user@example.com"), + (committed_counter(org_counter("org-1")), "spend:committed:org:org-1"), + (committed_counter(end_user_counter("customer é 1")), "spend:committed:end_user:customer é 1"), ) @@ -81,7 +87,17 @@ def test_a_read_decodes_what_redis_cache_decodes(wire: str, spend: float) -> Non assert type(decoded) is type(spend) -@pytest.mark.parametrize("key", ["spend_db_floor:spend:key:k", "virtual_key_spend:k:gpt:1d", "spend:unknown:k"]) +@pytest.mark.parametrize( + "key", + [ + "spend_db_floor:spend:key:k", + "virtual_key_spend:k:gpt:1d", + "spend:unknown:k", + "spend:committed:key:k:window:1d", + "spend:committed:tag:prod", + "spend:committed:committed:key:k", + ], +) def test_the_native_store_refuses_keys_outside_the_spend_counter_keyspace(key: str) -> None: with pytest.raises(ValueError, match="a spend counter is a spend:\* key"): NativeSpendCounters(None).get_local(key) @@ -153,7 +169,7 @@ async def test_post_call_increments_go_on_the_requests_native_post_call_batch_un targets: Final[list[str | None]] = [] class Recording(FakeSpendCounters): - def increment_many_post_call(self, request: object, increments: object) -> bool: + def increment_many_post_call(self, request: object, increments: object) -> FakeSettlement | None: targets.append(current_service_target()) return super().increment_many_post_call(request, increments) # pyright: ignore[reportArgumentType] # recorded as given @@ -165,8 +181,8 @@ async def test_post_call_increments_go_on_the_requests_native_post_call_batch_un inside: Final = store.increment_many_post_call((("spend:key:k", 1.0),)) empty: Final = store.increment_many_post_call(()) - assert (outside, inside, empty) == (False, True, False) - assert spend.post_call == [(natives.built[0], (("spend:key:k", 1.0),))] + assert (outside, empty) == (None, None) + assert spend.post_call == [(natives.built[0], (("spend:key:k", 1.0),), inside)] assert targets == [SPEND_COUNTERS_TARGET] diff --git a/tests/unit/router_strategy/test_complexity_router.py b/tests/unit/router_strategy/test_complexity_router.py index 444cd2ba6dd..f06cb0d282d 100644 --- a/tests/unit/router_strategy/test_complexity_router.py +++ b/tests/unit/router_strategy/test_complexity_router.py @@ -7772,6 +7772,163 @@ class TestChildSessionTierPins: router.discard() +class TestClaudeCodeTierPins: + SESSION: Final = "11111111-2222-3333-4444-555555555555" + AGENT_A: Final = "a7256c04dfd99c9b9" + AGENT_B: Final = "a71aeeb9c67f32cee" + TOOL_USE: Final[dict[str, object]] = { + "role": "assistant", + "content": [{"type": "tool_use", "id": "t1", "name": "Read", "input": {}}], + } + TOOL_RESULT: Final[dict[str, object]] = { + "type": "tool_result", + "tool_use_id": "t1", + "content": "def add(a, b): return a + b", + } + # Captured from Claude Code 2.1.316 on 2026-11-09, trimmed to the sentinel and a little context. + SUMMARY_ASK: Final = ( + "Describe your most recent action in 3-5 words using present tense (-ing). " + "Name the file or function, not the branch. Do not use tools." + ) + SUGGESTION_ASK: Final = ( + "[SUGGESTION MODE: Suggest what the user might naturally type next into Claude Code.]\n\n" + "Reply with ONLY the suggestion, no quotes or explanation." + ) + + @staticmethod + def _router(route_housekeeping_to_cheapest_tier: bool = True) -> Router: + return Router( + model_list=[ + { + "model_name": "claude-code-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "tiers": {"SIMPLE": "simple", "REASONING": "reasoning"}, + "classification_mode": "user_turn", + "keyword_tier_rules": [{"keywords": ["reason carefully"], "tier": "REASONING"}], + "route_housekeeping_to_cheapest_tier": route_housekeeping_to_cheapest_tier, + }, + }, + }, + *[ + { + "model_name": group, + "litellm_params": {"model": "openai/test-model", "api_key": "test-only"}, + "model_info": {"id": group}, + } + for group in ("simple", "reasoning") + ], + ], + num_retries=0, + ) + + @staticmethod + def _ask(prompt: str) -> list[dict[str, object]]: + return [{"role": "user", "content": prompt}] + + @classmethod + def _continuation(cls) -> list[dict[str, object]]: + return [*cls._ask("Hello!"), cls.TOOL_USE, {"role": "user", "content": [cls.TOOL_RESULT]}] + + @classmethod + def _side_request(cls, side_ask: str) -> list[dict[str, object]]: + return [ + *cls._ask("Hello!"), + cls.TOOL_USE, + {"role": "user", "content": [cls.TOOL_RESULT, {"type": "text", "text": side_ask}]}, + ] + + async def _route( + self, + router: Router, + messages: list[dict[str, object]], + agent_id: str | None = None, + request_class: str | None = None, + ) -> tuple[str, str]: + from litellm.proxy.litellm_pre_call_utils import ( + LiteLLMProxyRequestSetup, + metadata_variable_name_for_route, + ) + + headers: Final = { + "x-claude-code-session-id": self.SESSION, + **({"x-claude-code-agent-id": agent_id} if agent_id is not None else {}), + **({"x-claude-code-request-class": request_class} if request_class is not None else {}), + } + bucket: Final = metadata_variable_name_for_route("/v1/messages") + data: Final[dict[str, object]] = { + bucket: {"user_api_key_hash": "caller-a", "user_api_key_request_route": "/v1/messages"}, + "proxy_server_request": {"headers": headers}, + } + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers=headers, data=data, _metadata_variable_name=bucket + ) + await router.async_get_available_deployment(model="claude-code-router", request_kwargs=data, messages=messages) + metadata: Final = data[bucket] + assert isinstance(metadata, dict) + assert metadata["session_id"] == self.SESSION + decision: Final = TypeAdapter(StandardLoggingRoutingDecision).validate_python(metadata["routing_decision"]) + return str(decision.get("tier")), decision["cause"] + + @pytest.mark.asyncio + async def test_parent_and_each_subagent_keep_their_own_pin(self) -> None: + router: Final = self._router() + try: + for messages, agent_id, expected in ( + (self._ask("reason carefully"), None, ("REASONING", "literal_keyword_match")), + (self._ask("Hello!"), self.AGENT_A, ("SIMPLE", "heuristic_scorer")), + (self._continuation(), self.AGENT_A, ("SIMPLE", "user_turn_continuation")), + (self._continuation(), None, ("REASONING", "user_turn_continuation")), + (self._ask("reason carefully"), self.AGENT_B, ("REASONING", "literal_keyword_match")), + (self._continuation(), self.AGENT_A, ("SIMPLE", "user_turn_continuation")), + (self._continuation(), self.AGENT_B, ("REASONING", "user_turn_continuation")), + ): + assert await self._route(router, messages, agent_id) == expected + finally: + router.discard() + + @pytest.mark.asyncio + @pytest.mark.parametrize("hint_header", [True, False], ids=["request-class-hint", "sentinel-fallback"]) + @pytest.mark.parametrize( + "agent_id,side_ask", + [(AGENT_A, SUMMARY_ASK), (None, SUGGESTION_ASK)], + ids=["subagent-summary", "main-prompt-suggestion"], + ) + async def test_side_requests_route_cheapest_and_never_move_the_pin( + self, hint_header: bool, agent_id: str | None, side_ask: str + ) -> None: + own_class: Final = ("subagent" if agent_id is not None else "main") if hint_header else None + side_class: Final = "auxiliary" if hint_header else None + router: Final = self._router() + try: + for messages, request_class, expected in ( + (self._ask("reason carefully"), own_class, ("REASONING", "literal_keyword_match")), + (self._side_request(side_ask), side_class, ("SIMPLE", "housekeeping")), + (self._continuation(), own_class, ("REASONING", "user_turn_continuation")), + ): + assert await self._route(router, messages, agent_id, request_class) == expected + finally: + router.discard() + + @pytest.mark.asyncio + async def test_the_auxiliary_hint_routes_housekeeping_when_no_sentinel_matches(self) -> None: + router: Final = self._router() + try: + for messages, request_class, expected in ( + (self._ask("reason carefully"), "subagent", ("REASONING", "literal_keyword_match")), + ( + self._side_request("Summarize what you just did in a few words."), + "auxiliary", + ("SIMPLE", "housekeeping"), + ), + (self._continuation(), "subagent", ("REASONING", "user_turn_continuation")), + ): + assert await self._route(router, messages, self.AGENT_A, request_class) == expected + finally: + router.discard() + + class TestClassificationMode: """Test classification_mode='user_turn': classify only requests whose newest turn is a new human ask; tool-loop continuation turns replay the session's held routing decision.""" diff --git a/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py index 76a574bae24..adf08e43c95 100644 --- a/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py @@ -3293,3 +3293,194 @@ async def test_a_changed_named_credential_moves_the_pin_within_the_secret_manage assert pin_before_the_change == [peer] assert pin_after_the_change == [other_org] assert reasoning_after_the_change == _reasoning_input_minted_by("origin")[0] + + +# --------------------------------------------------------------------------- +# Operator-declared encryption boundary (model_info.encrypted_content_boundary) +# --------------------------------------------------------------------------- + + +def _regional_deployment(model_name: str, region: str, boundary: str | None) -> dict: + return { + "model_name": model_name, + "litellm_params": { + "model": "openai/gpt-5.6", + "api_base": f"https://{region}.example.invalid/v1", + "api_key": "sk-one-org", + }, + "model_info": {"id": region} | ({} if boundary is None else {"encrypted_content_boundary": boundary}), + } + + +async def _forwarded_reasoning_after_model_switch(origin_boundary: str | None, peer_boundary: str | None) -> tuple: + router: Final = litellm.Router( + model_list=[ + _regional_deployment("reasoning-a", "region-1", origin_boundary), + _regional_deployment("reasoning-b", "region-2", peer_boundary), + ], + optional_pre_call_checks=["encrypted_content_affinity"], + num_retries=0, + ) + handler: Final = AsyncMock(return_value=_build_mock_response(output_items=[], response_id="resp_second")) + with patch( + "litellm.llms.custom_httpx.llm_http_handler.BaseLLMHTTPHandler.async_response_api_handler", + new=handler, + ): + response: Final = await router.aresponses( + model="reasoning-b", + input=[ + { + "type": "reasoning", + "summary": [], + "encrypted_content": ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id( + "blob-from-region-1", "region-1" + ), + }, + {"role": "user", "content": "follow-up"}, + ], + litellm_metadata={}, + ) + forwarded_reasoning: Final = handler.await_args.kwargs["input"][0] + return response._hidden_params["model_id"], forwarded_reasoning.get("encrypted_content") + + +@pytest.mark.asyncio +async def test_declared_boundary_keeps_reasoning_across_endpoints_with_different_api_bases(): + assert await _forwarded_reasoning_after_model_switch("org-1", "org-1") == ("region-2", "blob-from-region-1") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("origin_boundary", "peer_boundary"), + [(None, None), ("org-1", "org-2"), ("org-1", None), (None, "org-1")], + ids=["undeclared", "different-boundaries", "only-origin-declared", "only-peer-declared"], +) +async def test_reasoning_is_stripped_unless_both_deployments_declare_the_same_boundary( + origin_boundary, peer_boundary +): + assert await _forwarded_reasoning_after_model_switch(origin_boundary, peer_boundary) == ("region-2", None) + + +def test_declared_boundary_overrides_a_shared_api_base_and_key(): + origin: Final = _make_originating_mock("https://region-1.example.invalid/v1", "sk-one-org") + origin.model_info = litellm.types.router.ModelInfo(id="region-1", encrypted_content_boundary="org-1") + check: Final = EncryptedContentAffinityCheck(ctx=_make_router_mock_with_cooldown(origin)) + other_org: Final = { + "model_info": {"id": "region-1-org-2", "encrypted_content_boundary": "org-2"}, + "litellm_params": {"api_base": "https://region-1.example.invalid/v1", "api_key": "sk-one-org"}, + } + + matches, _ = check._find_deployments_on_same_encryption_boundary([other_org], "region-1") + + assert matches == [] + + +@pytest.mark.asyncio +async def test_boundary_pin_only_selects_deployments_that_declare_the_origins_boundary(): + origin: Final = _make_originating_mock("https://region-1.example.invalid/v1", "sk-one-org") + origin.model_info = litellm.types.router.ModelInfo(id="region-1", encrypted_content_boundary="org-1") + check: Final = EncryptedContentAffinityCheck(ctx=_make_router_mock_with_cooldown(origin)) + peer: Final = _regional_deployment("reasoning-b", "region-2", "org-1") + stranger: Final = _regional_deployment("reasoning-b", "region-3", "org-2") + request_kwargs: Final = { + "input": [ + { + "type": "reasoning", + "encrypted_content": ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id( + "blob-from-region-1", "region-1" + ), + } + ] + } + + result: Final = await check.async_filter_deployments( + model="reasoning-b", + healthy_deployments=[stranger, peer], + messages=None, + request_kwargs=request_kwargs, + ) + + assert result == [peer] + assert request_kwargs["input"][0]["encrypted_content"] == ( + ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id("blob-from-region-1", "region-1") + ) + + +def test_a_boundary_declared_on_one_side_keeps_the_shared_api_base_and_key_match(): + origin: Final = _make_originating_mock("https://region-1.example.invalid/v1", "sk-one-org") + origin.model_info = litellm.types.router.ModelInfo(id="region-1", encrypted_content_boundary="org-1") + check: Final = EncryptedContentAffinityCheck(ctx=_make_router_mock_with_cooldown(origin)) + undeclared_twin: Final = { + "model_info": {"id": "region-1-twin"}, + "litellm_params": {"api_base": "https://region-1.example.invalid/v1", "api_key": "sk-one-org"}, + } + + matches, _ = check._find_deployments_on_same_encryption_boundary([undeclared_twin], "region-1") + + assert matches == [undeclared_twin] + + +def test_declared_boundary_matches_deployments_whose_credentials_cannot_be_resolved(): + origin: Final = _make_originating_mock("https://region-1.example.invalid/v1", "") + origin.litellm_params.model_dump.return_value = {"api_base": "https://region-1.example.invalid/v1"} + origin.model_info = litellm.types.router.ModelInfo(id="region-1", encrypted_content_boundary="org-1") + check: Final = EncryptedContentAffinityCheck(ctx=_make_router_mock_with_cooldown(origin)) + peer: Final = { + "model_info": {"id": "region-2", "encrypted_content_boundary": "org-1"}, + "litellm_params": {"api_base": "https://region-2.example.invalid/v1"}, + } + + matches, _ = check._find_deployments_on_same_encryption_boundary([peer], "region-1") + + assert matches == [peer] + + +async def _forwarded_reasoning_after_order_fallback(boundary: str | None) -> tuple: + primary: Final = _regional_deployment("reasoning", "region-1", boundary) + backup: Final = _regional_deployment("reasoning", "region-3", boundary) + router: Final = litellm.Router( + model_list=[ + primary | {"litellm_params": primary["litellm_params"] | {"order": 1}}, + backup | {"litellm_params": backup["litellm_params"] | {"order": 2}}, + ], + optional_pre_call_checks=["encrypted_content_affinity"], + disable_cooldowns=True, + num_retries=0, + ) + + async def upstream(*, litellm_params, **_kwargs): + if "region-1" in litellm_params.api_base: + raise litellm.InternalServerError("upstream 500", llm_provider="openai", model="gpt-5.6") + return _build_mock_response(output_items=[], response_id="resp_second") + + handler: Final = AsyncMock(side_effect=upstream) + with patch( + "litellm.llms.custom_httpx.llm_http_handler.BaseLLMHTTPHandler.async_response_api_handler", + new=handler, + ): + response: Final = await router.aresponses( + model="reasoning", + input=[ + { + "type": "reasoning", + "summary": [], + "encrypted_content": ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id( + "blob-from-region-1", "region-1" + ), + }, + {"role": "user", "content": "follow-up"}, + ], + litellm_metadata={}, + ) + forwarded_reasoning: Final = handler.await_args.kwargs["input"][0] + return response._hidden_params["model_id"], forwarded_reasoning.get("encrypted_content") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("boundary", "expected_reasoning"), + [("org-1", "blob-from-region-1"), (None, None)], + ids=["declared", "undeclared"], +) +async def test_order_fallback_keeps_reasoning_only_on_a_declared_boundary_peer(boundary, expected_reasoning): + assert await _forwarded_reasoning_after_order_fallback(boundary) == ("region-3", expected_reasoning) diff --git a/tests/unit/rust_bridge/host/routes/test_messages.py b/tests/unit/rust_bridge/host/routes/test_messages.py index 9d8b09dd1c4..aa7dcc74841 100644 --- a/tests/unit/rust_bridge/host/routes/test_messages.py +++ b/tests/unit/rust_bridge/host/routes/test_messages.py @@ -184,6 +184,28 @@ def test_body_params_keep_the_requested_effort_without_deployment_levels() -> No assert messages_host.body_params(request)["output_config"] == {"effort": "max"} +@pytest.mark.parametrize( + ("additional_drop_params", "expected_keys"), + [ + (["thinking"], {"max_tokens"}), + (["top_k"], {"max_tokens", "thinking", "output_config"}), + ], +) +def test_body_params_derive_no_effort_from_a_param_the_deployment_lists( + additional_drop_params: list[str], expected_keys: set[str] +) -> None: + request: Final[dict[str, object]] = { + "model": _ADAPTIVE_OPUS, + "messages": [], + "max_tokens": 8192, + "thinking": {"type": "enabled", "budget_tokens": 4096}, + "additional_drop_params": additional_drop_params, + "model_info": {"id": "deployment-1", "reasoning_effort_levels": ["low", "medium", "high"]}, + } + + assert set(messages_host.body_params(request)).difference({"messages", "model"}) == expected_keys + + def _marked(phase: str | None = None, *args: object) -> ValueError: error: Final = ValueError(*args) if phase is not None: @@ -207,6 +229,16 @@ def test_a_native_read_timeout_maps_to_the_public_timeout_with_python_text() -> assert "Read timed out. Timeout passed=1.0" in str(mapped) +@pytest.mark.parametrize(("phase", "billed"), (("Read", True), ("Connection", False))) +def test_a_native_timeout_is_billed_on_the_failure_row_as_python_bills_it(phase: str, billed: bool) -> None: + from litellm.proxy.utils import _provider_billed_the_failure + + mapped: Final = messages_host.map_failure(_marked(phase), {"model": _ADAPTIVE_OPUS, "timeout": 1.0}, "anthropic") + + assert isinstance(mapped, litellm.Timeout) + assert _provider_billed_the_failure(mapped) is billed + + def test_an_upstream_status_maps_through_the_provider_error_class() -> None: upstream: Final = ValueError(429, "slow down") upstream.headers = [("retry-after", "3")] # pyright: ignore[reportAttributeAccessIssue] # set by the native host @@ -218,6 +250,44 @@ def test_an_upstream_status_maps_through_the_provider_error_class() -> None: assert mapped.__context__ is upstream +@pytest.mark.parametrize( + ("body", "expected"), + [ + ("", {"content-type": "text/plain"}), + ("gateway", {"content-type": "text/plain", "content-length": "7"}), + ], +) +def test_upstream_failure_headers_match_the_python_clients_rebuilt_response( + body: str, expected: dict[str, str] +) -> None: + from litellm.litellm_core_utils.exception_mapping_utils import _get_response_headers + + upstream: Final = ValueError(504, body) + upstream.headers = [ # pyright: ignore[reportAttributeAccessIssue] # set by the native host + ("content-type", "text/plain"), + ("content-encoding", "gzip"), + ("content-length", "42"), + ] + + mapped: Final = messages_host.map_failure(upstream, {"model": _ADAPTIVE_OPUS}, "anthropic") + + assert isinstance(mapped, litellm.Timeout) + assert dict(_get_response_headers(mapped) or {}) == expected + + +@pytest.mark.parametrize("status", [200, 500]) +def test_upstream_failure_without_wire_headers_states_no_headers(status: int) -> None: + from litellm.litellm_core_utils.exception_mapping_utils import _get_response_headers + + upstream: Final = ValueError(status, "not a messages response") + upstream.headers = [] # pyright: ignore[reportAttributeAccessIssue] # set by the native host + + mapped: Final = messages_host.map_failure(upstream, {"model": _ADAPTIVE_OPUS}, "anthropic") + + assert isinstance(mapped, Exception) + assert not dict(_get_response_headers(mapped) or {}) + + def test_a_late_rejection_tells_the_logger_the_resolved_call_first() -> None: import datetime @@ -245,3 +315,133 @@ def test_a_late_rejection_tells_the_logger_the_resolved_call_first() -> None: assert isinstance(mapped, litellm.BadRequestError) assert logger.model_call_details["custom_llm_provider"] == "anthropic" assert logger.model_call_details["optional_params"] == {"max_tokens": 8} + + +@pytest.mark.parametrize( + ("aiohttp", "cut", "error_class", "text"), + ( + ( + True, + (60, None, False), + "ReadError", + "Response payload is not completed: ", + ), + ( + True, + (19, 100, True), + "ReadError", + "Response payload is not completed: . " + "ConnectionResetError(104, 'Connection reset by peer')", + ), + ( + False, + (60, None, False), + "RemoteProtocolError", + "peer closed connection without sending complete message body (incomplete chunked read)", + ), + ( + False, + (19, 100, False), + "RemoteProtocolError", + "peer closed connection without sending complete message body (received 19 bytes, expected 100)", + ), + (False, (19, 100, True), "ReadError", ""), + ), +) +def test_a_stream_body_cut_reaches_the_caller_as_the_transport_raised_it( + monkeypatch: pytest.MonkeyPatch, + aiohttp: bool, + cut: tuple[int, int | None, bool], + error_class: str, + text: str, +) -> None: + import httpx + + from litellm.rust_bridge.host import http_transport + + monkeypatch.setattr(http_transport, "uses_aiohttp_transport", lambda: aiohttp) + native: Final = RuntimeError("upstream network error") + native.messages_body_cut = cut # pyright: ignore[reportAttributeAccessIssue] # marker the native host sets + + mapped: Final = messages_host.map_failure(native, {"model": _ADAPTIVE_OPUS}, "anthropic") + + assert type(mapped) is getattr(httpx, error_class) + assert str(mapped) == text + + +_TRANSPORT: Final = "error sending request: client error" + + +@pytest.mark.parametrize( + ("aiohttp", "native", "text"), + ( + ( + True, + f"could not reach the provider: {_TRANSPORT} (Connect): tcp connect error: Connection refused (os error 111)", + "Cannot connect to host 127.0.0.1:9 ssl: [Connect call failed ('127.0.0.1', 9)]", + ), + ( + True, + f"could not reach the provider: {_TRANSPORT} (Connect): dns error: error resolving DNS: " + "failed to lookup address information: Name or service not known", + "Cannot connect to host 127.0.0.1:9 ssl: [Name or service not known]", + ), + ( + True, + f"could not reach the provider: {_TRANSPORT} (Connect): tls handshake eof", + "Cannot connect to host 127.0.0.1:9 ssl: [None]", + ), + ( + True, + f"upstream network error: {_TRANSPORT} (SendRequest): connection error: Connection reset by peer (os error 104)", + "[Errno 104] Connection reset by peer", + ), + ( + True, + f"upstream network error: {_TRANSPORT} (SendRequest): connection closed before message completed", + "Server disconnected", + ), + ( + False, + f"could not reach the provider: {_TRANSPORT} (Connect): tcp connect error: Connection refused (os error 111)", + "All connection attempts failed", + ), + ( + False, + f"could not reach the provider: {_TRANSPORT} (Connect): dns error: error resolving DNS: " + "failed to lookup address information: Name or service not known", + "[Errno -2] Name or service not known", + ), + (False, f"could not reach the provider: {_TRANSPORT} (Connect): tls handshake eof", ""), + ( + False, + f"upstream network error: {_TRANSPORT} (SendRequest): connection error: Connection reset by peer (os error 104)", + "", + ), + ( + False, + f"upstream network error: {_TRANSPORT} (SendRequest): connection closed before message completed", + "Server disconnected without sending a response.", + ), + ), +) +def test_a_failure_before_any_response_maps_like_the_python_transport_error( + monkeypatch: pytest.MonkeyPatch, aiohttp: bool, native: str, text: str +) -> None: + from litellm.rust_bridge.host import http_transport + + monkeypatch.setattr(http_transport, "uses_aiohttp_transport", lambda: aiohttp) + monkeypatch.setattr(messages_host, "get_ssl_configuration", lambda _ssl_verify: "") + + mapped: Final = messages_host.map_failure( + RuntimeError(native), {"model": _ADAPTIVE_OPUS, "api_base": "http://127.0.0.1:9"}, "anthropic" + ) + + assert type(mapped) is litellm.InternalServerError + assert mapped.status_code == 500 + assert mapped.message == ( + f"litellm.InternalServerError: AnthropicException - {text}. Handle with `litellm.InternalServerError`." + ) + assert "This looks like a bug in LiteLLM" not in str(mapped) diff --git a/tests/unit/rust_bridge/host/test_cache.py b/tests/unit/rust_bridge/host/test_cache.py index 9d532631ef2..c2b87cca1dc 100644 --- a/tests/unit/rust_bridge/host/test_cache.py +++ b/tests/unit/rust_bridge/host/test_cache.py @@ -1,10 +1,12 @@ +import asyncio from typing import Final import pytest import litellm from litellm.caching.caching import Cache -from litellm.rust_bridge.host.cache import configured_cache +from litellm.caching.caching_handler import _PENDING_CACHE_WRITES +from litellm.rust_bridge.host.cache import configured_cache, store_response from litellm.types.caching import LiteLLMCacheType @@ -19,3 +21,33 @@ def test_configured_cache_preserves_identity_and_observes_runtime_replacement(mo monkeypatch.setattr(litellm, "cache", None) assert configured_cache() is None + + +_REQUEST: Final = {"model": "claude-sonnet-4-5", "messages": [{"role": "user", "content": "hi"}]} + + +async def _stored(response: dict[str, object]) -> object: + cache: Final = Cache(type=LiteLLMCacheType.LOCAL) + store_response(cache, response, **_REQUEST) + await asyncio.gather(*_PENDING_CACHE_WRITES) + return await cache.async_get_cache(**_REQUEST) + + +@pytest.mark.asyncio +async def test_store_response_drops_only_the_rust_marker_header_from_the_cached_entry() -> None: + stored: Final = await _stored( + { + "id": "msg_1", + "_hidden_params": {"additional_headers": {"x-litellm-rust": "true", "llm_provider-x": "1"}, "model": "m"}, + } + ) + assert stored == { + "id": "msg_1", + "_hidden_params": {"additional_headers": {"llm_provider-x": "1"}, "model": "m"}, + } + + +@pytest.mark.asyncio +async def test_store_response_keeps_an_entry_without_the_rust_marker_unchanged() -> None: + response: Final = {"id": "msg_2", "_hidden_params": {"additional_headers": {"llm_provider-x": "1"}}} + assert await _stored(response) == response diff --git a/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py b/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py index 1ba0188e76d..cb594767908 100644 --- a/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py +++ b/tests/unit/rust_bridge/host/test_callbacks_legacy_python.py @@ -1,3 +1,4 @@ +import asyncio import datetime import inspect from collections.abc import Mapping @@ -101,6 +102,22 @@ def _kwargs_with_a_budget_reservation(reservation: dict) -> dict[str, object]: return {**_OCR_KWARGS, "metadata": {"user_api_key_budget_reservation": reservation}} +def test_a_deferred_success_waits_for_the_proxy_release() -> None: + logger: Final = _supplied_logger() + released: Final[list[str]] = [] + start: Final = datetime.datetime.now() + end: Final = start + datetime.timedelta(seconds=1) + + legacy.defer_success(logger, lambda: released.append("released"), "response", start, end) + + assert logger.deferred_success_result == ("response", start, end) + assert released == [] + release: Final = logger.enqueue_deferred_logging + assert release is not None + release() + assert released == ["released"] + + def test_setup_claims_the_budget_reservation_for_an_async_call() -> None: reservation: Final = _budget_reservation() @@ -247,6 +264,30 @@ def test_the_rust_contract_matches_the_shim_signatures() -> None: assert contract == {name: list(inspect.signature(getattr(legacy, name)).parameters) for name in contract} +def test_messages_input_is_the_body_as_the_python_handler_builds_and_serializes_it() -> None: + body: Final = { + "model": "claude-sonnet-4-5", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 16, + "system": "brief", + "stream": True, + "context_management": {"edits": []}, + "output_config": {"effort": "low"}, + } + + logged: Final = legacy.messages_input( + body, {"output_config": {"effort": "low"}, "context_management": {"edits": []}} + ) + + assert logged == [ + { + "role": "user", + "content": '{"messages": [{"role": "user", "content": "hi"}], "max_tokens": 16, "model": "claude-sonnet-4-5", ' + '"stream": true, "system": "brief", "output_config": {"effort": "low"}, "context_management": {"edits": []}}', + } + ] + + def _messages_logger() -> Logging: return Logging( model="claude-sonnet-4-5", @@ -298,3 +339,67 @@ def test_stream_head_is_stamped_in_place_with_the_call_hidden_params() -> None: assert head["additional_headers"] == {"request-id": "req_native", "llm_provider-request-id": "req_native"} assert head["litellm_call_id"] == "messages" assert "_response_ms" in head + + +@pytest.mark.asyncio +async def test_a_native_stream_failure_logs_like_the_python_chunk_processor(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm._logging import verbose_proxy_logger + from litellm.proxy.pass_through_endpoints.streaming_handler import PassThroughStreamingHandler + + logged: Final[list[tuple[object, ...]]] = [] + scheduled: Final[list[Mapping[str, object]]] = [] + + async def schedule(**kwargs: object) -> None: + scheduled.append(MappingProxyType(dict(kwargs))) + + monkeypatch.setattr(verbose_proxy_logger, "error", lambda *args: logged.append(args)) + monkeypatch.setattr(PassThroughStreamingHandler, "schedule_stream_failure_logging", schedule) + logger: Final = _supplied_logger() + logger.model_call_details["cache_hit"] = False + cut: Final = httpx.ReadError("cut") + + await legacy.stream_failure(logger, "anthropic", {}, [b"event: ping\n\n"], cut) + + assert logged == [("Error in chunk_processor: %s", cut)] + assert "cache_hit" not in logger.model_call_details + assert scheduled[0]["exception"] is cut + + +@pytest.mark.asyncio +async def test_a_proxy_stream_abandoned_mid_cancel_bills_after_the_disconnect_cleanup( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy.pass_through_endpoints.streaming_handler import PassThroughStreamingHandler + + logger: Final = _messages_logger() + logger.update_environment_variables( + litellm_params={"metadata": {}, "proxy_server_request": {"url": "/v1/messages"}}, optional_params={} + ) + + async def billing() -> None: + return None + + events: Final[list[str]] = [] + monkeypatch.setattr(PassThroughStreamingHandler, "_route_streaming_logging_to_handler", lambda **_: billing()) + monkeypatch.setattr(legacy, "enqueue_logging", lambda coroutine: (events.append("billed"), coroutine.close())) + started: Final = asyncio.Event() + now: Final = datetime.datetime.now() + + async def relay() -> None: + try: + started.set() + await asyncio.Event().wait() + except asyncio.CancelledError: + legacy.stream_success(logger, "/v1/messages", "anthropic", {}, [b"event: ping\n\n"], now, now, None, True) + await asyncio.sleep(0) + events.append("disconnect recorded") + raise + + task: Final = asyncio.create_task(relay()) + await started.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + await asyncio.sleep(0) + + assert events == ["disconnect recorded", "billed"] diff --git a/tests/unit/rust_bridge/test_detectors.py b/tests/unit/rust_bridge/test_detectors.py index a6e6b2c0dca..fccc01dd8e6 100644 --- a/tests/unit/rust_bridge/test_detectors.py +++ b/tests/unit/rust_bridge/test_detectors.py @@ -111,7 +111,7 @@ def _wrap_nested_level(value: object, depth: int) -> object: (detectors.prompt_management, {"model": "litellm_agent/gpt-5"}, True), (detectors.prompt_management, {"model": "gpt-5"}, False), (detectors.caller_client, {"client": object()}, True), - (detectors.caller_client, {"shared_session": object()}, True), + (detectors.caller_client, {"shared_session": object()}, False), (detectors.caller_client, {"client": None, "shared_session": None}, False), (detectors.caller_client, {}, False), (detectors.extra_body, {"extra_body": {"stream": True}}, True), diff --git a/tests/unit/rust_bridge/test_public_call.py b/tests/unit/rust_bridge/test_public_call.py index 5f50db6ea5a..27c3c119c1b 100644 --- a/tests/unit/rust_bridge/test_public_call.py +++ b/tests/unit/rust_bridge/test_public_call.py @@ -1,9 +1,22 @@ -from collections.abc import Mapping, Sequence -from typing import Final +import asyncio +from collections.abc import Awaitable, Callable, Mapping, Sequence +from typing import Final, TypeAlias import pytest -from litellm.rust_bridge.public_call import AsyncFlag, native_call, signature +import litellm +from litellm.rust_bridge.public_call import ( + AsyncFlag, + NativeCall, + native_acall_hook, + native_acall_hook_on_loop, + native_call, + signature, +) + +NativeAcallHook: TypeAlias = Callable[ + [Callable[[NativeCall], Awaitable[None]], NativeCall, tuple[object, ...], Mapping[str, object]], Awaitable[None] +] def _messages( @@ -76,3 +89,54 @@ def test_flag_keywords_name_positionals_and_drop_only_the_flag() -> None: def test_flag_keywords_reject_a_parameter_given_twice() -> None: with pytest.raises(TypeError): _FLAG.keywords((128,), {"max_tokens": 64}) + + +_CALL: Final = native_call(_SIGNATURE, (128, [], "model"), {}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("kwargs", "expected_timeout", "expected_num_retries"), + ( + ({"timeout": 30.0, "stream_timeout": 0.5, "num_retries": 0}, 30.0, 0), + ({"stream": True, "timeout": 30.0, "stream_timeout": 0.5}, 0.5, None), + ({"stream": True, "timeout": 30.0}, 30.0, None), + ({}, None, None), + ), +) +@pytest.mark.parametrize("hook", (native_acall_hook, native_acall_hook_on_loop)) +async def test_async_hook_failures_carry_the_client_wrapper_stamp( + hook: NativeAcallHook, + kwargs: Mapping[str, object], + expected_timeout: float | None, + expected_num_retries: int | None, +) -> None: + async def native(call: NativeCall) -> None: + raise litellm.Timeout(message="timed out", model="model", llm_provider="anthropic") + + with pytest.raises(litellm.Timeout) as raised: + await hook(native, _CALL, (), kwargs) + + assert getattr(raised.value, "timeout", "unset") == expected_timeout + assert getattr(raised.value, "num_retries", None) == expected_num_retries + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("hook", "yields"), ((native_acall_hook, True), (native_acall_hook_on_loop, False))) +async def test_only_the_executor_hook_lets_a_task_started_before_the_call_run_first( + hook: NativeAcallHook, yields: bool +) -> None: + messages: Final = [{"role": "user", "content": "during about secret"}] + sent: Final[list[object]] = [] + + async def mask() -> None: + messages[0]["content"] = "during about [REDACTED]" + + async def native(call: NativeCall) -> None: + sent.append(messages[0]["content"]) + + guardrail: Final = asyncio.create_task(mask()) + await hook(native, _CALL, (), {}) + await guardrail + + assert sent == ["during about [REDACTED]" if yields else "during about secret"] diff --git a/tests/unit/test_claude_fable_5_config.py b/tests/unit/test_claude_fable_5_config.py index 7ec76742bb8..2157364efb3 100644 --- a/tests/unit/test_claude_fable_5_config.py +++ b/tests/unit/test_claude_fable_5_config.py @@ -1,5 +1,5 @@ """ -Validate Claude Fable 5 and Claude Fable 5.1 model configuration entries. +Validate Claude Fable 5, Claude Fable 5.1 and Claude Fable 5.2 model configuration entries. Fable 5 is a new tier above Opus ($10/$50 per MTok) with the same adaptive-only API surface as Opus 4.7/4.8. The cost-map entries below are what make the model @@ -11,9 +11,11 @@ shape, which Fable 5 rejects with a 400. import json import os +from typing import Final import pytest +import litellm from litellm.constants import BEDROCK_CONVERSE_MODELS from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap @@ -92,3 +94,81 @@ def test_adaptive_thinking_detected_for_fable_5_1(local_model_cost_map, model): from litellm.llms.anthropic.common_utils import AnthropicModelInfo assert AnthropicModelInfo.is_adaptive_thinking_model(model, "anthropic") is True + + +FABLE_5_2_VARIANTS = ( + "claude-fable-5-2", + "anthropic.claude-fable-5-2", + "global.anthropic.claude-fable-5-2", + "us.anthropic.claude-fable-5-2", + "eu.anthropic.claude-fable-5-2", + "us-gov.anthropic.claude-fable-5-2", + "bedrock/us-gov-west-1/anthropic.claude-fable-5-2", + "bedrock/us-gov-east-1/anthropic.claude-fable-5-2", + "vertex_ai/claude-fable-5-2", + "vertex_ai/claude-fable-5-2@default", + "bedrock/claude_platform/claude-fable-5-2", +) + + +def test_fable_5_2_present_in_bundled_backup() -> None: + backup = GetModelCostMap.load_local_model_cost_map() + root = _load_root_cost_map() + for model_name in FABLE_5_2_VARIANTS: + assert model_name in backup, f"Missing from backup cost map: {model_name}" + assert backup[model_name] == root[model_name], model_name + + +@pytest.mark.parametrize( + "model", + [ + "claude-fable-5-2", + "anthropic/claude-fable-5-2", + "claude-fable-latest", + "anthropic.claude-fable-5-2", + "bedrock/us.anthropic.claude-fable-5-2", + "bedrock/invoke/eu.anthropic.claude-fable-5-2", + "bedrock/global.anthropic.claude-fable-5-2", + "bedrock/claude-fable-latest", + "vertex_ai/claude-fable-5-2", + "vertex_ai/claude-fable-5-2@default", + "bedrock/claude_platform/claude-fable-5-2", + ], +) +def test_adaptive_thinking_detected_for_fable_5_2(local_model_cost_map, model: str) -> None: + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + assert AnthropicModelInfo.is_adaptive_thinking_model(model, "anthropic") is True + + +def test_fable_5_2_registered_for_bedrock_converse() -> None: + assert "anthropic.claude-fable-5-2" in BEDROCK_CONVERSE_MODELS + + +def test_fable_5_2_fast_mode_doubles_every_token_price_on_the_anthropic_api(local_model_cost_map) -> None: + """Anthropic bills fast mode at 2x the standard rate on every token type + (platform.claude.com/docs/en/about-claude/pricing, read 2026-11-09), and the + Anthropic cost path is the one that reads the row's `fast` multiplier.""" + from litellm.litellm_core_utils.llm_cost_calc.utils import get_provider_specific_speed_multiplier + from litellm.types.utils import Usage + + info: Final = litellm.get_model_info("claude-fable-5-2") + assert get_provider_specific_speed_multiplier(info, Usage(speed="fast")) == 2.0 + assert get_provider_specific_speed_multiplier(info, Usage()) == 1.0 + + +def test_fable_5_2_vertex_rows_do_not_advertise_fast_mode(local_model_cost_map) -> None: + """Vertex cost calculation ignores the `fast` multiplier, so a Vertex row that + advertised fast mode would bill fast calls at the standard rate.""" + for model in ("vertex_ai/claude-fable-5-2", "vertex_ai/claude-fable-5-2@default"): + info: Final = litellm.get_model_info(model) + assert info.get("supports_fast_mode") is not True, model + assert "fast" not in (info.get("provider_specific_entry") or {}), model + + +def test_fable_5_2_accepts_the_speed_param_on_the_anthropic_api(local_model_cost_map) -> None: + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + assert AnthropicModelInfo.supports_fast_mode("claude-fable-5-2", "anthropic") is True + assert AnthropicConfig._model_supports_speed_param("claude-fable-5-2", "anthropic") is True diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 49d51e99790..1d9f79aafc6 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -7532,82 +7532,65 @@ def test_whisper_azure(): assert round(cost, 5) == round(expected_cost, 5) -@pytest.mark.parametrize( - ("provider", "requested_model", "response_extras", "expected_seconds"), - [ - ("openai", "openai/gpt-fake-transcribe", {}, 2.0), - ("openai", "openai/gpt-fake-transcribe", {"_audio_transcription_duration": 1.35}, 1.35), - ("openai", "openai/gpt-fake-transcribe", {"duration": 1.34}, 1.34), - ("openai", "openai/whisper-fake", {}, 0.0), - ("azure", "gpt-fake-transcribe", {}, 0.0), - ], - ids=[ - "undecodable_upload_bills_provider_seconds", - "measured_duration_still_wins", - "provider_duration_still_wins", - "whisper_models_never_bill_provider_seconds", - "azure_never_bills_provider_seconds", - ], -) -def test_transcription_cost_falls_back_to_provider_billed_seconds( - monkeypatch: pytest.MonkeyPatch, - provider: str, - requested_model: str, - response_extras: Mapping[str, float], - expected_seconds: float, -) -> None: - model: Final = requested_model.removeprefix(f"{provider}/") +def _per_second_transcription( + monkeypatch: pytest.MonkeyPatch, provider: str, model: str, response_object: Mapping[str, object] +) -> TranscriptionResponse: monkeypatch.setitem( litellm.model_cost, model, {"input_cost_per_second": 0.01, "litellm_provider": provider, "mode": "audio_transcription"}, ) - transcription: Final = convert_to_model_response_object( - response_object={ - "text": "What is the weather in Paris?", - "usage": {"type": "duration", "seconds": 2}, - **response_extras, - }, + return convert_to_model_response_object( + response_object={"text": "What is the weather in Paris?", **response_object}, model_response_object=TranscriptionResponse(), hidden_params={"model": model, "custom_llm_provider": provider}, response_type="audio_transcription", ) - cost: Final = completion_cost( - completion_response=transcription, model=requested_model, call_type="atranscription" + +@pytest.mark.parametrize( + ("provider", "model", "measured"), + [ + ("openai", "whisper-fake", {}), + ("openai", "whisper-fake", {"_audio_transcription_duration": 2.4}), + ("openai", "whisper-fake", {"duration": 2.4}), + ("openai", "gpt-fake-transcribe", {"_audio_transcription_duration": 2.4}), + ("azure", "whisper-fake", {"_audio_transcription_duration": 2.4}), + ], + ids=[ + "unmeasured_upload", + "measured_upload", + "verbose_json_duration", + "gpt_transcribe_model", + "azure", + ], +) +def test_transcription_cost_bills_the_seconds_the_provider_billed( + monkeypatch: pytest.MonkeyPatch, provider: str, model: str, measured: Mapping[str, float] +) -> None: + transcription: Final = _per_second_transcription( + monkeypatch, provider, model, {"usage": {"type": "duration", "seconds": 3}, **measured} ) - assert cost == pytest.approx(0.01 * expected_seconds) + cost: Final = completion_cost( + completion_response=transcription, model=f"{provider}/{model}", call_type="atranscription" + ) + + assert cost == pytest.approx(0.01 * 3) @pytest.mark.parametrize( - ("provider", "measured_seconds", "expected_seconds"), - [("openai", 1.0, 1.0), ("azure", 1.0, 1.0), ("azure", None, 7.0)], - ids=["openai_measured_upload_wins", "azure_measured_upload_wins", "unmeasured_upload_bills_provider_seconds"], + ("measured", "expected_seconds"), + [({"_audio_transcription_duration": 2.4}, 2.4), ({"duration": 2.4}, 2.4), ({}, 0.0)], + ids=["measured_upload", "reply_duration", "nothing_known"], ) -def test_transcription_priced_by_provider_duration_still_bills_the_measured_upload( - monkeypatch: pytest.MonkeyPatch, provider: str, measured_seconds: float | None, expected_seconds: float +def test_transcription_cost_without_provider_seconds_bills_the_measured_length( + monkeypatch: pytest.MonkeyPatch, measured: Mapping[str, float], expected_seconds: float ) -> None: - monkeypatch.setitem( - litellm.model_cost, - "gpt-fake-transcribe", - { - "input_cost_per_second": 0.01, - "litellm_provider": provider, - "mode": "audio_transcription", - "provider_specific_entry": {"transcription_usage_duration": 1}, - }, - ) - measured: Final = {} if measured_seconds is None else {"audio_transcription_duration": measured_seconds} - transcription: Final = convert_to_model_response_object( - response_object={"text": "What is the weather in Paris?", "usage": {"type": "duration", "seconds": 7}}, - model_response_object=TranscriptionResponse(), - hidden_params={"model": "gpt-fake-transcribe", "custom_llm_provider": provider, **measured}, - response_type="audio_transcription", - ) + transcription: Final = _per_second_transcription(monkeypatch, "openai", "whisper-fake", measured) cost: Final = completion_cost( - completion_response=transcription, model=f"{provider}/gpt-fake-transcribe", call_type="atranscription" + completion_response=transcription, model="openai/whisper-fake", call_type="atranscription" ) assert cost == pytest.approx(0.01 * expected_seconds) diff --git a/tests/unit/test_router/core/test_endpoints_characterization_inference.py b/tests/unit/test_router/core/test_endpoints_characterization_inference.py index 8db0d6bea12..c975e45b1d3 100644 --- a/tests/unit/test_router/core/test_endpoints_characterization_inference.py +++ b/tests/unit/test_router/core/test_endpoints_characterization_inference.py @@ -205,6 +205,17 @@ async def test_sync_text_completion_merges_default_metadata_into_request_metadat assert (metadata["x"], metadata["y"], metadata["model_group"], sent["user"]) == (1, 2, "g", "u") +async def test_sync_text_completion_keeps_the_request_metadata_over_the_default_metadata( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core( + model_list=[deployment("g", "d1")], + default_litellm_params={"metadata": {"requester_metadata": {"env": "prod", "region": "us"}}}, + ) + core.text_completion(model="g", prompt="p", metadata={"requester_metadata": {"env": "dev", "team": "t1"}}) + assert _bucket(_sent(host))["requester_metadata"] == {"env": "dev", "region": "us", "team": "t1"} + + async def test_sync_text_completion_sends_the_model_group_name_not_the_deployment_model( make_core: MakeCore, host: EndpointHost ): @@ -329,12 +340,12 @@ async def test_default_litellm_params_metadata_merges_into_request_metadata(make assert (metadata["team"], metadata["trace"], metadata["model_group"]) == ("t", "x", "g") -async def test_default_litellm_params_metadata_overwrites_the_same_request_metadata_key( +async def test_the_request_metadata_key_wins_over_the_same_default_litellm_params_metadata_key( make_core: MakeCore, host: EndpointHost ): core: Final = make_core(model_list=[deployment("g", "d1")], default_litellm_params={"metadata": {"team": "t"}}) await core.acompletion(model="g", messages=MESSAGES, metadata={"team": "request"}) - assert _bucket(_sent(host))["team"] == "t" + assert _bucket(_sent(host))["team"] == "request" async def test_default_litellm_params_metadata_goes_to_litellm_metadata_for_generic_endpoints( diff --git a/tests/unit/test_router/core/test_gaps_l2800.py b/tests/unit/test_router/core/test_gaps_l2800.py index 78e237f0696..196ce20d4b3 100644 --- a/tests/unit/test_router/core/test_gaps_l2800.py +++ b/tests/unit/test_router/core/test_gaps_l2800.py @@ -246,7 +246,7 @@ async def test_responses_replace_kubernetes_pod_routing_in_the_request_metadata( """On the Responses path the routing goes to ``litellm_metadata``; a pod routing key the caller put in ``metadata`` is overwritten with the serving deployment's (none here).""" core: Final = make_core(model_list=[deployment("g", "d1")]) - await core.aresponses( # pyright: ignore[reportAttributeAccessIssue] + await core.aresponses( # pyright: ignore[reportAttributeAccessIssue] # endpoint method model="g", input="hi", metadata={KUBERNETES_POD_ROUTING_KEY: {"pod": "stale"}, "keep": 1} ) sent: Final = host.calls[0].kwargs @@ -332,7 +332,7 @@ async def test_the_process_request_timeout_outranks_the_router_timeout_once_a_ro process settings), and only when the router has a ``timeout`` of its own.""" monkeypatch.setattr(litellm, "request_timeout", 33.0) core: Final = make_core(model_list=[deployment("g", "d1")], timeout=42) - await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] # endpoint method assert host.calls[0].kwargs["timeout"] == 33.0 @@ -353,7 +353,7 @@ async def test_default_litellm_params_timeout_applies_when_no_router_or_process_ ) -> None: host.set_dynamic(request_timeout=None) core: Final = make_core(model_list=[deployment("g", "d1")], default_litellm_params={"timeout": 9}) - await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] # endpoint method assert host.calls[0].kwargs["timeout"] == 9 @@ -368,6 +368,257 @@ async def test_default_litellm_params_stream_timeout_is_the_last_stream_timeout_ assert host.calls[0].kwargs["stream_timeout"] == 4 +_PROXY_REQUESTER_METADATA: Final = {"metadata": {"requester_metadata": {"env": "prod"}}} + + +@pytest.mark.parametrize( + ("request_metadata", "sent_requester_metadata"), + [ + ({"requester_metadata": {"team": "t1"}}, {"env": "prod", "team": "t1"}), + ({"requester_metadata": {"env": "dev", "team": "t1"}}, {"env": "dev", "team": "t1"}), + ({"requester_metadata": None}, {"env": "prod"}), + ({"requester_metadata": {}}, {"env": "prod"}), + ({"requester_metadata": "x"}, {"env": "prod"}), + ({}, {"env": "prod"}), + ], +) +async def test_default_requester_metadata_merges_under_the_request_requester_metadata( + make_core: MakeCore, + host: ScriptedHost, + request_metadata: Mapping[str, object], + sent_requester_metadata: Mapping[str, str], +) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")], default_litellm_params=_PROXY_REQUESTER_METADATA) + await core.acompletion(model="g", messages=MESSAGES, metadata=dict(request_metadata)) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + assert _metadata(host.calls[0].kwargs)["requester_metadata"] == sent_requester_metadata + + +async def test_responses_default_requester_metadata_merges_under_the_request_litellm_metadata( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")], default_litellm_params=_PROXY_REQUESTER_METADATA) + await core.aresponses( # pyright: ignore[reportAttributeAccessIssue] # endpoint method + model="g", input="hi", litellm_metadata={"requester_metadata": {"team": "t1"}} + ) + sent: Final = _metadata(host.calls[0].kwargs, "litellm_metadata") + assert sent["requester_metadata"] == {"env": "prod", "team": "t1"} + + +async def test_default_requester_metadata_does_not_carry_one_request_labels_into_the_next( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")], default_litellm_params=_PROXY_REQUESTER_METADATA) + await core.acompletion(model="g", messages=MESSAGES, metadata={"requester_metadata": {"team": "t1"}}) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + assert _metadata(host.calls[1].kwargs)["requester_metadata"] == {"env": "prod"} + + +@pytest.mark.parametrize( + ("request_metadata", "sent_regions"), + [({"regions": ["eu"]}, ["eu"]), ({"regions": None}, ["us"]), ({}, ["us"])], +) +async def test_a_default_metadata_list_applies_only_when_the_request_sets_none( + make_core: MakeCore, host: ScriptedHost, request_metadata: Mapping[str, object], sent_regions: list[str] +) -> None: + core: Final = make_core( + model_list=[deployment("g", "d1")], default_litellm_params={"metadata": {"regions": ["us"]}} + ) + await core.acompletion(model="g", messages=MESSAGES, metadata=dict(request_metadata)) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + assert _metadata(host.calls[0].kwargs)["regions"] == sent_regions + + +@pytest.mark.parametrize("mask_key", ["mask_input", "mask_output"]) +@pytest.mark.parametrize( + ("default_mask", "request_mask", "sent_mask"), + [ + (True, False, True), + (True, None, True), + ("true", False, "true"), + (False, True, True), + (False, None, False), + ("false", True, True), + ], +) +async def test_default_metadata_masking_stays_on_once_either_side_turns_it_on( + make_core: MakeCore, + host: ScriptedHost, + mask_key: str, + default_mask: object, + request_mask: object, + sent_mask: object, +) -> None: + core: Final = make_core( + model_list=[deployment("g", "d1")], default_litellm_params={"metadata": {mask_key: default_mask}} + ) + await core.acompletion(model="g", messages=MESSAGES, metadata={mask_key: request_mask}) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + assert _metadata(host.calls[0].kwargs)[mask_key] == sent_mask + + +@pytest.mark.parametrize( + ("default_public", "request_public", "sent_public"), + [ + (False, True, False), + ("false", True, "false"), + (False, None, False), + (True, False, False), + (True, None, True), + ("true", True, True), + ], +) +async def test_a_private_default_trace_stays_private_whatever_the_request_sends( + make_core: MakeCore, + host: ScriptedHost, + default_public: object, + request_public: object, + sent_public: object, +) -> None: + core: Final = make_core( + model_list=[deployment("g", "d1")], default_litellm_params={"metadata": {"trace_public": default_public}} + ) + await core.acompletion(model="g", messages=MESSAGES, metadata={"trace_public": request_public}) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + assert _metadata(host.calls[0].kwargs)["trace_public"] == sent_public + + +async def test_a_default_redis_namespace_beats_the_request_redis_namespace( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=[deployment("g", "d1")], default_litellm_params={"metadata": {"redis_namespace": "tenant-a"}} + ) + await core.acompletion(model="g", messages=MESSAGES, metadata={"redis_namespace": "tenant-b"}) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + assert _metadata(host.calls[0].kwargs)["redis_namespace"] == "tenant-a" + + +def _redact(_value: object) -> str: + return "redacted" + + +@pytest.mark.parametrize("request_value", ["", {"caller": "value"}]) +async def test_a_default_masking_function_beats_whatever_the_request_sends( + make_core: MakeCore, host: ScriptedHost, request_value: object +) -> None: + core: Final = make_core( + model_list=[deployment("g", "d1")], default_litellm_params={"metadata": {"langfuse_masking_function": _redact}} + ) + await core.acompletion(model="g", messages=MESSAGES, metadata={"langfuse_masking_function": request_value}) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + assert _metadata(host.calls[0].kwargs)["langfuse_masking_function"] is _redact + + +_PROXY_DEFAULT_TAGS: Final = {"metadata": {"tags": ["proxy-default"]}} + + +@pytest.mark.parametrize( + ("request_metadata", "sent_tags"), + [ + ({"tags": ["caller"]}, ["caller", "proxy-default"]), + ({"tags": ["proxy-default", "caller"]}, ["proxy-default", "caller"]), + ({"tags": None}, ["proxy-default"]), + ({}, ["proxy-default"]), + ({"tags": "caller"}, ["proxy-default"]), + ], +) +async def test_default_metadata_tags_follow_the_request_tags( + make_core: MakeCore, host: ScriptedHost, request_metadata: Mapping[str, object], sent_tags: object +) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")], default_litellm_params=_PROXY_DEFAULT_TAGS) + await core.acompletion(model="g", messages=MESSAGES, metadata=dict(request_metadata)) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + assert _metadata(host.calls[0].kwargs)["tags"] == sent_tags + + +@pytest.mark.parametrize( + ("request_metadata", "sent_tags"), + [ + ({"tags": ["caller"]}, ["caller"]), + ({}, "proxy-default"), + ], +) +async def test_a_default_tags_that_is_not_a_list_yields_to_the_request_tags( + make_core: MakeCore, host: ScriptedHost, request_metadata: Mapping[str, object], sent_tags: object +) -> None: + core: Final = make_core( + model_list=[deployment("g", "d1")], default_litellm_params={"metadata": {"tags": "proxy-default"}} + ) + await core.acompletion(model="g", messages=MESSAGES, metadata=dict(request_metadata)) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + assert _metadata(host.calls[0].kwargs)["tags"] == sent_tags + + +@pytest.mark.parametrize( + ("request_metadata", "sent_tags"), + [ + ({}, ["t1", "Credential: cred", "proxy-default"]), + ({"tags": ["caller"]}, ["caller", "t1", "Credential: cred", "proxy-default"]), + ], +) +async def test_default_metadata_tags_join_the_deployment_and_credential_tags( + make_core: MakeCore, host: ScriptedHost, request_metadata: Mapping[str, object], sent_tags: list[str] +) -> None: + core: Final = make_core( + model_list=[deployment("g", "d1", tags=["t1"], litellm_credential_name="cred")], + default_litellm_params=_PROXY_DEFAULT_TAGS, + ) + await core.acompletion(model="g", messages=MESSAGES, metadata=dict(request_metadata)) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + assert _metadata(host.calls[0].kwargs)["tags"] == sent_tags + + +async def test_default_metadata_tags_do_not_carry_one_request_tags_into_the_next( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")], default_litellm_params=_PROXY_DEFAULT_TAGS) + await core.acompletion(model="g", messages=MESSAGES, metadata={"tags": ["caller"]}) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + assert _metadata(host.calls[1].kwargs)["tags"] == ["proxy-default"] + + +@pytest.mark.parametrize( + ("request_metadata", "sent_guardrails"), + [ + ({"guardrails": []}, ["admin-guard"]), + ({"guardrails": ["caller-guard"]}, ["caller-guard", "admin-guard"]), + ({"guardrails": "caller-guard"}, ["admin-guard"]), + ({}, ["admin-guard"]), + ({"guardrails": [{"admin-guard": {"extra_body": {"threshold": 1}}}]}, ["admin-guard"]), + ( + { + "guardrails": [ + {"admin-guard": {"extra_body": {"threshold": 1}}, "caller-guard": {"extra_body": {"k": 2}}} + ] + }, + [{"caller-guard": {"extra_body": {"k": 2}}}, "admin-guard"], + ), + ], +) +async def test_a_default_guardrail_stays_on_with_its_own_settings_whatever_the_request_lists( + make_core: MakeCore, host: ScriptedHost, request_metadata: Mapping[str, object], sent_guardrails: list[object] +) -> None: + core: Final = make_core( + model_list=[deployment("g", "d1")], default_litellm_params={"metadata": {"guardrails": ["admin-guard"]}} + ) + await core.acompletion(model="g", messages=MESSAGES, metadata=dict(request_metadata)) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + assert _metadata(host.calls[0].kwargs)["guardrails"] == sent_guardrails + + +async def test_a_default_guardrail_entry_keeps_the_admins_settings_over_the_callers( + make_core: MakeCore, host: ScriptedHost +) -> None: + admin_entry: Final = {"admin-guard": {"extra_body": {"threshold": 0.9}}} + core: Final = make_core( + model_list=[deployment("g", "d1")], default_litellm_params={"metadata": {"guardrails": [admin_entry]}} + ) + await core.acompletion( # pyright: ignore[reportAttributeAccessIssue] # endpoint method + model="g", messages=MESSAGES, metadata={"guardrails": [{"admin-guard": {"extra_body": {"threshold": 0}}}]} + ) + assert _metadata(host.calls[0].kwargs)["guardrails"] == [admin_entry] + + +@pytest.mark.parametrize("router_caching_groups", [[("g", "other")], None]) +async def test_the_router_caching_groups_replace_the_request_caching_groups( + make_core: MakeCore, host: ScriptedHost, router_caching_groups: list[tuple[str, str]] | None +) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")], caching_groups=router_caching_groups) + await core.acompletion(model="g", messages=MESSAGES, metadata={"caching_groups": [("g", "premium")]}) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + assert _metadata(host.calls[0].kwargs)["caching_groups"] == router_caching_groups + + # --- sync completion runs the routing pre-call checks for a model group only ------------------- @@ -412,7 +663,7 @@ async def test_acompletion_counts_total_success_and_failed_calls_per_deployment_ make_core: MakeCore, host: ScriptedHost ) -> None: core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=0) - await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] # endpoint method host.answer("openai/d1", rate_limit()) with pytest.raises(litellm.RateLimitError): await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] diff --git a/ui/litellm-dashboard/eslint.config.mjs b/ui/litellm-dashboard/eslint.config.mjs index acf6f624f96..828c0830cd0 100644 --- a/ui/litellm-dashboard/eslint.config.mjs +++ b/ui/litellm-dashboard/eslint.config.mjs @@ -5,8 +5,157 @@ import prettier from "eslint-config-prettier/flat"; import unusedImports from "eslint-plugin-unused-imports"; import testingLibrary from "eslint-plugin-testing-library"; import jestDom from "eslint-plugin-jest-dom"; +import i18next from "eslint-plugin-i18next"; +import htmlEntities from "eslint-plugin-i18next/lib/options/htmlEntities.js"; import local from "./scripts/eslint-rules/index.mjs"; +const rawFetchRestriction = { + selector: "CallExpression[callee.name='fetch']", + message: + "Raw fetch() is only allowed in src/lib/http/. Use the shared client (createApiClient / apiClient) from @/lib/http/client instead.", +}; + +const toastLiteralRestriction = { + selector: + "CallExpression[callee.object.name='toast'][callee.property.name=/^(success|error|warning|info)$/] > Literal.arguments:first-child", + message: "Pass toast a t() string, not an English literal, so the message can be translated.", +}; + +const noLiteralStringOptions = { + mode: "jsx-only", + "should-validate-template": true, + "jsx-components": { exclude: ["Trans", "code", "pre", "kbd", "Kbd", "CodeBlock", "SyntaxHighlighter"] }, + "jsx-attributes": { + exclude: [ + "className", + ".*ClassName", + "style", + ".*Style", + "key", + "id", + ".*Id", + ".*ID", + "htmlFor", + "name", + "value", + "defaultValue", + "data-.*", + "aria-(describedby|labelledby|controls|owns|haspopup|live|current|hidden|expanded|selected)", + "type", + "role", + "href", + "to", + "src", + "target", + "rel", + "accept", + "autoComplete", + "inputMode", + "step", + "min", + "max", + "width", + "height", + "size", + "color", + ".*Color", + "variant", + "language", + "index", + ".*Key", + "keys", + "feature", + "fill", + "stroke", + "scope", + "delay", + "stackId", + "category", + "categories", + "gridCols", + "defaultTheme", + "defaultRole", + "schemaComponent", + "storageKey", + "code", + ], + }, + callees: { + exclude: [ + "i18n(ext)?", + "t", + "i18n\\.t", + "require", + "addEventListener", + "removeEventListener", + "postMessage", + "getElementById", + "dispatch", + "commit", + "includes", + "indexOf", + "endsWith", + "startsWith", + "cn", + "clsx", + "cva", + ".*\\.(get|set|has|delete|getAll|append|getColumn|format|split|replace|closest|querySelector|getItem|setItem|removeItem|toLocaleDateString|toLocaleString|hasOwnProperty)", + "Object\\.hasOwn", + "uiHref", + "hasVisited", + "toggleKey", + "emitLocalStorageChange", + "setLocalStorageItem", + "removeLocalStorageItem", + "getLocalStorageItem", + "window\\.open", + "Error", + "new Error", + ], + }, + words: { + exclude: [ + "[0-9!-/:-@[-`{-~]+", + "[A-Z_-]+", + "v", + htmlEntities, + /^\p{Emoji}+$/u, + /^[\p{P}\p{S}\p{N}\p{M}\s]+$/u, + /^[a-z0-9]+(?:[_\-.:/][a-z0-9*[\]]+)+:?$/, + /^(?:https?:\/\/|\/|~\/|\.\/)\S*$/, + ], + }, + "object-properties": { + exclude: [ + "[A-Z_-]+", + "value", + "key", + "id", + "name", + "dataIndex", + "className", + "width", + "fontSize", + "maxHeight", + "borderRadius", + "bg", + "color", + "background", + "padding", + "fill", + "role", + "class", + "spellcheck", + "toast", + "scope", + "origin", + "classifierType", + "defaultValue", + "1", + ], + }, +}; + const eslintConfig = [ { ignores: [".next/**", "out/**", "build/**", "coverage/**", "next-env.d.ts", "src/lib/http/schema.d.ts"], @@ -41,14 +190,7 @@ const eslintConfig = [ "max-depth": ["warn", 4], "max-params": ["error", 4], "max-nested-callbacks": ["error", 4], - "no-restricted-syntax": [ - "error", - { - selector: "CallExpression[callee.name='fetch']", - message: - "Raw fetch() is only allowed in src/lib/http/. Use the shared client (createApiClient / apiClient) from @/lib/http/client instead.", - }, - ], + "no-restricted-syntax": ["error", rawFetchRestriction], "no-restricted-imports": [ "error", { @@ -81,6 +223,26 @@ const eslintConfig = [ "max-lines": ["error", { max: 800, skipBlankLines: true, skipComments: true }], }, }, + { + files: ["src/**/*.{ts,tsx}"], + ignores: [ + "src/**/*.test.{ts,tsx}", + "src/**/*.spec.{ts,tsx}", + "src/**/*.test-d.ts", + "src/**/__tests__/**", + "src/**/*TestHarness.tsx", + ], + languageOptions: { parserOptions: { projectService: true, tsconfigRootDir: import.meta.dirname } }, + plugins: { i18next }, + rules: { + "i18next/no-literal-string": ["error", noLiteralStringOptions], + "no-restricted-syntax": ["error", rawFetchRestriction, toastLiteralRestriction], + }, + }, + { + files: ["src/i18n/untranslated/**"], + rules: { "i18next/no-literal-string": "off" }, + }, { files: ["src/lib/http/**"], rules: { diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index d506d114d8c..5c530c4ac80 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -77,6 +77,7 @@ "eslint": "9.39.2", "eslint-config-next": "16.3.3", "eslint-config-prettier": "10.1.8", + "eslint-plugin-i18next": "6.1.5", "eslint-plugin-jest-dom": "5.10.1", "eslint-plugin-testing-library": "7.16.2", "eslint-plugin-unused-imports": "4.3.0", @@ -6647,6 +6648,29 @@ "ms": "^2.1.1" } }, + "node_modules/eslint-plugin-i18next": { + "version": "6.1.5", + "resolved": "https://registry.npmjs.org/eslint-plugin-i18next/-/eslint-plugin-i18next-6.1.5.tgz", + "integrity": "sha512-xCTfstbK9ZpQ6UFT5S1s6zbZMhKn2o2jRSQYiU2UkV7wt8y1m3WjkUYGkGlipgRdFInYI9+LAqE+JQU/L73HmA==", + "dev": true, + "license": "ISC", + "dependencies": { + "requireindex": "~1.1.0" + }, + "engines": { + "node": ">=18.10.0" + } + }, + "node_modules/eslint-plugin-i18next/node_modules/requireindex": { + "version": "1.1.0", + "resolved": "https://registry.npmjs.org/requireindex/-/requireindex-1.1.0.tgz", + "integrity": "sha512-LBnkqsDE7BZKvqylbmn7lTIVdpx4K/QCduRATpO5R+wtPmky/a8pN1bO2D6wXppn1497AJF9mNjqAXr6bdl9jg==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=0.10.5" + } + }, "node_modules/eslint-plugin-import": { "version": "2.32.0", "resolved": "https://registry.npmjs.org/eslint-plugin-import/-/eslint-plugin-import-2.32.0.tgz", diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index 9fc0242f040..2174d7cc2ad 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -94,6 +94,7 @@ "eslint": "9.39.2", "eslint-config-next": "16.3.3", "eslint-config-prettier": "10.1.8", + "eslint-plugin-i18next": "6.1.5", "eslint-plugin-jest-dom": "5.10.1", "eslint-plugin-testing-library": "7.16.2", "eslint-plugin-unused-imports": "4.3.0", diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index e2f932b44a2..634ef66f36f 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -8683,7 +8683,7 @@ export interface paths { * - 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"}]. @@ -8998,7 +8998,7 @@ export interface paths { * - 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 @@ -9150,7 +9150,7 @@ export interface paths { * - 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"}]. @@ -11715,6 +11715,64 @@ export interface paths { patch: operations["openrouter_proxy_route_openrouter__endpoint__patch"]; trace?: never; }; + "/opensearch/{endpoint}": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Opensearch Proxy Route + * @description Call an OpenSearch index through the proxy, e.g. `POST /opensearch//_search`. + * + * `` 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. + */ + get: operations["opensearch_proxy_route_opensearch__endpoint__get"]; + /** + * Opensearch Proxy Route + * @description Call an OpenSearch index through the proxy, e.g. `POST /opensearch//_search`. + * + * `` 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. + */ + put: operations["opensearch_proxy_route_opensearch__endpoint__put"]; + /** + * Opensearch Proxy Route + * @description Call an OpenSearch index through the proxy, e.g. `POST /opensearch//_search`. + * + * `` 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. + */ + post: operations["opensearch_proxy_route_opensearch__endpoint__post"]; + /** + * Opensearch Proxy Route + * @description Call an OpenSearch index through the proxy, e.g. `POST /opensearch//_search`. + * + * `` 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. + */ + delete: operations["opensearch_proxy_route_opensearch__endpoint__delete"]; + options?: never; + /** + * Opensearch Proxy Route + * @description Call an OpenSearch index through the proxy, e.g. `POST /opensearch//_search`. + * + * `` 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. + */ + head: operations["opensearch_proxy_route_opensearch__endpoint__head"]; + /** + * Opensearch Proxy Route + * @description Call an OpenSearch index through the proxy, e.g. `POST /opensearch//_search`. + * + * `` 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. + */ + patch: operations["opensearch_proxy_route_opensearch__endpoint__patch"]; + trace?: never; + }; "/organization/daily/activity": { parameters: { query?: never; @@ -17853,7 +17911,7 @@ export interface paths { * - 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"]. @@ -18089,7 +18147,7 @@ export interface paths { * - 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"]. @@ -26181,6 +26239,17 @@ export interface components { /** Updated By */ updated_by?: string | null; }; + /** AccessGroupRule */ + AccessGroupRule: { + /** Access Group */ + access_group: string; + /** Attribute */ + attribute: string; + /** Contains */ + contains?: string | number | boolean | null; + /** Equals */ + equals?: string | number | boolean | null; + }; /** AccessGroupUpdateRequest */ AccessGroupUpdateRequest: { /** Access Agent Ids */ @@ -30750,6 +30819,8 @@ export interface components { * @description Documents all the fields supported by `general_settings` in config.yaml */ ConfigGeneralSettings: { + /** Access Group Rules */ + access_group_rules?: components["schemas"]["AccessGroupRule"][] | null; /** * Admission Queue Timeout Seconds * @description maximum time a request waits for a worker slot @@ -43961,7 +44032,7 @@ export interface components { heuristic_v2_success_threshold?: number | null; /** * Housekeeping Patterns - * @description Additional case-sensitive literal sentinels that mark a request as client housekeeping, on top of the built-in conversation-title ones. For clients whose wording the built-ins don't cover, or after a client release changes its strings. + * @description Additional case-sensitive literal sentinels that mark a request as client housekeeping, on top of the built-in conversation-title, agent-summary and prompt-suggestion ones. For clients whose wording the built-ins don't cover, or after a client release changes its strings. */ housekeeping_patterns?: string[] | null; /** @@ -52301,6 +52372,8 @@ export interface components { discoverable?: boolean | null; /** Enable Tag Filtering */ enable_tag_filtering?: boolean | null; + /** Encrypted Content Boundary */ + encrypted_content_boundary?: string | null; /** Enforce Supported Endpoints */ enforce_supported_endpoints?: boolean | null; /** Explicit Cache Read Input Token Cost */ @@ -70085,6 +70158,246 @@ export interface operations { }; }; }; + opensearch_proxy_route_opensearch__endpoint__get: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: string; + }; + cookie?: never; + }; + requestBody?: { + content: { + "application/json": { + [key: string]: unknown; + }; + "application/x-ndjson": string; + }; + }; + responses: { + /** @description OpenSearch's response, unmodified */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": { + [key: string]: unknown; + }; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + opensearch_proxy_route_opensearch__endpoint__put: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: string; + }; + cookie?: never; + }; + requestBody?: { + content: { + "application/json": { + [key: string]: unknown; + }; + "application/x-ndjson": string; + }; + }; + responses: { + /** @description OpenSearch's response, unmodified */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": { + [key: string]: unknown; + }; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + opensearch_proxy_route_opensearch__endpoint__post: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: string; + }; + cookie?: never; + }; + requestBody?: { + content: { + "application/json": { + [key: string]: unknown; + }; + "application/x-ndjson": string; + }; + }; + responses: { + /** @description OpenSearch's response, unmodified */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": { + [key: string]: unknown; + }; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + opensearch_proxy_route_opensearch__endpoint__delete: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: string; + }; + cookie?: never; + }; + requestBody?: { + content: { + "application/json": { + [key: string]: unknown; + }; + "application/x-ndjson": string; + }; + }; + responses: { + /** @description OpenSearch's response, unmodified */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": { + [key: string]: unknown; + }; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + opensearch_proxy_route_opensearch__endpoint__head: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: string; + }; + cookie?: never; + }; + requestBody?: { + content: { + "application/json": { + [key: string]: unknown; + }; + "application/x-ndjson": string; + }; + }; + responses: { + /** @description OpenSearch's response, unmodified */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": { + [key: string]: unknown; + }; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + opensearch_proxy_route_opensearch__endpoint__patch: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: string; + }; + cookie?: never; + }; + requestBody?: { + content: { + "application/json": { + [key: string]: unknown; + }; + "application/x-ndjson": string; + }; + }; + responses: { + /** @description OpenSearch's response, unmodified */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": { + [key: string]: unknown; + }; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; get_organization_daily_activity_organization_daily_activity_get: { parameters: { query?: { diff --git a/ui/litellm-dashboard/tests/fetch-location-rule.test.ts b/ui/litellm-dashboard/tests/fetch-location-rule.test.ts index bd495ad3c40..329012724e8 100644 --- a/ui/litellm-dashboard/tests/fetch-location-rule.test.ts +++ b/ui/litellm-dashboard/tests/fetch-location-rule.test.ts @@ -18,7 +18,7 @@ describe("location-based fetch() rule", () => { }); it("flags a raw fetch() in a normal source file", async () => { - const messages = await fetchMessages("src/components/some_feature.tsx"); + const messages = await fetchMessages("src/app/layout.tsx"); expect(messages).toHaveLength(1); expect(messages[0].message).toMatch(/@\/lib\/http\/client/); }); diff --git a/uv.lock b/uv.lock index 8712b78bbd6..62187d596ec 100644 --- a/uv.lock +++ b/uv.lock @@ -1117,6 +1117,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/bf/0f/f897abe4ea0a8c408ae65c8c83bffab4936ad65d6032d4fb4cd35bbdc3ee/cachetools-7.1.1-py3-none-any.whl", hash = "sha256:0335cd7a0952d2b22327441fb0628139e234c565559eeb91a8a4ac7551c5353d", size = 16775, upload-time = "2026-05-03T20:00:27.857Z" }, ] +[[package]] +name = "casbin" +version = "1.43.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "simpleeval" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ad/df/ff2aa55cf0d7c14622ce4f9252cdc34c828c81d4213965d73207ac5434ae/casbin-1.43.0.tar.gz", hash = "sha256:d2e90ce8e72f912877851e94d37999f32c558c6ba7aba0437d483275262e86e0", size = 425727, upload-time = "2025-05-10T06:57:18.902Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/08/07/facef6abd81378e6153b757dc1848621675d971fbc88ebb5d182ebc1c37f/casbin-1.43.0-py3-none-any.whl", hash = "sha256:63a3d1228870250e859ccd94133fe478821093f71dd37f05e0baa0c6fea26623", size = 475059, upload-time = "2025-05-10T06:57:16.89Z" }, +] + [[package]] name = "certifi" version = "2026.10.19" @@ -5096,6 +5108,7 @@ proxy = [ { name = "backoff" }, { name = "backports-zstd", marker = "python_full_version < '3.14'" }, { name = "boto3" }, + { name = "casbin" }, { name = "cryptography" }, { name = "expression" }, { name = "fastapi" }, @@ -5312,6 +5325,7 @@ requires-dist = [ { name = "backports-zstd", marker = "python_full_version < '3.14' and extra == 'proxy'", specifier = ">=1.0.0,<2.0" }, { name = "boto3", specifier = ">=1.43.117,<2.0" }, { name = "boto3", marker = "extra == 'proxy'", specifier = ">=1.43.117,<2.0" }, + { name = "casbin", marker = "extra == 'proxy'", specifier = "==1.43.0" }, { name = "certifi", specifier = ">=2026.10.19" }, { name = "click", specifier = ">=8.0.0,<9.0" }, { name = "cryptography", marker = "extra == 'proxy'", specifier = ">=49.0.0,<51.0" }, @@ -10377,6 +10391,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/52/59/0782e51887ac6b07ffd1570e0364cf901ebc36345fea669969d2084baebb/simple_websocket-1.1.0-py3-none-any.whl", hash = "sha256:4af6069630a38ed6c561010f0e11a5bc0d4ca569b36306eb257cd9a192497c8c", size = 13842, upload-time = "2024-10-10T22:39:29.645Z" }, ] +[[package]] +name = "simpleeval" +version = "1.0.8" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/89/3e/25326f373871a34e78cac231c0d3c18a0a5c3a80aa4813b50b9f26a5d460/simpleeval-1.0.8.tar.gz", hash = "sha256:b417f869c11f2de034eb919820a825f35d12ce99f63a43971947ce73679eb9ba", size = 31846, upload-time = "2026-09-12T10:16:08.103Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/12/33/d1ecd9f4f5f272b36ad7b268e307ba0a59c109d8ad134a231cda780baac9/simpleeval-1.0.8-py3-none-any.whl", hash = "sha256:c9c4ff42879ebf0c05f6562cb81031869948916be343bcc950b81fcc081fe9ce", size = 18925, upload-time = "2026-09-12T10:16:06.862Z" }, +] + [[package]] name = "six" version = "1.17.0"