diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260916000000_add_key_total_spend/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260916000000_add_key_total_spend/migration.sql new file mode 100644 index 00000000000..daacd66db39 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260916000000_add_key_total_spend/migration.sql @@ -0,0 +1,5 @@ +-- AlterTable +ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "total_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0; + +-- AlterTable +ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN IF NOT EXISTS "total_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index d2375903c47..139fb031671 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -426,6 +426,7 @@ model LiteLLM_VerificationToken { key_alias String? soft_budget_cooldown Boolean @default(false) // key-level state on if budget alerts need to be cooled down spend Float @default(0.0) + total_spend Float @default(0.0) expires DateTime? models String[] aliases Json @default("{}") @@ -528,6 +529,7 @@ model LiteLLM_DeletedVerificationToken { key_alias String? soft_budget_cooldown Boolean @default(false) spend Float @default(0.0) + total_spend Float @default(0.0) expires DateTime? models String[] aliases Json @default("{}") diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index e2a3af77594..7397742369b 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1951,6 +1951,7 @@ dependencies = [ name = "litellm-python-bridge" version = "0.1.0" dependencies = [ + "bytes", "criterion", "futures-util", "litellm-auth", diff --git a/litellm-rust/crates/core/src/ocr/client.rs b/litellm-rust/crates/core/src/ocr/client.rs index 00bfeb2b7b2..9a30b2f8e04 100644 --- a/litellm-rust/crates/core/src/ocr/client.rs +++ b/litellm-rust/crates/core/src/ocr/client.rs @@ -54,9 +54,16 @@ impl OcrClient { match call.resume(result.take()).await? { OcrCallStep::Host(OcrHostOperation::ProjectRequest) => { result = Some(OcrHostResult::Request(Ok(( - Box::new(request.take().ok_or_else(|| { - Error::InvalidRequest("OCR request was already projected".into()) - })?), + Box::new( + request + .take() + .ok_or_else(|| { + Error::InvalidRequest( + "OCR request was already projected".into(), + ) + })? + .into(), + ), false, )))) } diff --git a/litellm-rust/crates/core/src/ocr/document.rs b/litellm-rust/crates/core/src/ocr/document.rs index a7afdaf8793..1b3d2dada44 100644 --- a/litellm-rust/crates/core/src/ocr/document.rs +++ b/litellm-rust/crates/core/src/ocr/document.rs @@ -1,3 +1,6 @@ +use std::io::Read; +use std::path::Path; + use base64::{Engine, engine::general_purpose::STANDARD}; use data_url::mime::Mime; use data_url::{DataUrl, DataUrlError, forgiving_base64::DecodeError}; @@ -5,12 +8,52 @@ use reqwest::Url; use serde_json::Map; use super::error::{OcrError, OcrRequestError, OcrResponseError}; -use super::types::{OcrConnection, OcrDocument}; +use super::types::{OcrConnection, OcrDocument, OcrDocumentInput}; use crate::constants::{OCR_INLINE_MAX_BYTES, OCR_MAX_FETCH_REDIRECTS}; use crate::media::Error as MediaError; use crate::media::{DownloadPolicy, MediaFetcher}; use crate::transport::Error as TransportError; +pub fn prepare_document(input: OcrDocumentInput) -> Result { + match input { + OcrDocumentInput::Document(document) => Ok(document), + OcrDocumentInput::Path { path, mime_type } => { + read_path_document(&path, mime_type.as_deref()) + } + OcrDocumentInput::Bytes { + bytes, + file_name, + mime_type, + } => Ok(encode_file_document( + &bytes, + file_name.as_deref(), + mime_type.as_deref(), + )?), + OcrDocumentInput::HostReader { .. } => Err(super::Error::InvalidRequest( + "OCR file reader was not read by the host".into(), + )), + } +} + +pub fn read_path_document( + path: &Path, + mime_type: Option<&str>, +) -> Result { + let mut bytes = Vec::new(); + std::fs::File::open(path) + .and_then(|file| { + file.take(OCR_INLINE_MAX_BYTES as u64 + 1) + .read_to_end(&mut bytes) + }) + .map_err(|source| super::Error::FileRead { + path: path.to_owned(), + kind: source.kind(), + message: source.to_string(), + })?; + let name = path.file_name().map(|name| name.to_string_lossy()); + Ok(encode_file_document(&bytes, name.as_deref(), mime_type)?) +} + pub fn encode_file_document( bytes: &[u8], file_name: Option<&str>, @@ -75,18 +118,6 @@ pub fn mime_type_for_name(name: &str) -> &'static str { } } -pub fn upload_mime_type<'a>(file_name: Option<&str>, content_type: Option<&'a str>) -> &'a str { - match content_type - .and_then(|value| value.split(';').next()) - .map(str::trim) - { - Some(value) if !value.is_empty() && value != "application/octet-stream" => value, - _ => file_name - .map(mime_type_for_name) - .unwrap_or("application/octet-stream"), - } -} - pub(crate) struct InlineDocument<'a>(DataUrl<'a>); impl<'a> InlineDocument<'a> { @@ -230,24 +261,65 @@ mod tests { } #[test] - fn upload_mime_mapping_matches_python() { + fn path_documents_are_read_and_named_by_core() { + let dir = std::env::temp_dir().join(format!("litellm-ocr-{}", rand::random::())); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.join("scan.png"); + std::fs::write(&path, b"abc").unwrap(); assert_eq!( - upload_mime_type(Some("report.pdf"), Some("application/octet-stream")), - "application/pdf" - ); - assert_eq!(upload_mime_type(Some("image.png"), None), "image/png"); - assert_eq!(upload_mime_type(None, None), "application/octet-stream"); - assert_eq!( - upload_mime_type(Some("doc.pdf"), Some("application/pdf; charset=utf-8")), - "application/pdf" + prepare_document(OcrDocumentInput::Path { + path: path.clone(), + mime_type: None, + }) + .unwrap(), + OcrDocument::ImageUrl { + image_url: "data:image/png;base64,YWJj".into(), + extra_fields: Map::new(), + } ); assert_eq!( - upload_mime_type( - Some("img.png"), - Some("image/png; charset=utf-8; boundary=something") - ), - "image/png" + prepare_document(OcrDocumentInput::Path { + path: path.clone(), + mime_type: Some("application/pdf".into()), + }) + .unwrap(), + document("data:application/pdf;base64,YWJj") ); + std::fs::write(&path, vec![b'a'; OCR_INLINE_MAX_BYTES + 1]).unwrap(); + assert_eq!( + prepare_document(OcrDocumentInput::Path { + path: path.clone(), + mime_type: None, + }), + Err(OcrRequestError::InlineDocumentTooLarge.into()) + ); + std::fs::remove_dir_all(&dir).unwrap(); + + let missing = dir.join("missing.pdf"); + let Err(super::super::Error::FileRead { path, kind, .. }) = + prepare_document(OcrDocumentInput::Path { + path: missing.clone(), + mime_type: None, + }) + else { + panic!("missing paths must surface a file read error"); + }; + assert_eq!(path, missing); + assert_eq!(kind, std::io::ErrorKind::NotFound); + } + + #[test] + fn byte_documents_are_encoded_and_host_readers_must_be_read_first() { + assert_eq!( + prepare_document(OcrDocumentInput::Bytes { + bytes: b"abc".as_slice().into(), + file_name: Some("scan.pdf".into()), + mime_type: None, + }) + .unwrap(), + document("data:application/pdf;base64,YWJj") + ); + assert!(prepare_document(OcrDocumentInput::HostReader { mime_type: None }).is_err()); } #[test] diff --git a/litellm-rust/crates/core/src/ocr/error.rs b/litellm-rust/crates/core/src/ocr/error.rs index 1c21edb6c91..0c92b511a38 100644 --- a/litellm-rust/crates/core/src/ocr/error.rs +++ b/litellm-rust/crates/core/src/ocr/error.rs @@ -50,6 +50,12 @@ pub enum Error { Connect(String), #[error("routing error: {0}")] Routing(String), + #[error("Failed to read OCR file {}: {message}", path.display())] + FileRead { + path: std::path::PathBuf, + kind: std::io::ErrorKind, + message: String, + }, /// The request is outside the surface this route covers in Rust. Hosts that /// keep a reference implementation treat this as "fall back", not "fail". #[error("unsupported by the rust path: {0}")] diff --git a/litellm-rust/crates/core/src/ocr/lifecycle.rs b/litellm-rust/crates/core/src/ocr/lifecycle.rs index efa2b1f2873..994a9698459 100644 --- a/litellm-rust/crates/core/src/ocr/lifecycle.rs +++ b/litellm-rust/crates/core/src/ocr/lifecycle.rs @@ -9,6 +9,7 @@ use super::hooks::{ OcrDuringCallRequest, OcrHookFuture, OcrHooks, OcrLogFuture, OcrPostCallRequest, OcrPreCallRequest, }; +use super::types::{OcrDocumentInput, OcrFileContent}; use super::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrClient}; use crate::call_lifecycle::host::{ HostCall, HostCallFuture, HostCallStep, HostFailure, HostLifecycle, HostPhase, @@ -52,6 +53,7 @@ impl OcrAdmission { #[derive(Clone, Debug)] pub enum OcrHostOperation { ProjectRequest, + ReadDocument, Lifecycle(HostPhase), ConstructResponse(Arc), MapFailure(Error), @@ -83,7 +85,8 @@ impl OcrHostOperation { } pub enum OcrHostResult { - Request(Result<(Box, bool), Error>), + Request(Result<(Box>, bool), Error>), + Document(Result), Lifecycle(Result<(), HostFailure>), AzureAdToken(Result), PreCall(Result), @@ -313,7 +316,7 @@ struct PendingOperation { struct OcrExecution { client: Option, - request: Option, + request: Option>, operations_tx: mpsc::UnboundedSender, operations_rx: mpsc::UnboundedReceiver, pending_result: Option>, @@ -397,12 +400,14 @@ impl OcrExecution { }, ))); } - request.hooks = Arc::new(ProtocolHooks { + let hooks = Arc::new(ProtocolHooks { operations: self.operations_tx.clone(), intercepts_requests, terminal: self.terminal.clone(), }); + request.hooks = hooks.clone(); self.execution = Some(tokio::spawn(async move { + let request = prepare_request_document(request, &hooks).await?; perform_ocr_request(&client, request).await })); } @@ -423,6 +428,39 @@ impl OcrExecution { } } +async fn prepare_request_document( + request: LiteLLMOcrRequest, + hooks: &ProtocolHooks, +) -> Result { + let request = match &request.document { + OcrDocumentInput::HostReader { mime_type } => { + let mime_type = mime_type.clone(); + let content = match hooks.invoke(OcrHostOperation::ReadDocument).await? { + OcrHostResult::Document(result) => result?, + _ => { + return Err(Error::InvalidRequest( + "invalid OCR document read host result".into(), + )); + } + }; + request.with_document(OcrDocumentInput::Bytes { + bytes: content.bytes, + file_name: content.file_name, + mime_type, + }) + } + _ => request, + }; + if let OcrDocumentInput::Document(_) = &request.document { + return request.map_document(super::document::prepare_document); + } + tokio::task::spawn_blocking(move || request.map_document(super::document::prepare_document)) + .await + .map_err(|error| { + Error::InvalidRequest(format!("OCR document preparation task failed: {error}")) + })? +} + impl Drop for OcrExecution { fn drop(&mut self) { if let Some(execution) = &self.execution { @@ -567,6 +605,9 @@ impl OcrHost for NoopOcrHost { OcrHostOperation::ProjectRequest => OcrHostResult::Request(Err( Error::InvalidRequest("OCR host has no request projection".into()), )), + OcrHostOperation::ReadDocument => OcrHostResult::Document(Err( + Error::InvalidRequest("OCR host has no document reader".into()), + )), OcrHostOperation::Lifecycle(_) | OcrHostOperation::ConstructResponse(_) | OcrHostOperation::MapFailure(_) @@ -602,6 +643,9 @@ impl OcrHost for OcrHookHost { OcrHostOperation::ProjectRequest => OcrHostResult::Request(Err( Error::InvalidRequest("OCR hook host has no request projection".into()), )), + OcrHostOperation::ReadDocument => OcrHostResult::Document(Err( + Error::InvalidRequest("OCR hook host has no document reader".into()), + )), OcrHostOperation::Success { context, response, diff --git a/litellm-rust/crates/core/src/ocr/mod.rs b/litellm-rust/crates/core/src/ocr/mod.rs index 3b51ff98356..f2e7aa4f46d 100644 --- a/litellm-rust/crates/core/src/ocr/mod.rs +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -13,12 +13,15 @@ pub mod types; pub mod wire; pub use client::{OcrClient, ocr}; -pub use document::{encode_file_document, mime_type_for_name, upload_mime_type}; +pub use document::{encode_file_document, mime_type_for_name, read_path_document}; pub use lifecycle::{ NativeOutcome, NativeResult, NoopOcrHost, OcrAdmission, OcrCall, OcrCallStep, OcrDecline, OcrHookHost, OcrHost, OcrHostOperation, OcrHostResult, }; -pub use types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrDocument}; +pub use types::{ + LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrDocumentInput, + OcrFileContent, +}; #[cfg(test)] #[path = "../../tests/azure_ai_ocr.rs"] diff --git a/litellm-rust/crates/core/src/ocr/types.rs b/litellm-rust/crates/core/src/ocr/types.rs index 69e6982414b..bb212674b33 100644 --- a/litellm-rust/crates/core/src/ocr/types.rs +++ b/litellm-rust/crates/core/src/ocr/types.rs @@ -1,7 +1,10 @@ use std::collections::BTreeMap; +use std::convert::Infallible; +use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; +use bytes::Bytes; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; @@ -50,6 +53,35 @@ impl OcrDocument { } } +#[derive(Clone, Debug, PartialEq)] +pub enum OcrDocumentInput { + Document(OcrDocument), + Path { + path: PathBuf, + mime_type: Option, + }, + Bytes { + bytes: Bytes, + file_name: Option, + mime_type: Option, + }, + HostReader { + mime_type: Option, + }, +} + +impl From for OcrDocumentInput { + fn from(document: OcrDocument) -> Self { + Self::Document(document) + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct OcrFileContent { + pub bytes: Bytes, + pub file_name: Option, +} + #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "lowercase")] pub enum OcrResponseFormat { @@ -89,9 +121,9 @@ impl Default for OcrConnection { } } -pub struct LiteLLMOcrRequest { +pub struct LiteLLMOcrRequest { pub model: String, - pub document: OcrDocument, + pub document: D, pub connection: OcrConnection, pub hooks: Arc, pub litellm_call_id: Option, @@ -101,10 +133,10 @@ pub struct LiteLLMOcrRequest { pub(crate) adapter: OcrAdapterKind, } -impl LiteLLMOcrRequest { +impl LiteLLMOcrRequest { pub fn new( model: String, - document: OcrDocument, + document: D, custom_llm_provider: Option<&str>, optional_params: Map, ) -> Result { @@ -151,6 +183,36 @@ impl LiteLLMOcrRequest { ..self } } + + pub fn map_document( + self, + map: impl FnOnce(D) -> Result, + ) -> Result, E> { + Ok(LiteLLMOcrRequest { + model: self.model, + document: map(self.document)?, + connection: self.connection, + hooks: self.hooks, + litellm_call_id: self.litellm_call_id, + optional_params: self.optional_params, + input_sources: self.input_sources, + azure_ad_token_provider: self.azure_ad_token_provider, + adapter: self.adapter, + }) + } + + pub fn with_document(self, document: T) -> LiteLLMOcrRequest { + let Ok(request) = self.map_document(|_| Ok::(document)); + request + } +} + +impl From for LiteLLMOcrRequest { + fn from(request: LiteLLMOcrRequest) -> Self { + let Ok(request) = request + .map_document(|document| Ok::<_, Infallible>(OcrDocumentInput::Document(document))); + request + } } #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] diff --git a/litellm-rust/crates/core/src/ocr/wire.rs b/litellm-rust/crates/core/src/ocr/wire.rs index 93816effcb1..f0cad2b4e93 100644 --- a/litellm-rust/crates/core/src/ocr/wire.rs +++ b/litellm-rust/crates/core/src/ocr/wire.rs @@ -68,9 +68,9 @@ pub struct DecodedOcrResponse { #[derive(Deserialize)] #[serde(deny_unknown_fields)] -pub struct OcrWireRequest { +pub struct OcrWireRequest { pub model: String, - pub document: Value, + pub document: D, pub api_key: Option, pub api_base: Option, pub custom_llm_provider: Option, @@ -141,10 +141,34 @@ pub fn consumed_optional_params( } pub fn decode_request(wire: OcrWireRequest) -> Result { + let OcrWireRequest { + model, + document, + api_key, + api_base, + custom_llm_provider, + extra_headers, + optional_params, + input_sources, + timeout_seconds, + } = wire; + decode_request_input(OcrWireRequest { + model, + document: decode_document(document)?, + api_key, + api_base, + custom_llm_provider, + extra_headers, + optional_params, + input_sources, + timeout_seconds, + }) +} + +pub fn decode_request_input(wire: OcrWireRequest) -> Result, Error> { let api_key_source = source_for(&wire.input_sources, "api_key"); let api_base_source = source_for(&wire.input_sources, "api_base"); let extra_headers_source = source_for(&wire.input_sources, "extra_headers"); - let document = decode_document(wire.document)?; let headers = wire .extra_headers .unwrap_or_default() @@ -183,7 +207,7 @@ pub fn decode_request(wire: OcrWireRequest) -> Result .unwrap_or(defaults.max_response_bytes); let request = LiteLLMOcrRequest::new( wire.model, - document, + wire.document, wire.custom_llm_provider.as_deref(), wire.optional_params .into_iter() @@ -209,14 +233,14 @@ pub fn decode_request(wire: OcrWireRequest) -> Result }) } -fn decode_document(value: Value) -> Result { +pub fn decode_document(value: Value) -> Result { let kind = value.get("type").and_then(Value::as_str); let missing_url = matches!(kind, Some("document_url")) && value.get("document_url").is_none() || matches!(kind, Some("image_url")) && value.get("image_url").is_none(); if missing_url { - return Err(OcrRequestError::MissingDocumentUrl); + return Err(OcrRequestError::MissingDocumentUrl.into()); } - decode_request_value(value, "document") + Ok(decode_request_value(value, "document")?) } fn source_for(sources: &BTreeMap, name: &str) -> InputSource { @@ -334,10 +358,7 @@ mod tests { serde_json::json!({"type": "document_url"}), serde_json::json!({"type": "image_url"}), ] { - assert_eq!( - decode_document(document), - Err(OcrRequestError::MissingDocumentUrl) - ); + assert_eq!(decode_document(document), Err(Error::MissingDocumentUrl)); } } } diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index 373972cf68b..a24d960422d 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -348,13 +348,14 @@ async fn fallible_host_phases_do_not_replay_or_reach_transport() { } OcrHostOperation::ProjectRequest => { result = Some(OcrHostResult::Request(Ok(( - Box::new(request.take().unwrap()), + Box::new(request.take().unwrap().into()), false, )))) } OcrHostOperation::AcquireAzureAdToken => { panic!("test request has no token provider") } + OcrHostOperation::ReadDocument => panic!("test request has no file reader"), OcrHostOperation::PreCall(request) => { phases.push("pre"); result = Some(OcrHostResult::PreCall(if failure_phase == "pre" { @@ -405,7 +406,7 @@ async fn invalid_provider_response_runs_post_call_before_normalization_failure() match call.resume(result.take()).await { Ok(OcrCallStep::Host(OcrHostOperation::ProjectRequest)) => { result = Some(OcrHostResult::Request(Ok(( - Box::new(request.take().unwrap()), + Box::new(request.take().unwrap().into()), false, )))); } @@ -467,9 +468,10 @@ async fn direct_native_host_drives_the_same_state_machine() { _ => panic!("unexpected OCR operation"), }); result = Some(match operation { - OcrHostOperation::ProjectRequest => { - OcrHostResult::Request(Ok((Box::new(request.take().unwrap()), false))) - } + OcrHostOperation::ProjectRequest => OcrHostResult::Request(Ok(( + Box::new(request.take().unwrap().into()), + false, + ))), operation => host.invoke(operation).await, }); } @@ -501,6 +503,137 @@ async fn direct_native_host_drives_the_same_state_machine() { )); } +async fn drive_native_file_call( + request: super::LiteLLMOcrRequest, + content: Result, +) -> (Result, usize) { + let NativeOutcome::Completed(mut call) = + OcrCall::admit(super::test_support::ocr_client(), OcrAdmission::all()) + else { + panic!("supported call declined") + }; + let mut request = Some(request); + let mut content = Some(content); + let mut result = None; + let mut reads = 0; + let outcome = loop { + match call.resume(result.take()).await { + Ok(OcrCallStep::Host(OcrHostOperation::ProjectRequest)) => { + result = Some(OcrHostResult::Request(Ok(( + Box::new(request.take().unwrap()), + false, + )))); + } + Ok(OcrCallStep::Host(OcrHostOperation::ReadDocument)) => { + reads += 1; + result = Some(OcrHostResult::Document(content.take().unwrap())); + } + Ok(OcrCallStep::Host(operation)) => result = Some(NoopOcrHost.invoke(operation).await), + Ok(OcrCallStep::Complete(response)) => break Ok(response), + Err(error) => break Err(error), + } + }; + (outcome, reads) +} + +#[tokio::test] +async fn host_reader_documents_are_read_once_at_the_core_selected_point_and_encoded() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "pages":[{"index":0,"markdown":"file"}] + }))]) + .await; + let request = wire_request("mistral/model", &base, json!({})).with_document( + super::OcrDocumentInput::HostReader { + mime_type: Some("application/pdf".into()), + }, + ); + let (response, reads) = drive_native_file_call( + request, + Ok(super::OcrFileContent { + bytes: b"abc".as_slice().into(), + file_name: Some("scan.png".into()), + }), + ) + .await; + server.await.unwrap(); + assert_eq!(response.unwrap().pages[0]["markdown"], "file"); + assert_eq!(reads, 1); + assert!(seen.lock().unwrap()[0].contains("data:application/pdf;base64,YWJj")); +} + +#[tokio::test] +async fn host_reader_failures_and_empty_files_fail_before_the_provider_is_called() { + let (base, seen, _server) = mock_server(vec![]).await; + let request = wire_request("mistral/model", &base, json!({})); + let failure = crate::ocr::Error::InvalidRequest("reader exploded".into()); + let (response, reads) = drive_native_file_call( + request.with_document(super::OcrDocumentInput::HostReader { mime_type: None }), + Err(failure.clone()), + ) + .await; + assert_eq!(response.unwrap_err(), failure); + assert_eq!(reads, 1); + + let request = wire_request("mistral/model", &base, json!({})); + let (response, _) = drive_native_file_call( + request.with_document(super::OcrDocumentInput::HostReader { mime_type: None }), + Ok(super::OcrFileContent { + bytes: Default::default(), + file_name: None, + }), + ) + .await; + assert!(matches!( + response.unwrap_err(), + crate::ocr::Error::InvalidRequest(_) + )); + assert!(seen.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn path_documents_are_read_by_core_without_a_host_operation() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "pages":[{"index":0,"markdown":"path"}] + }))]) + .await; + let dir = std::env::temp_dir().join(format!("litellm-ocr-{}", rand::random::())); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.join("scan.png"); + std::fs::write(&path, b"abc").unwrap(); + let request = wire_request("mistral/model", &base, json!({})).with_document( + super::OcrDocumentInput::Path { + path: path.clone(), + mime_type: None, + }, + ); + let (response, reads) = drive_native_file_call( + request, + Err(crate::ocr::Error::InvalidRequest("unused".into())), + ) + .await; + server.await.unwrap(); + std::fs::remove_dir_all(&dir).unwrap(); + assert_eq!(response.unwrap().pages[0]["markdown"], "path"); + assert_eq!(reads, 0); + assert!(seen.lock().unwrap()[0].contains("data:image/png;base64,YWJj")); + + let (base, seen, _server) = mock_server(vec![]).await; + let request = wire_request("mistral/model", &base, json!({})); + let (response, _) = drive_native_file_call( + request.with_document(super::OcrDocumentInput::Path { + path: path.clone(), + mime_type: None, + }), + Err(crate::ocr::Error::InvalidRequest("unused".into())), + ) + .await; + assert!(matches!( + response.unwrap_err(), + crate::ocr::Error::FileRead { path: failed, kind: std::io::ErrorKind::NotFound, .. } if failed == path + )); + assert!(seen.lock().unwrap().is_empty()); +} + #[tokio::test] async fn public_finalization_failure_never_dispatches_success_or_replays_provider() { use crate::call_lifecycle::host::{HostFailure, HostPhase}; @@ -543,9 +676,10 @@ async fn public_finalization_failure_never_dispatches_success_or_replays_provide | OcrHostOperation::Lifecycle(HostPhase::DeploymentFailure) => { panic!("finalization failure used provider/success dispatch") } - OcrHostOperation::ProjectRequest => { - OcrHostResult::Request(Ok((Box::new(request.take().unwrap()), false))) - } + OcrHostOperation::ProjectRequest => OcrHostResult::Request(Ok(( + Box::new(request.take().unwrap().into()), + false, + ))), operation => host.invoke(operation).await, }); } @@ -582,7 +716,7 @@ async fn cancellation_at_provider_hook_prevents_execution_and_further_resumption OcrCallStep::Host(OcrHostOperation::PreCall(_)) => break, OcrCallStep::Host(OcrHostOperation::ProjectRequest) => { result = Some(OcrHostResult::Request(Ok(( - Box::new(request.take().unwrap()), + Box::new(request.take().unwrap().into()), false, )))) } @@ -799,7 +933,7 @@ async fn cancellation_waits_for_provider_capture_drop_even_when_acknowledgement_ _ = entered.notified() => break, step = call.resume(result.take()) => { result = Some(match step.unwrap() { - OcrCallStep::Host(OcrHostOperation::ProjectRequest) => OcrHostResult::Request(Ok((Box::new(request.take().unwrap()), false))), + OcrCallStep::Host(OcrHostOperation::ProjectRequest) => OcrHostResult::Request(Ok((Box::new(request.take().unwrap().into()), false))), OcrCallStep::Host(operation) => NoopOcrHost.invoke(operation).await, OcrCallStep::Complete(_) => panic!("pending provider completed"), }); diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 1562d4c1021..6dde7c71af6 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -16,6 +16,7 @@ extension-module = ["pyo3/extension-module"] panic-test = [] [dependencies] +bytes.workspace = true futures-util.workspace = true litellm-core.workspace = true litellm-auth.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/document.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/document.rs index d43c2f88775..33c0561184d 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/document.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/document.rs @@ -1,97 +1,56 @@ -use std::io::Read; use std::path::PathBuf; -use pyo3::exceptions::{PyFileNotFoundError, PyTypeError, PyValueError}; +use bytes::Bytes; +use pyo3::exceptions::{PyTypeError, PyValueError}; +use pyo3::gc::{PyTraverseError, PyVisit}; use pyo3::prelude::*; use pyo3::pybacked::PyBackedBytes; -#[cfg(test)] -use pyo3::types::PyDict; use pyo3::types::{PyBytes, PyString}; -use litellm_core::constants::OCR_INLINE_MAX_BYTES; -use litellm_core::ocr::{OcrDocument, encode_file_document, mime_type_for_name, upload_mime_type}; -use litellm_python_interop::to_py_preserving_errors; +use litellm_core::ocr::{OcrDocumentInput, OcrFileContent}; -enum FileBytes { - Python(PyBackedBytes), - Native(Vec), +#[derive(Debug)] +pub(super) struct PythonFileReader { + reader: Py, + name: Option, } -impl AsRef<[u8]> for FileBytes { - fn as_ref(&self) -> &[u8] { - match self { - Self::Python(bytes) => bytes, - Self::Native(bytes) => bytes, - } +impl PythonFileReader { + pub(super) fn read(&self, py: Python<'_>) -> PyResult { + let value = self.reader.bind(py).call0()?; + let bytes = if value.is_instance_of::() { + Bytes::from(value.extract::()?) + } else if value.is_instance_of::() { + extract_bytes(&value)? + } else { + return Err(PyTypeError::new_err(format!( + "OCR file read must return bytes or str, got {}", + value.get_type(), + ))); + }; + Ok(OcrFileContent { + bytes, + file_name: self.name.clone(), + }) + } + + pub(super) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.reader) } } -fn read_file_input( - py: Python<'_>, - file: &Bound<'_, PyAny>, -) -> PyResult<(FileBytes, Option)> { - if file.is_instance_of::() { - return Err(PyValueError::new_err( - "OCR file input does not accept bare str values. Pass bytes, a pathlib.Path, or a file-like object.", - )); +fn extract_bytes(value: &Bound<'_, PyAny>) -> PyResult { + if value.is_exact_instance_of::() { + return Ok(Bytes::from_owner(value.extract::()?)); } - if file.is_instance(&py.import("os")?.getattr("PathLike")?)? { - let path: PathBuf = file.extract()?; - let name = path - .file_name() - .map(|value| value.to_string_lossy().into_owned()); - let bytes = py - .detach(|| { - let mut bytes = Vec::new(); - std::fs::File::open(&path)? - .take(OCR_INLINE_MAX_BYTES as u64 + 1) - .read_to_end(&mut bytes)?; - Ok::<_, std::io::Error>(bytes) - }) - .map_err(|error| { - if error.kind() == std::io::ErrorKind::NotFound { - PyFileNotFoundError::new_err(format!("File not found: {}", path.display())) - } else { - error.into() - } - })?; - return Ok((FileBytes::Native(bytes), name)); - } - if file.is_instance_of::() { - return Ok((FileBytes::Python(file.extract()?), None)); - } - let reader = file - .getattr_opt("read")? - .filter(|value| value.is_callable()); - let Some(reader) = reader else { - return Err(PyValueError::new_err(format!( - "Unsupported file input type: {}. Expected pathlib.Path, bytes, or a file-like object.", - file.get_type(), - ))); - }; - let name = file - .getattr_opt("name")? - .filter(|value| !value.is_none()) - .map(|value| value.extract::()) - .transpose()?; - let value = reader.call0()?; - let bytes = if value.is_instance_of::() { - FileBytes::Native(value.extract::()?.into_bytes()) - } else if value.is_instance_of::() { - FileBytes::Python(value.extract()?) - } else { - return Err(PyTypeError::new_err(format!( - "OCR file read must return bytes or str, got {}", - value.get_type(), - ))); - }; - Ok((bytes, name)) + Ok(Bytes::copy_from_slice( + value.extract::()?.as_ref(), + )) } pub(super) struct FileDocumentInput { - bytes: FileBytes, - name: Option, - mime_type: Option, + pub input: OcrDocumentInput, + pub reader: Option, } impl FromPyObject<'_, '_> for FileDocumentInput { @@ -104,79 +63,79 @@ impl FromPyObject<'_, '_> for FileDocumentInput { Err(error) if error.is_instance_of::(py) => None, Err(error) => return Err(error), }; + let missing = || { + PyValueError::new_err( + "document with type='file' must include a 'file' field containing a pathlib.Path, file-like object, or bytes", + ) + }; let file = document.get_item("file").map_err(|error| { if error.is_instance_of::(py) { - PyValueError::new_err("document with type='file' must include a 'file' field containing a pathlib.Path, file-like object, or bytes") + missing() } else { error } })?; if file.is_none() { + return Err(missing()); + } + if file.is_instance_of::() { return Err(PyValueError::new_err( - "document with type='file' must include a 'file' field containing a pathlib.Path, file-like object, or bytes", + "OCR file input does not accept bare str values. Pass bytes, a pathlib.Path, or a file-like object.", )); } - let (bytes, name) = read_file_input(py, &file)?; + if file.is_instance(&py.import("os")?.getattr("PathLike")?)? { + return Ok(Self { + input: OcrDocumentInput::Path { + path: file.extract::()?, + mime_type, + }, + reader: None, + }); + } + if file.is_instance_of::() { + return Ok(Self { + input: OcrDocumentInput::Bytes { + bytes: extract_bytes(&file)?, + file_name: None, + mime_type, + }, + reader: None, + }); + } + let reader = file + .getattr_opt("read")? + .filter(|value| value.is_callable()); + let Some(reader) = reader else { + return Err(PyValueError::new_err(format!( + "Unsupported file input type: {}. Expected pathlib.Path, bytes, or a file-like object.", + file.get_type(), + ))); + }; + let name = file + .getattr_opt("name")? + .filter(|value| !value.is_none()) + .map(|value| value.extract::()) + .transpose()?; Ok(Self { - bytes, - name, - mime_type, + input: OcrDocumentInput::HostReader { mime_type }, + reader: Some(PythonFileReader { + reader: reader.unbind(), + name, + }), }) } } -pub(super) fn file_document(py: Python<'_>, document: FileDocumentInput) -> PyResult { - py.detach(|| { - encode_file_document( - document.bytes.as_ref(), - document.name.as_deref(), - document.mime_type.as_deref(), - ) - }) - .map_err(|error| PyValueError::new_err(error.to_string())) -} - -#[pyfunction] -fn _ocr_file_document(py: Python<'_>, document: Bound<'_, PyAny>) -> PyResult> { - to_py_preserving_errors(py, &file_document(py, document.extract()?)?) -} - -#[pyfunction] -fn _ocr_mime_type(file_name: &str) -> String { - mime_type_for_name(file_name).into() -} - -#[pyfunction] -#[pyo3(signature = (file_content, file_name=None, content_type=None))] -fn _ocr_upload_document( - py: Python<'_>, - file_content: &Bound<'_, PyBytes>, - file_name: Option<&str>, - content_type: Option<&str>, -) -> PyResult> { - let bytes: PyBackedBytes = file_content.extract()?; - let document = py - .detach(|| { - encode_file_document( - &bytes, - None, - Some(upload_mime_type(file_name, content_type)), - ) - }) - .map_err(|error| PyValueError::new_err(error.to_string()))?; - to_py_preserving_errors(py, &document) -} - -pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { - module.add("_OCR_MAX_FILE_BYTES", OCR_INLINE_MAX_BYTES)?; - module.add_function(wrap_pyfunction!(_ocr_upload_document, module)?)?; - module.add_function(wrap_pyfunction!(_ocr_file_document, module)?)?; - module.add_function(wrap_pyfunction!(_ocr_mime_type, module)?) -} - #[cfg(test)] mod tests { use super::*; + use pyo3::types::PyDict; + + fn eval<'py>(py: Python<'py>, source: &std::ffi::CStr) -> Bound<'py, PyDict> { + let locals = PyDict::new(py); + py.run(source, Some(&locals), Some(&locals)).unwrap(); + locals + } #[test] fn extraction_validates_required_file_and_optional_mime_type() { @@ -196,69 +155,148 @@ mod tests { let error = document.extract::().err().unwrap(); assert!(error.is_instance_of::(py)); } - let document = py.eval(c"{'file': b'abc'}", None, None).unwrap(); + let error = py + .eval(c"{'file': 'scan.pdf'}", None, None) + .unwrap() + .extract::() + .err() + .unwrap(); + assert!(error.is_instance_of::(py)); + assert!(error.to_string().contains("bare str")); + let document = py + .eval(c"{'file': b'abc', 'mime_type': 'image/png'}", None, None) + .unwrap(); let input: FileDocumentInput = document.extract().unwrap(); - assert_eq!(input.bytes.as_ref(), b"abc"); - assert_eq!(input.name, None); - assert_eq!(input.mime_type, None); + assert!(input.reader.is_none()); + assert_eq!( + input.input, + OcrDocumentInput::Bytes { + bytes: b"abc".as_slice().into(), + file_name: None, + mime_type: Some("image/png".into()), + } + ); }); } #[test] - fn extraction_validates_mime_type_before_consuming_file() { + fn paths_and_readers_are_projected_without_io() { Python::initialize(); Python::attach(|py| { - let locals = PyDict::new(py); - py.run( - c"class Reader: + let locals = eval( + py, + c"from pathlib import Path +class Reader: + name = 'scan.png' def __init__(self): self.reads = 0 def read(self): self.reads += 1 return b'abc' reader = Reader() -document = {'file': reader, 'mime_type': 7}", - Some(&locals), - Some(&locals), - ) - .unwrap(); +document = {'file': reader, 'mime_type': 7} +reader_document = {'file': reader} +path_document = {'file': Path('/nonexistent/ocr-projection-test.pdf'), 'mime_type': 'image/png'}", + ); let document = locals.get_item("document").unwrap().unwrap(); let error = document.extract::().err().unwrap(); assert!(error.is_instance_of::(py)); - let reads: usize = locals - .get_item("reader") - .unwrap() - .unwrap() - .getattr("reads") - .unwrap() - .extract() - .unwrap(); - assert_eq!(reads, 0); + + let document = locals.get_item("reader_document").unwrap().unwrap(); + let input: FileDocumentInput = document.extract().unwrap(); + assert_eq!( + input.input, + OcrDocumentInput::HostReader { mime_type: None } + ); + let reads = || { + locals + .get_item("reader") + .unwrap() + .unwrap() + .getattr("reads") + .unwrap() + .extract::() + .unwrap() + }; + assert_eq!(reads(), 0); + let content = input.reader.unwrap().read(py).unwrap(); + assert_eq!(reads(), 1); + assert_eq!( + content, + OcrFileContent { + bytes: b"abc".as_slice().into(), + file_name: Some("scan.png".into()), + } + ); + + let document = locals.get_item("path_document").unwrap().unwrap(); + let input: FileDocumentInput = document.extract().unwrap(); + assert!(input.reader.is_none()); + assert_eq!( + input.input, + OcrDocumentInput::Path { + path: PathBuf::from("/nonexistent/ocr-projection-test.pdf"), + mime_type: Some("image/png".into()), + } + ); }); } #[test] - fn extraction_preserves_reader_key_error_identity() { + fn reader_results_are_normalized_and_exceptions_keep_their_identity() { Python::initialize(); Python::attach(|py| { - let locals = PyDict::new(py); - py.run( + let locals = eval( + py, c"failure = KeyError('reader failed') -class Reader: +class Raising: def read(self): raise failure -document = {'file': Reader()}", - Some(&locals), - Some(&locals), - ) - .unwrap(); - let document = locals.get_item("document").unwrap().unwrap(); - let error = document.extract::().err().unwrap(); +class Text: + def read(self): + return 'héllo' +class Wrong: + def read(self): + return 7 +raising = {'file': Raising()} +text = {'file': Text()} +wrong = {'file': Wrong()}", + ); + let reader = |name: &str| { + locals + .get_item(name) + .unwrap() + .unwrap() + .extract::() + .unwrap() + .reader + .unwrap() + }; + let error = reader("raising").read(py).unwrap_err(); assert!( error .value(py) .is(locals.get_item("failure").unwrap().unwrap()) ); + assert_eq!( + reader("text").read(py).unwrap().bytes.as_ref(), + "héllo".as_bytes() + ); + let error = reader("wrong").read(py).unwrap_err(); + assert!(error.is_instance_of::(py)); + assert!(error.to_string().contains("bytes or str")); }); } + + #[test] + fn exact_python_bytes_transfer_without_copying_and_outlive_the_input() { + Python::initialize(); + let (bytes, pointer) = Python::attach(|py| { + let value = PyBytes::new(py, b"document bytes"); + let pointer = value.as_bytes().as_ptr() as usize; + (extract_bytes(value.as_any()).unwrap(), pointer) + }); + assert_eq!(bytes.as_ptr() as usize, pointer); + assert_eq!(bytes.as_ref(), b"document bytes"); + } } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs index e4ce813d297..7dbc35289ff 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs @@ -1,4 +1,5 @@ use litellm_core::ocr::Error; +use pyo3::exceptions::{PyFileNotFoundError, PyOSError}; use pyo3::prelude::*; use crate::errors::{RustUpstreamError, core_error_to_pyerr}; @@ -7,6 +8,12 @@ pub(super) fn to_pyerr(error: Error) -> PyErr { let status = error.http_status_code(); let mapped = match error { Error::Http { status, body } => RustUpstreamError::new_err((status, body)), + Error::FileRead { + path, + kind: std::io::ErrorKind::NotFound, + .. + } => PyFileNotFoundError::new_err(format!("File not found: {}", path.display())), + Error::FileRead { message, .. } => PyOSError::new_err(message), other => core_error_to_pyerr(other.into()), }; attach_status(mapped, status) diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs index 32794936899..e710b0d82f9 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs @@ -66,13 +66,27 @@ impl PythonOcrHost { retained_fields.set_item(name, value)?; } } - retained_fields.set_item("document", &self.projected()?.fields.document)?; let projected = self.projected_mut()?; + let document = match &projected.fields.document { + Some(document) => document.clone_ref(py), + None => to_py(py, &request.document)?, + }; + retained_fields.set_item("document", &document)?; + projected.fields.document = Some(document); projected.retained_fields = Some(retained_fields.unbind()); projected.pre_call = Some((&request).into()); Ok(request) } + fn read_document(&self, py: Python<'_>) -> PyResult { + self.projected()? + .fields + .reader + .as_ref() + .ok_or_else(missing_state)? + .read(py) + } + fn acquire_azure_ad_token(&self, py: Python<'_>) -> PyResult { let provider = self .projected()? @@ -193,7 +207,7 @@ impl PythonRoute for PythonOcrHost { let OcrHostData::Unprojected { request } = &self.data else { return Err(missing_state()); }; - let projected = project_request(py, request.bind(py), self.state.kwargs.bind(py))?; + let projected = project_request(request.bind(py), self.state.kwargs.bind(py))?; let has_token_provider = projected.fields.azure_ad_token_provider.is_some(); let request = projected.request; self.data = OcrHostData::Projected(Box::new(ProjectedOcrHost { @@ -205,6 +219,7 @@ impl PythonRoute for PythonOcrHost { })); OcrHostResult::Request(Ok((Box::new(request), has_token_provider))) } + OcrHostOperation::ReadDocument => OcrHostResult::Document(Ok(self.read_document(py)?)), OcrHostOperation::AcquireAzureAdToken => { OcrHostResult::AzureAdToken(Ok(self.acquire_azure_ad_token(py)?)) } @@ -258,6 +273,9 @@ impl PythonRoute for PythonOcrHost { OcrHostData::Projected(projected) => { visit.call(&projected.fields.boundary_request)?; visit.call(&projected.fields.document)?; + if let Some(reader) = &projected.fields.reader { + reader.traverse(visit)?; + } visit.call(&projected.fields.api_key)?; if let Some(provider) = &projected.fields.azure_ad_token_provider { provider.traverse(visit)?; diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs index f17bf249b7f..5eae8ccf33f 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs @@ -9,6 +9,5 @@ use pyo3::prelude::*; pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { value::register(module)?; - document::register(module)?; lifecycle::register(module) } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs index 8d8d5f8c518..ad223645c62 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs @@ -1,14 +1,15 @@ use std::sync::Arc; -use litellm_core::ocr::wire::{OcrWireRequest, consumed_optional_params, decode_request}; -use litellm_core::ocr::{LiteLLMOcrRequest, NativeOutcome, OcrCall}; -use litellm_python_interop::{ - from_py_preserving_errors as from_py, to_py_preserving_errors as to_py, +use litellm_core::ocr::wire::{ + OcrWireRequest, consumed_optional_params, decode_document, decode_request_input, }; +use litellm_core::ocr::{LiteLLMOcrRequest, NativeOutcome, OcrCall, OcrDocumentInput}; +use litellm_python_interop::from_py_preserving_errors as from_py; use pyo3::prelude::*; use pyo3::types::PyDict; use serde_json::{Map, Value}; +use super::document::{FileDocumentInput, PythonFileReader}; use super::errors::to_pyerr as ocr_error_to_pyerr; use super::lifecycle::BridgeOcrHooks; use crate::auth::{AZURE_AD_TOKEN_PROVIDER, PythonTokenProvider}; @@ -17,7 +18,8 @@ use crate::marshal::{project_optional_fields, python_timeout_seconds, request_in pub(super) struct ProjectedOcrFields { pub boundary_request: Py, - pub document: Py, + pub document: Option>, + pub reader: Option, pub api_key: Py, pub azure_ad_token_provider: Option, pub provider: &'static str, @@ -25,7 +27,7 @@ pub(super) struct ProjectedOcrFields { } pub(super) struct ProjectedOcrCall { - pub request: LiteLLMOcrRequest, + pub request: LiteLLMOcrRequest, pub fields: ProjectedOcrFields, } @@ -80,12 +82,12 @@ impl<'py> OcrArguments<'_, 'py> { } enum ProjectedDocument { - File { wire: Value, retained: Py }, + File(FileDocumentInput), Other { wire: Value, retained: Py }, } impl ProjectedDocument { - fn project(py: Python<'_>, document: &Bound<'_, PyAny>) -> PyResult { + fn project(document: &Bound<'_, PyAny>) -> PyResult { let kind: String = document.get_item("type")?.extract()?; if kind != "file" { return Ok(Self::Other { @@ -93,25 +95,28 @@ impl ProjectedDocument { retained: document.clone().unbind(), }); } - let input = document.extract()?; - let encoded = super::document::file_document(py, input)?; - let wire = serde_json::to_value(encoded) - .map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string()))?; - Ok(Self::File { - retained: to_py(py, &wire)?, - wire, - }) + Ok(Self::File(document.extract()?)) } - fn into_parts(self) -> (Value, Py) { + fn into_parts( + self, + ) -> PyResult<( + OcrDocumentInput, + Option>, + Option, + )> { match self { - Self::File { wire, retained } | Self::Other { wire, retained } => (wire, retained), + Self::File(FileDocumentInput { input, reader }) => Ok((input, None, reader)), + Self::Other { wire, retained } => Ok(( + decode_document(wire).map_err(ocr_error_to_pyerr)?.into(), + Some(retained), + None, + )), } } } pub(super) fn project_request( - py: Python<'_>, request: &Bound<'_, PyAny>, kwargs: &Bound<'_, PyDict>, ) -> PyResult { @@ -119,8 +124,7 @@ pub(super) fn project_request( let arguments = OcrArguments { request, kwargs }; let model = arguments.model()?; let custom_llm_provider = arguments.custom_llm_provider()?; - let (wire_document, retained_document) = - ProjectedDocument::project(py, &arguments.document()?)?.into_parts(); + let document = ProjectedDocument::project(&arguments.document()?)?; let api_key = arguments.api_key()?; let specs = consumed_optional_params(&model, custom_llm_provider.as_deref()) .map_err(ocr_error_to_pyerr)?; @@ -136,9 +140,10 @@ pub(super) fn project_request( let azure_ad_token_provider = kwargs .get_item("azure_ad_token_provider")? .and_then(|provider| PythonTokenProvider::select(provider, AZURE_AD_TOKEN_PROVIDER)); + let (document, retained_document, reader) = document.into_parts()?; let wire = OcrWireRequest { model, - document: wire_document, + document, api_key: api_key.extract()?, api_base: arguments.api_base()?, custom_llm_provider, @@ -147,13 +152,14 @@ pub(super) fn project_request( input_sources, timeout_seconds: arguments.timeout_seconds()?, }; - let request = decode_request(wire).map_err(ocr_error_to_pyerr)?; + let request = decode_request_input(wire).map_err(ocr_error_to_pyerr)?; let provider = request.provider_name(); Ok(ProjectedOcrCall { request: request.with_host_hooks(Arc::new(BridgeOcrHooks), None), fields: ProjectedOcrFields { boundary_request, document: retained_document, + reader, api_key: api_key.unbind(), azure_ad_token_provider, provider, @@ -197,10 +203,21 @@ mod tests { } fn project_document( - py: Python<'_>, document: &Bound<'_, PyAny>, - ) -> PyResult<(Value, Py)> { - ProjectedDocument::project(py, document).map(ProjectedDocument::into_parts) + ) -> PyResult<( + OcrDocumentInput, + Option>, + Option, + )> { + ProjectedDocument::project(document)?.into_parts() + } + + fn url_document(url: &str) -> OcrDocumentInput { + litellm_core::ocr::OcrDocument::DocumentUrl { + document_url: url.into(), + extra_fields: Map::new(), + } + .into() } fn stub_timeout_conversion(py: Python<'_>) { @@ -374,7 +391,7 @@ kwargs = {} } #[test] - fn document_reader_mutations_are_visible_to_later_field_reads() { + fn document_readers_are_not_consumed_during_projection() { Python::initialize(); Python::attach(|py| { stub_timeout_conversion(py); @@ -406,7 +423,12 @@ kwargs = {} .unwrap(); let arguments = arguments(&request, &kwargs); let document = arguments.document().unwrap(); - project_document(py, &document).unwrap(); + let (input, retained, reader) = project_document(&document).unwrap(); + assert_eq!(input, OcrDocumentInput::HostReader { mime_type: None }); + assert!(retained.is_none()); + assert_eq!(arguments.api_base().unwrap().as_deref(), Some("original")); + assert_eq!(arguments.timeout_seconds().unwrap(), Some(1.0)); + reader.unwrap().read(py).unwrap(); assert_eq!(arguments.api_base().unwrap().as_deref(), Some("mutated")); assert_eq!(arguments.timeout_seconds().unwrap(), Some(9.0)); }); @@ -444,7 +466,7 @@ kwargs = {'api_key': key} } #[test] - fn file_documents_are_encoded_and_other_documents_keep_the_python_object() { + fn file_documents_become_typed_inputs_and_other_documents_keep_the_python_object() { Python::initialize(); Python::attach(|py| { let file = py @@ -454,13 +476,17 @@ kwargs = {'api_key': key} None, ) .unwrap(); + let (input, retained, reader) = project_document(&file).unwrap(); assert_eq!( - project_document(py, &file).unwrap().0, - serde_json::json!({ - "type": "document_url", - "document_url": "data:application/pdf;base64,JVBERi0xLjQ=", - }) + input, + OcrDocumentInput::Bytes { + bytes: b"%PDF-1.4".as_slice().into(), + file_name: None, + mime_type: Some("application/pdf".into()), + } ); + assert!(retained.is_none()); + assert!(reader.is_none()); let original = py .eval( @@ -469,44 +495,21 @@ kwargs = {'api_key': key} None, ) .unwrap(); - let (wire, retained) = project_document(py, &original).unwrap(); - assert_eq!( - wire, - serde_json::json!({ - "type": "document_url", - "document_url": "https://example.com/a.pdf", - }) - ); - assert!(retained.bind(py).is(&original)); + let (input, retained, _) = project_document(&original).unwrap(); + assert_eq!(input, url_document("https://example.com/a.pdf")); + assert!(retained.unwrap().bind(py).is(&original)); }); } #[test] - fn unknown_document_types_reach_existing_downstream_validation() { + fn unknown_document_types_reach_existing_core_validation() { Python::initialize(); Python::attach(|py| { let document = py .eval(c"{'type': 'mystery', 'mystery': 'x'}", None, None) .unwrap(); - let wire_document = project_document(py, &document).unwrap().0; - assert_eq!( - wire_document, - serde_json::json!({"type": "mystery", "mystery": "x"}) - ); - let error = match decode_request(OcrWireRequest { - model: "mistral/mistral-ocr-latest".into(), - document: wire_document, - api_key: None, - api_base: None, - custom_llm_provider: None, - extra_headers: None, - optional_params: Map::new(), - input_sources: Default::default(), - timeout_seconds: None, - }) { - Ok(_) => panic!("unknown discriminators belong to core validation"), - Err(error) => error, - }; + let error = project_document(&document).unwrap_err(); + assert!(error.is_instance_of::(py)); assert!(error.to_string().contains("document")); }); } @@ -517,14 +520,14 @@ kwargs = {'api_key': key} Python::attach(|py| { let missing = py.eval(c"{}", None, None).unwrap(); assert!( - project_document(py, &missing) + project_document(&missing) .unwrap_err() .is_instance_of::(py) ); let non_string = py.eval(c"{'type': 1}", None, None).unwrap(); assert!( - project_document(py, &non_string) + project_document(&non_string) .unwrap_err() .is_instance_of::(py) ); @@ -540,7 +543,7 @@ document = Document() ", ); let error = - project_document(py, &locals.get_item("document").unwrap().unwrap()).unwrap_err(); + project_document(&locals.get_item("document").unwrap().unwrap()).unwrap_err(); assert!( error .value(py) @@ -569,9 +572,9 @@ document = Document() ", ); let document = locals.get_item("document").unwrap().unwrap(); - let (wire, retained) = project_document(py, &document).unwrap(); - assert_eq!(wire["type"], "document_url"); - assert!(!retained.bind(py).is(&document)); + let (input, retained, _) = project_document(&document).unwrap(); + assert!(matches!(input, OcrDocumentInput::Bytes { .. })); + assert!(retained.is_none()); let reads: Vec = document.getattr("reads").unwrap().extract().unwrap(); assert_eq!(reads, ["type", "mime_type", "file"]); }); diff --git a/litellm/constants.py b/litellm/constants.py index ce5b65080ee..338fe0f6b85 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -40,6 +40,7 @@ ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG: Final[frozenset[str]] = frozenset( "router_general_settings", "ignore_invalid_deployments", "fallback_access_check", + "fallback_budget_check", "auto_router_capability_limit", } ) @@ -53,6 +54,7 @@ S3_BOUNDED_OBJECT_KEY_HEAD_BYTES: Final = 64 S3_PREFIX_DIGEST_CHARS: Final = 16 # s3 allows 2048 bytes of combined metadata headers, which Content-Disposition counts against MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES: Final = 1024 +S3_LOG_PROMPTS_ONLY_ENV_VAR: Final = "S3_LOG_PROMPTS_ONLY" MAX_FILE_LIST_LIMIT: Final = 10000 DEFAULT_SQS_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10)) DEFAULT_NUM_WORKERS_LITELLM_PROXY: Final = int(os.getenv("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1)) @@ -1499,6 +1501,7 @@ OUTPUT_TOKEN_CEILING_PARAMS: Final = frozenset({"max_tokens", "max_completion_to CLIENT_OUTPUT_CEILING_METADATA_KEY: Final = "_client_output_ceiling" CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags" ROUTING_REQUEST_TAGS_METADATA_KEY: Final = "_routing_request_tags" +ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: Final = "_litellm_router_usage_counted_tokens" INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin" SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated" SESSION_ID_OMITTED_METADATA_KEY: Final = "litellm_session_id_omitted" diff --git a/litellm/integrations/callback_configs.json b/litellm/integrations/callback_configs.json index 85bfcc6e7ed..6806188c97c 100644 --- a/litellm/integrations/callback_configs.json +++ b/litellm/integrations/callback_configs.json @@ -446,6 +446,12 @@ "ui_name": "S3 Path Prefix", "description": "Path prefix within the bucket for organizing logs", "required": false + }, + "s3_log_prompts_only": { + "type": "boolean", + "ui_name": "Log Prompts Only", + "description": "Log request messages to S3 but drop the model response from each logged object", + "required": false } }, "description": "S3 Bucket (AWS) Logging Integration" diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 70d2f3ae5c3..5b5261fab6b 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -118,6 +118,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac alias_map: Final = { "langfuse_otel": "langfuse", + "s3_v2": "s3", } lookup_name: Final = alias_map.get(normalized_name, normalized_name) diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index 9456817a205..34b712d47e3 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -21,6 +21,7 @@ from litellm.integrations.opentelemetry_utils.gen_ai_semconv import ( OTELSemconvCategory, parse_semconv_opt_in, ) +from litellm.integrations.otel.mappers.utils import drop_none from litellm.integrations.otel.model.baggage import promoted_metadata from litellm.integrations.otel.model.db_endpoint import db_span_attributes from litellm.integrations.otel.model.metadata import flatten_metadata @@ -208,6 +209,20 @@ def _resolve_metric_attribute_filter( ) +def _provider_label(custom_llm_provider: object) -> str | None: + """The provider label for one call's metrics and events, or None when the + call carries no provider. + + Every attribute set drops None before export, so the label is simply absent + in that case: the OTLP encoder rejects a None attribute value outright, and a + placeholder would mint a permanent metric series that no operator can act + on. Mirrors the v2 integration's ``_provider_attributes``. + """ + if not isinstance(custom_llm_provider, str) or not custom_llm_provider: + return None + return custom_llm_provider + + def _normalize_team_metadata_keys(value: str | Iterable[object] | None) -> list[str]: """Coerce a team-metadata allowlist from a list or comma-separated string. @@ -1616,19 +1631,22 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): ) = _resolve_metric_attribute_filter(attributes) self._metric_attr_filter_resolved = True - def _filter_metric_attributes(self, attrs: dict[str, str]) -> dict[str, str]: + def _filter_metric_attributes(self, attrs: Mapping[str, str | None]) -> dict[str, str]: if not self._metric_attr_filter_resolved: self._ensure_metric_attribute_filter() + return {k: v for k, v in attrs.items() if v is not None and self._metric_attribute_allowed(k)} + + def _metric_attribute_allowed(self, key: str) -> bool: if self._metric_attr_include is not None: - return {k: v for k, v in attrs.items() if k in self._metric_attr_include} + return key in self._metric_attr_include if self._metric_attr_exclude is not None: - return {k: v for k, v in attrs.items() if k not in self._metric_attr_exclude} - return attrs + return key not in self._metric_attr_exclude + return True def _record_metrics(self, kwargs, response_obj, start_time, end_time): duration_s: Final = (end_time - start_time).total_seconds() params: Final = kwargs.get("litellm_params") or {} - provider: Final = params.get("custom_llm_provider", "Unknown") + provider: Final = _provider_label(params.get("custom_llm_provider")) common_attrs = { "gen_ai.operation.name": ( @@ -1872,7 +1890,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): otel_logger: Final = self._logger_provider.get_logger(LITELLM_LOGGER_NAME) parent_ctx: Final = span.get_span_context() - provider: Final = (kwargs.get("litellm_params") or {}).get("custom_llm_provider", "Unknown") + provider: Final = _provider_label((kwargs.get("litellm_params") or {}).get("custom_llm_provider")) if self._gen_ai_semconv_latest_experimental: self._emit_inference_details_event( @@ -1909,7 +1927,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): severity_number=SeverityNumber.INFO, severity_text="INFO", body=body, - attributes=attrs, + attributes=drop_none(attrs), ) otel_logger.emit(log_record) @@ -1941,7 +1959,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): severity_number=SeverityNumber.INFO, severity_text="INFO", body=body, - attributes=attrs, + attributes=drop_none(attrs), ) otel_logger.emit(log_record) diff --git a/litellm/integrations/opentelemetry_utils/gen_ai_semconv.py b/litellm/integrations/opentelemetry_utils/gen_ai_semconv.py index b5eedc42fe9..81d9a947da7 100644 --- a/litellm/integrations/opentelemetry_utils/gen_ai_semconv.py +++ b/litellm/integrations/opentelemetry_utils/gen_ai_semconv.py @@ -33,6 +33,7 @@ from datetime import datetime from enum import Enum from typing import TYPE_CHECKING, Any, Final +from litellm.integrations.otel.mappers.utils import drop_none from litellm.litellm_core_utils.safe_json_dumps import safe_dumps if TYPE_CHECKING: @@ -195,13 +196,16 @@ class OTELGenAISemconvMixin: if value: self.safe_set_attribute(span=span, key=semconv_key, value=value) - def _build_inference_details_attrs(self, kwargs: dict, response_obj: dict, provider: str) -> dict[str, str]: + def _build_inference_details_attrs( + self, kwargs: dict, response_obj: dict, provider: str | None + ) -> dict[str, str | None]: """Build the attribute payload for the inference-details event. - Always includes provider/operation; input/output messages are added + Always includes operation and provider (None when the call carries none, + dropped before the event is emitted); input/output messages are added only when content capture is enabled and non-empty. Mixin-internal. """ - attrs: Final[dict[str, str]] = { + attrs: Final[dict[str, str | None]] = { "event_name": _INFERENCE_DETAILS_EVENT_NAME, "gen_ai.provider.name": provider, "gen_ai.operation.name": self._gen_ai_operation_name(kwargs), @@ -221,7 +225,7 @@ class OTELGenAISemconvMixin: self, kwargs: dict, response_obj: dict, - provider: str, + provider: str | None, otel_logger, parent_ctx, ) -> None: @@ -239,6 +243,6 @@ class OTELGenAISemconvMixin: severity_number=SeverityNumber.INFO, severity_text="INFO", body=None, - attributes=self._build_inference_details_attrs(kwargs, response_obj, provider), + attributes=drop_none(self._build_inference_details_attrs(kwargs, response_obj, provider)), ) otel_logger.emit(log_record) diff --git a/litellm/integrations/s3.py b/litellm/integrations/s3.py index 8ce461eea5b..796784fb993 100644 --- a/litellm/integrations/s3.py +++ b/litellm/integrations/s3.py @@ -2,19 +2,42 @@ # On success + failure, log events to Supabase import hashlib +import os +from collections.abc import Mapping from datetime import datetime from typing import Final, cast +from pydantic import TypeAdapter, ValidationError + import litellm from litellm._logging import print_verbose, verbose_logger from litellm.constants import ( MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES, MAX_S3_OBJECT_KEY_BYTES, S3_BOUNDED_OBJECT_KEY_HEAD_BYTES, + S3_LOG_PROMPTS_ONLY_ENV_VAR, S3_PREFIX_DIGEST_CHARS, ) from litellm.types.utils import StandardLoggingPayload +_S3_LOG_PROMPTS_ONLY: Final = TypeAdapter(bool) + + +def resolve_s3_log_prompts_only(configured: object, environ: Mapping[str, str] | None = None) -> bool: + env: Final = os.environ if environ is None else environ + raw: Final = env.get(S3_LOG_PROMPTS_ONLY_ENV_VAR) if configured is None else configured + if raw is None or raw == "": + return False + try: + return _S3_LOG_PROMPTS_ONLY.validate_python(raw.strip() if isinstance(raw, str) else raw) + except ValidationError: + verbose_logger.warning("s3 logging: s3_log_prompts_only=%r is not a boolean, logging prompts only", raw) + return True + + +def prompts_only_payload(payload: StandardLoggingPayload) -> StandardLoggingPayload: + return {**payload, "response": None} + class S3Logger: # Class variables or attributes @@ -33,6 +56,7 @@ class S3Logger: s3_config=None, s3_server_side_encryption: str | None = None, s3_sse_kms_key_id: str | None = None, + s3_log_prompts_only: bool | None = None, **kwargs, ): import boto3 @@ -41,29 +65,30 @@ class S3Logger: verbose_logger.debug("in init s3 logger - s3_callback_params %s", litellm.s3_callback_params) s3_use_team_prefix = False + params: Final = { + key: litellm.get_secret(value) if isinstance(value, str) and value.startswith("os.environ/") else value + for key, value in (litellm.s3_callback_params or {}).items() + } if litellm.s3_callback_params is not None: - # read in .env variables - example os.environ/AWS_BUCKET_NAME - for key, value in litellm.s3_callback_params.items(): - if isinstance(value, str) and value.startswith("os.environ/"): - litellm.s3_callback_params[key] = litellm.get_secret(value) - # now set s3 params from litellm.s3_logger_params - s3_bucket_name = litellm.s3_callback_params.get("s3_bucket_name") - s3_region_name = litellm.s3_callback_params.get("s3_region_name") - s3_api_version = litellm.s3_callback_params.get("s3_api_version") - s3_use_ssl = litellm.s3_callback_params.get("s3_use_ssl", True) - s3_verify = litellm.s3_callback_params.get("s3_verify") - s3_endpoint_url = litellm.s3_callback_params.get("s3_endpoint_url") - s3_aws_access_key_id = litellm.s3_callback_params.get("s3_aws_access_key_id") - s3_aws_secret_access_key = litellm.s3_callback_params.get("s3_aws_secret_access_key") - s3_aws_session_token = litellm.s3_callback_params.get("s3_aws_session_token") - s3_config = litellm.s3_callback_params.get("s3_config") - s3_path = litellm.s3_callback_params.get("s3_path") - s3_server_side_encryption = litellm.s3_callback_params.get("s3_server_side_encryption") - s3_sse_kms_key_id = litellm.s3_callback_params.get("s3_sse_kms_key_id") - # done reading litellm.s3_callback_params - s3_use_team_prefix = bool(litellm.s3_callback_params.get("s3_use_team_prefix", False)) + s3_bucket_name = params.get("s3_bucket_name") + s3_region_name = params.get("s3_region_name") + s3_api_version = params.get("s3_api_version") + s3_use_ssl = params.get("s3_use_ssl", True) + s3_verify = params.get("s3_verify") + s3_endpoint_url = params.get("s3_endpoint_url") + s3_aws_access_key_id = params.get("s3_aws_access_key_id") + s3_aws_secret_access_key = params.get("s3_aws_secret_access_key") + s3_aws_session_token = params.get("s3_aws_session_token") + s3_config = params.get("s3_config") + s3_path = params.get("s3_path") + s3_server_side_encryption = params.get("s3_server_side_encryption") + s3_sse_kms_key_id = params.get("s3_sse_kms_key_id") + s3_use_team_prefix = bool(params.get("s3_use_team_prefix", False)) self.s3_use_team_prefix = s3_use_team_prefix + self.s3_log_prompts_only: object = ( + params.get("s3_log_prompts_only") if s3_log_prompts_only is None else s3_log_prompts_only + ) self.bucket_name = s3_bucket_name self.s3_path = s3_path self.s3_server_side_encryption, self.s3_sse_kms_key_id = resolve_sse_params( @@ -144,7 +169,9 @@ class S3Logger: from litellm.litellm_core_utils.safe_json_dumps import safe_dumps - payload_str: Final = safe_dumps(payload) + payload_str: Final = safe_dumps( + prompts_only_payload(payload) if resolve_s3_log_prompts_only(self.s3_log_prompts_only) else payload + ) print_verbose(f"\ns3 Logger - Logging payload = {payload_str}") diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index 972ac79e306..826f55cc798 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -21,6 +21,8 @@ from litellm.constants import DEFAULT_S3_BATCH_SIZE, DEFAULT_S3_FLUSH_INTERVAL_S from litellm.integrations.s3 import ( get_s3_object_download_filename, get_s3_object_key, + prompts_only_payload, + resolve_s3_log_prompts_only, resolve_sse_params, ) from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix @@ -68,6 +70,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): s3_use_virtual_hosted_style: bool = False, s3_server_side_encryption: str | None = None, s3_sse_kms_key_id: str | None = None, + s3_log_prompts_only: bool | None = None, s3_callback_params_override: dict | None = None, **kwargs, ): @@ -108,6 +111,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): s3_use_virtual_hosted_style=s3_use_virtual_hosted_style, s3_server_side_encryption=s3_server_side_encryption, s3_sse_kms_key_id=s3_sse_kms_key_id, + s3_log_prompts_only=s3_log_prompts_only, ) verbose_logger.debug("s3 logger using endpoint url %s", s3_endpoint_url) @@ -163,6 +167,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): s3_use_virtual_hosted_style: bool = False, s3_server_side_encryption: str | None = None, s3_sse_kms_key_id: str | None = None, + s3_log_prompts_only: bool | None = None, params_source: dict | None = None, ): """ @@ -212,6 +217,10 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): bool(params.get("s3_use_virtual_hosted_style", False)) or s3_use_virtual_hosted_style ) + self.s3_log_prompts_only: object = ( + params.get("s3_log_prompts_only") if s3_log_prompts_only is None else s3_log_prompts_only + ) + self.s3_server_side_encryption, self.s3_sse_kms_key_id = resolve_sse_params( params.get("s3_server_side_encryption") or s3_server_side_encryption, params.get("s3_sse_kms_key_id") or s3_sse_kms_key_id, @@ -489,8 +498,13 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): s3_object_download_filename: Final = get_s3_object_download_filename(start_time, standard_logging_payload["id"]) + payload: Final = ( + prompts_only_payload(standard_logging_payload) + if resolve_s3_log_prompts_only(self.s3_log_prompts_only) + else standard_logging_payload + ) return s3BatchLoggingElement( - payload=dict(standard_logging_payload), + payload=dict(payload), s3_object_key=s3_object_key, s3_object_download_filename=s3_object_download_filename, ) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 2ea20143f0c..93d87585496 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -1561,10 +1561,12 @@ class AnthropicMessagesHandler(BaseTranslation): def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: stream_ended: Final = self._check_streaming_has_ended(responses_so_far) + tool_use_fingerprints: Final = self._streamed_tool_use_fingerprints(responses_so_far) return StreamingScanKey( texts=(self.get_streaming_string_so_far(responses_so_far),), - tool_calls=self._streamed_tool_use_fingerprints(responses_so_far) if stream_ended else (), + tool_calls=tool_use_fingerprints if stream_ended else (), stream_ended=stream_ended, + tool_calls_in_flight=bool(tool_use_fingerprints) and not stream_ended, ) @classmethod diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index 4dd0deeb62b..58ed9b38ccb 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -632,6 +632,7 @@ class ModelResponseIterator: self.tool_name_reverse_map: dict[str, str] = tool_name_reverse_map or {} # Generate response ID once per stream to match OpenAI-compatible behavior self.response_id = _generate_id() + self.served_model: str | None = None # Track if we're currently streaming a response_format tool self.is_response_format_tool: bool = False @@ -1067,6 +1068,9 @@ class ModelResponseIterator: } """ message_start_block: Final = MessageStartBlock(**chunk) + start_message: Final = message_start_block["message"] + if "model" in start_message: + self.served_model = start_message["model"] if "usage" in message_start_block["message"]: usage = self._handle_usage(anthropic_usage_chunk=message_start_block["message"]["usage"]) elif type_chunk == "error": @@ -1098,6 +1102,7 @@ class ModelResponseIterator: ], usage=usage, id=self.response_id, + model=self.served_model, ) return returned_chunk diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py index fb6a1c40253..ecaf8f2e7e1 100644 --- a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py +++ b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py @@ -78,6 +78,7 @@ _PROPAGATED_METADATA_KEYS: Final = ( "user_api_key_end_user_id", "user_api_end_user_max_budget", "user_api_key_model_max_budget", + "user_api_key_team_model_max_budget", "user_api_key_user_model_max_budget", "user_api_key_end_user_model_max_budget", "litellm_call_id", @@ -395,9 +396,9 @@ async def _check_summary_model_budget( ``user_api_key_auth`` runs for the client-requested model. Returns True outside the proxy or when no per-model budget is configured. - All three scopes are checked because the summary's spend is charged to all - three: this file propagates the key, user and end-user budgets into the - subrequest's metadata, so enforcing only two of them would let compaction + Every scope is checked because the summary's spend is charged to every + scope: this file propagates the key, team, user and end-user budgets into the + subrequest's metadata, so skipping one of them would let compaction increment a counter it can never be refused by. """ if user_api_key_auth is None: @@ -444,6 +445,26 @@ async def _check_summary_model_budget( ) return False + team_model_max_budget: Final = user_api_key_auth.team_model_max_budget + team_id: Final = user_api_key_auth.team_id + if isinstance(team_model_max_budget, dict) and team_model_max_budget and team_id is not None: + try: + await model_max_budget_limiter.is_team_within_model_budget( + team_id=team_id, + team_model_max_budget=team_model_max_budget, + key_model_max_budget=model_max_budget if isinstance(model_max_budget, dict) else None, + model=summary_model, + ) + except litellm.BudgetExceededError: + return False + except Exception as e: # noqa: BLE001 # a budget gate denies on any failure, as the other scopes do + verbose_logger.warning( + "compact_20260112: unexpected error during team model-budget check for summary_model=%s; denying: %s", + summary_model, + e, + ) + return False + end_user_model_max_budget: Final[dict[str, object] | None] = getattr( user_api_key_auth, "end_user_model_max_budget", None ) diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index f1143425ced..89ad67f0485 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -40,11 +40,15 @@ class StreamingScanKey: """What a streaming guardrail round would hand to ``apply_guardrail``. Two keys compare equal when the round would scan the same content again; ``stream_ended`` stays out of the comparison and only says whether the handler is on its - end-of-stream path, where an empty payload is still scanned today.""" + end-of-stream path, where an empty payload is still scanned today. + ``tool_calls_in_flight`` also stays out of the comparison: it flags that tool + calls have streamed which this round cannot scan yet, so a buffered window + holding them must stay withheld until the end-of-stream scan covers them.""" texts: tuple[str, ...] tool_calls: tuple[str, ...] = () stream_ended: bool = field(default=False, compare=False) + tool_calls_in_flight: bool = field(default=False, compare=False) @property def has_nothing_to_scan(self) -> bool: diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 01e14f2248d..5b232dd0bf4 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -792,10 +792,12 @@ class OpenAIChatCompletionsHandler(BaseTranslation): def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: chunks: Final = tuple(chunk for chunk in responses_so_far if isinstance(chunk, ModelResponseStream)) stream_ended: Final = self._first_choice_has_finished(responses_so_far) + tool_call_fingerprints: Final = self._streamed_tool_call_fingerprints(responses_so_far) return StreamingScanKey( texts=tuple(self._combine_streaming_texts(chunks).values()), - tool_calls=self._streamed_tool_call_fingerprints(responses_so_far) if stream_ended else (), + tool_calls=tool_call_fingerprints if stream_ended else (), stream_ended=stream_ended, + tool_calls_in_flight=bool(tool_call_fingerprints) and not stream_ended, ) @staticmethod @@ -804,7 +806,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): stream_item_fingerprint(tool_call) for chunk in responses_so_far for choice in _stream_chunk_choices(chunk) - for tool_call in stream_item_items(stream_item_field(choice, "delta"), "tool_calls") + for tool_call in _streamed_delta_tool_calls(stream_item_field(choice, "delta")) ) @staticmethod @@ -1342,6 +1344,12 @@ def _stream_chunk_choices(item: object) -> Sequence[object]: return () +def _streamed_delta_tool_calls(delta: object) -> tuple[object, ...]: + function_call: Final = stream_item_field(delta, "function_call") + legacy: Final = () if function_call is None else (function_call,) + return stream_item_items(delta, "tool_calls") + legacy + + def _blocked_stream_identity( exc: "ModifyResponseException", responses_so_far: Sequence[object] ) -> tuple[str, int, str]: diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 27ff55f120c..1ab4811b8a0 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -1175,11 +1175,22 @@ class OpenAIResponsesHandler(BaseTranslation): last_event_type: Final = stream_item_field(last_event, "type") if last_event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE.value: return None - if last_event_type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value: + if last_event_type in _TERMINAL_ENVELOPE_EVENT_TYPES: return self._completed_response_scan_key(stream_item_field(last_event, "response")) return StreamingScanKey( texts=(self.get_streaming_string_so_far(responses_so_far),), - stream_ended=self._check_streaming_has_ended(responses_so_far), + tool_calls_in_flight=self._has_streamed_tool_call_events(responses_so_far), + ) + + @staticmethod + def _has_streamed_tool_call_events(responses_so_far: Sequence[object]) -> bool: + return any( + stream_item_field(event, "type") in _TOOL_CALL_PAYLOAD_EVENT_TYPES + or ( + stream_item_field(event, "type") in _OUTPUT_ITEM_EVENT_TYPES + and stream_item_field(stream_item_field(event, "item"), "type") in _TOOL_CALL_ITEM_TYPES + ) + for event in responses_so_far ) @staticmethod diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 3abf80c74b7..92e5b1c4ff7 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -352,6 +352,17 @@ "supports_function_calling": true, "supports_pdf_input": true }, + "writer.palmyra-vision-7b": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 4096, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-writer-palmyra-vision-7b.html", + "supports_vision": true + }, "amazon.nova-lite-v1:0": { "cache_read_input_token_cost": 1.5e-08, "input_cost_per_token": 6e-08, @@ -3796,6 +3807,53 @@ "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false }, + "azure_ai/gpt-5.5-2026-04-24": { + "deprecation_date": "2027-10-26", + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "cache_read_input_token_cost_priority": 1e-06, + "cache_read_input_token_cost_above_272k_tokens_priority": 2e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, + "input_cost_per_token_priority": 1e-05, + "input_cost_per_token_above_272k_tokens_priority": 2e-05, + "litellm_provider": "azure_ai", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "output_cost_per_token_priority": 6e-05, + "output_cost_per_token_above_272k_tokens_priority": 9e-05, + "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false + }, "azure_ai/gpt-5.4": { "deprecation_date": "2027-09-02", "cache_read_input_token_cost": 2.5e-07, @@ -5311,6 +5369,7 @@ "supports_vision": true }, "azure/gpt-4o": { + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 1.25e-06, "input_cost_per_token": 2.5e-06, "litellm_provider": "azure", @@ -7510,6 +7569,70 @@ "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false }, + "azure/gpt-5.6-sol-2026-07-09": { + "cache_creation_input_token_cost": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost_above_272k_tokens_flex": 5e-06, + "cache_creation_input_token_cost_priority": 1e-05, + "cache_creation_input_token_cost_above_272k_tokens_priority": 2e-05, + "cache_creation_input_token_cost_flex": 2.5e-06, + "cache_read_input_token_cost": 4e-07, + "cache_read_input_token_cost_above_272k_tokens": 8e-07, + "cache_read_input_token_cost_above_272k_tokens_flex": 4e-07, + "cache_read_input_token_cost_priority": 8e-07, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.6e-06, + "cache_read_input_token_cost_flex": 2e-07, + "deprecation_date": "2028-01-11", + "input_cost_per_token": 4e-06, + "input_cost_per_token_above_272k_tokens": 8e-06, + "input_cost_per_token_above_272k_tokens_flex": 4e-06, + "input_cost_per_token_priority": 8e-06, + "input_cost_per_token_above_272k_tokens_priority": 1.6e-05, + "input_cost_per_token_flex": 2e-06, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2e-05, + "output_cost_per_token_above_272k_tokens": 3e-05, + "output_cost_per_token_above_272k_tokens_flex": 1.5e-05, + "output_cost_per_token_priority": 4e-05, + "output_cost_per_token_above_272k_tokens_priority": 6e-05, + "output_cost_per_token_flex": 1e-05, + "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://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false + }, "azure/gpt-5.6-terra": { "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, @@ -7574,6 +7697,70 @@ "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false }, + "azure/gpt-5.6-terra-2026-07-09": { + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06, + "cache_creation_input_token_cost_priority": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-05, + "cache_creation_input_token_cost_flex": 1.25e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_272k_tokens": 4e-07, + "cache_read_input_token_cost_above_272k_tokens_flex": 2e-07, + "cache_read_input_token_cost_priority": 4e-07, + "cache_read_input_token_cost_above_272k_tokens_priority": 8e-07, + "cache_read_input_token_cost_flex": 1e-07, + "deprecation_date": "2028-01-11", + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_272k_tokens": 4e-06, + "input_cost_per_token_above_272k_tokens_flex": 2e-06, + "input_cost_per_token_priority": 4e-06, + "input_cost_per_token_above_272k_tokens_priority": 8e-06, + "input_cost_per_token_flex": 1e-06, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_272k_tokens": 1.8e-05, + "output_cost_per_token_above_272k_tokens_flex": 9e-06, + "output_cost_per_token_priority": 2.4e-05, + "output_cost_per_token_above_272k_tokens_priority": 3.6e-05, + "output_cost_per_token_flex": 6e-06, + "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://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false + }, "azure/gpt-5.6-luna": { "cache_creation_input_token_cost": 2.5e-07, "cache_creation_input_token_cost_above_272k_tokens": 5e-07, @@ -7638,6 +7825,70 @@ "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false }, + "azure/gpt-5.6-luna-2026-07-09": { + "cache_creation_input_token_cost": 2.5e-07, + "cache_creation_input_token_cost_above_272k_tokens": 5e-07, + "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-07, + "cache_creation_input_token_cost_priority": 5e-07, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-06, + "cache_creation_input_token_cost_flex": 1.25e-07, + "cache_read_input_token_cost": 2e-08, + "cache_read_input_token_cost_above_272k_tokens": 4e-08, + "cache_read_input_token_cost_above_272k_tokens_flex": 2e-08, + "cache_read_input_token_cost_priority": 4e-08, + "cache_read_input_token_cost_above_272k_tokens_priority": 8e-08, + "cache_read_input_token_cost_flex": 1e-08, + "deprecation_date": "2028-01-11", + "input_cost_per_token": 2e-07, + "input_cost_per_token_above_272k_tokens": 4e-07, + "input_cost_per_token_above_272k_tokens_flex": 2e-07, + "input_cost_per_token_priority": 4e-07, + "input_cost_per_token_above_272k_tokens_priority": 8e-07, + "input_cost_per_token_flex": 1e-07, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "output_cost_per_token_above_272k_tokens": 1.8e-06, + "output_cost_per_token_above_272k_tokens_flex": 9e-07, + "output_cost_per_token_priority": 2.4e-06, + "output_cost_per_token_above_272k_tokens_priority": 3.6e-06, + "output_cost_per_token_flex": 6e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models" + }, "azure/gpt-6-astra": { "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, @@ -7686,6 +7937,54 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure/gpt-6-astra-2026-09-03": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, + "cache_read_input_token_cost": 1e-06, + "cache_read_input_token_cost_above_272k_tokens": 2e-06, + "input_cost_per_token": 1e-05, + "input_cost_per_token_above_272k_tokens": 2e-05, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "output_cost_per_token_above_272k_tokens": 7.5e-05, + "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://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": false, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "azure/gpt-chat-latest": { "cache_read_input_token_cost": 5e-07, "deprecation_date": "2026-12-02", @@ -8502,6 +8801,63 @@ "output_cost_per_token_batches": 1.5e-05, "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-5.5-2026-04-24": { + "deprecation_date": "2027-10-26", + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "cache_read_input_token_cost_priority": 1.25e-06, + "cache_read_input_token_cost_above_272k_tokens_priority": 2e-06, + "cache_read_input_token_cost_flex": 2.5e-07, + "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, + "input_cost_per_token_above_272k_tokens_flex": 5e-06, + "input_cost_per_token_priority": 1.25e-05, + "input_cost_per_token_above_272k_tokens_priority": 2e-05, + "input_cost_per_token_flex": 2.5e-06, + "litellm_provider": "azure", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "output_cost_per_token_batches": 1.5e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "output_cost_per_token_priority": 7.5e-05, + "output_cost_per_token_above_272k_tokens_priority": 9e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models" + }, "azure/us/gpt-5.5-2026-04-23": { "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, @@ -8550,6 +8906,57 @@ "output_cost_per_token_batches": 1.65e-05, "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/us/gpt-5.5-2026-04-24": { + "deprecation_date": "2027-10-26", + "cache_read_input_token_cost": 5.5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, + "cache_read_input_token_cost_priority": 1.375e-06, + "input_cost_per_token": 5.5e-06, + "input_cost_per_token_batches": 2.75e-06, + "input_cost_per_token_above_272k_tokens": 1.1e-05, + "input_cost_per_token_priority": 1.375e-05, + "output_cost_per_token": 3.3e-05, + "output_cost_per_token_batches": 1.65e-05, + "output_cost_per_token_above_272k_tokens": 4.95e-05, + "output_cost_per_token_priority": 8.25e-05, + "litellm_provider": "azure", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models" + }, "azure/eu/gpt-5.5-2026-04-23": { "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, @@ -8598,6 +9005,57 @@ "output_cost_per_token_batches": 1.65e-05, "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/eu/gpt-5.5-2026-04-24": { + "deprecation_date": "2027-10-26", + "cache_read_input_token_cost": 5.5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, + "cache_read_input_token_cost_priority": 1.375e-06, + "input_cost_per_token": 5.5e-06, + "input_cost_per_token_batches": 2.75e-06, + "input_cost_per_token_above_272k_tokens": 1.1e-05, + "input_cost_per_token_priority": 1.375e-05, + "output_cost_per_token": 3.3e-05, + "output_cost_per_token_batches": 1.65e-05, + "output_cost_per_token_above_272k_tokens": 4.95e-05, + "output_cost_per_token_priority": 8.25e-05, + "litellm_provider": "azure", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models" + }, "azure/gpt-5.5-pro": { "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, @@ -9055,6 +9513,36 @@ "supports_vision": true, "supports_pdf_input": true }, + "azure/gpt-image-2.5-flare": { + "deprecation_date": "2027-09-09", + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_image_token": 8e-06, + "litellm_provider": "azure", + "mode": "image_generation", + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "azure/gpt-image-2.5-sunburst": { + "deprecation_date": "2027-09-09", + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_image_token": 8e-06, + "litellm_provider": "azure", + "mode": "image_generation", + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, "azure/gpt-image-2-2026-04-21": { "cache_read_input_token_cost": 1.25e-06, "deprecation_date": "2027-10-21", @@ -11371,6 +11859,7 @@ "output_cost_per_token": 2.5e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, @@ -11389,6 +11878,7 @@ "output_cost_per_token": 2.5e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, @@ -13904,7 +14394,7 @@ "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "anthropic", - "max_input_tokens": 200000, + "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", @@ -13939,7 +14429,7 @@ "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "anthropic", - "max_input_tokens": 200000, + "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", @@ -13960,7 +14450,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "source": "https://docs.anthropic.com/en/docs/about-claude/pricing" }, "claude-sonnet-5": { "deprecation_date": "2027-06-30", @@ -15089,6 +15580,21 @@ "supports_function_calling": true, "supports_tool_choice": true }, + "command-a-plus-05-2026": { + "input_cost_per_token": 0.0, + "litellm_provider": "cohere_chat", + "max_input_tokens": 128000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 0.0, + "source": "https://docs.cohere.com/docs/command-a-plus", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "command-light": { "input_cost_per_token": 3e-07, "litellm_provider": "cohere_chat", @@ -26884,13 +27390,71 @@ "supports_image_size": false }, "gemini/gemini-flash-latest": { + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_token": 7.5e-07, + "litellm_provider": "gemini", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 3.75e-06, + "output_cost_per_token": 3.75e-06, + "rpm": 15, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_output": false, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_vision": true, + "supports_web_search": true, + "tpm": 250000, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "google_maps_grounding_cost_per_query": 0.014, + "cache_read_input_token_cost_flex": 3.75e-08, + "cache_read_input_token_cost_priority": 1.35e-07, + "input_cost_per_token_batches": 3.75e-07, + "input_cost_per_token_flex": 3.75e-07, + "input_cost_per_token_priority": 1.35e-06, + "output_cost_per_token_batches": 1.875e-06, + "output_cost_per_token_flex": 1.875e-06, + "output_cost_per_token_priority": 6.75e-06, + "prompt_cache_min_tokens": 4096, + "supports_audio_input": true, + "supports_native_streaming": true, + "supports_video_input": true, + "web_search_billing_unit": "per_query" + }, + "gemini/gemini-flash-lite-latest": { "cache_read_input_token_cost": 3e-08, - "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", "output_cost_per_reasoning_token": 2.5e-06, "output_cost_per_token": 2.5e-06, @@ -26924,58 +27488,23 @@ "supports_web_search": true, "tpm": 250000, "search_context_cost_per_query": { - "search_context_size_low": 0.035, - "search_context_size_medium": 0.035, - "search_context_size_high": 0.035 + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 }, - "google_maps_grounding_cost_per_query": 0.025 - }, - "gemini/gemini-flash-lite-latest": { - "cache_read_input_token_cost": 1e-08, - "input_cost_per_audio_token": 3e-07, - "input_cost_per_token": 1e-07, - "litellm_provider": "gemini", - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, - "mode": "chat", - "output_cost_per_reasoning_token": 4e-07, - "output_cost_per_token": 4e-07, - "rpm": 15, - "source": "https://ai.google.dev/gemini-api/docs/pricing", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 250000, - "search_context_cost_per_query": { - "search_context_size_low": 0.035, - "search_context_size_medium": 0.035, - "search_context_size_high": 0.035 - }, - "google_maps_grounding_cost_per_query": 0.025 + "google_maps_grounding_cost_per_query": 0.014, + "cache_read_input_token_cost_flex": 2e-08, + "cache_read_input_token_cost_priority": 5e-08, + "input_cost_per_token_batches": 1.5e-07, + "input_cost_per_token_flex": 1.5e-07, + "input_cost_per_token_priority": 5.4e-07, + "output_cost_per_token_batches": 1.25e-06, + "output_cost_per_token_flex": 1.25e-06, + "output_cost_per_token_priority": 4.5e-06, + "supports_audio_input": true, + "supports_native_streaming": true, + "supports_video_input": true, + "web_search_billing_unit": "per_query" }, "gemini/gemini-2.5-flash-lite-preview-06-17": { "deprecation_date": "2025-11-18", @@ -32153,7 +32682,7 @@ "input_cost_per_token": 4e-06, "input_cost_per_token_above_272k_tokens": 8e-06, "litellm_provider": "openai", - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -32192,7 +32721,7 @@ "input_cost_per_token": 4e-06, "input_cost_per_token_above_272k_tokens": 8e-06, "litellm_provider": "openai", - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", @@ -34977,6 +35506,7 @@ "supports_tool_choice": true }, "inception/mercury-2.5": { + "cache_read_input_token_cost": 2e-08, "input_cost_per_token": 2e-07, "litellm_provider": "inception", "max_input_tokens": 260000, @@ -34986,6 +35516,7 @@ "output_cost_per_token": 7.5e-07, "source": "https://docs.inceptionlabs.ai/get-started/models", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true @@ -36525,6 +37056,57 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "mistral/zai-glm-5-3": { + "cache_read_input_token_cost": 1.4e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "mistral", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://docs.mistral.ai/models/zai-glm-5-3", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "mistral/zai-glm-5": { + "cache_read_input_token_cost": 1.4e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "mistral", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://docs.mistral.ai/models/zai-glm-5-3", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "mistral/zai-glm-latest": { + "cache_read_input_token_cost": 1.4e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "mistral", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://docs.mistral.ai/models/zai-glm-5-3", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "mistral/glm-5-2": { "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.4e-06, @@ -38151,6 +38733,15 @@ "supports_reasoning": true, "source": "https://tokenfactory.nebius.com/models/catalog/text2text/deepseek-ai%2FDeepSeek-V4-Pro" }, + "nebius/deepseek-ai/DeepSeek-V4-Pro-0813": { + "input_cost_per_token": 1.32e-06, + "litellm_provider": "nebius", + "mode": "chat", + "output_cost_per_token": 3.96e-06, + "source": "https://tokenfactory.nebius.com/models/catalog/text2text/deepseek-ai%2FDeepSeek-V4-Pro-0813", + "supports_function_calling": true, + "supports_reasoning": true + }, "nebius/MiniMaxAI/MiniMax-M2.5": { "max_tokens": 196608, "max_input_tokens": 196608, @@ -38417,6 +39008,17 @@ "supports_reasoning": true, "source": "https://tokenfactory.nebius.com/models/catalog/text2text/zai-org%2FGLM-5.2" }, + "nebius/zai-org/GLM-5.3": { + "input_cost_per_token": 1.4e-06, + "litellm_provider": "nebius", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://tokenfactory.nebius.com/models/catalog/text2text/zai-org%2FGLM-5.3", + "supports_function_calling": true, + "supports_reasoning": true + }, "nebius/zai-org/GLM-5.3-Flash": { "max_tokens": 1024000, "max_input_tokens": 1024000, @@ -40250,15 +40852,16 @@ "supports_tool_choice": true }, "openrouter/deepseek/deepseek-chat": { - "input_cost_per_token": 3.2e-07, + "input_cost_per_token": 2.574e-07, "litellm_provider": "openrouter", "max_input_tokens": 65536, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 8.9e-07, + "output_cost_per_token": 1.0287e-06, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "source": "https://openrouter.ai/api/v1/models" }, "openrouter/deepseek/deepseek-chat-v3-0324": { "input_cost_per_token": 2.5e-07, @@ -40351,21 +40954,21 @@ "supports_tool_choice": true }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 8.59908e-07, + "input_cost_per_token": 1.6e-06, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.719816e-06, + "output_cost_per_token": 3.2e-06, "source": "https://openrouter.ai/deepseek/deepseek-v4-pro", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.1659e-08 + "cache_read_input_token_cost": 1.35e-07 }, "openrouter/deepseek/deepseek-v4.1-flash": { "input_cost_per_token": 1.5e-07, @@ -40660,12 +41263,13 @@ "supports_vision": true }, "openrouter/gryphe/mythomax-l2-13b": { - "input_cost_per_token": 6e-08, + "input_cost_per_token": 8e-08, "litellm_provider": "openrouter", "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 6e-08, - "supports_tool_choice": true + "output_cost_per_token": 1.1e-07, + "supports_tool_choice": true, + "source": "https://openrouter.ai/api/v1/models" }, "openrouter/mancer/weaver": { "input_cost_per_token": 4e-07, @@ -40801,14 +41405,15 @@ "max_output_tokens": 131072 }, "openrouter/mistralai/mistral-small-3.2-24b-instruct": { - "input_cost_per_token": 7.5e-08, + "input_cost_per_token": 9.375e-08, "litellm_provider": "openrouter", "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 2e-07, + "output_cost_per_token": 2.5e-07, "supports_tool_choice": true, "max_input_tokens": 128000, - "max_output_tokens": 128000 + "max_output_tokens": 128000, + "source": "https://openrouter.ai/api/v1/models" }, "openrouter/mistralai/mixtral-8x22b-instruct": { "input_cost_per_token": 2e-06, @@ -41329,13 +41934,13 @@ "supports_tool_choice": true }, "openrouter/qwen/qwen3-235b-a22b-2507": { - "input_cost_per_token": 2.2e-07, + "input_cost_per_token": 8.75e-08, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 262144, "max_tokens": 262144, "mode": "chat", - "output_cost_per_token": 8.8e-07, + "output_cost_per_token": 3.5e-07, "source": "https://openrouter.ai/qwen/qwen3-235b-a22b-2507", "supports_function_calling": true, "supports_tool_choice": true @@ -41368,13 +41973,13 @@ "supports_vision": true }, "openrouter/qwen/qwen3.5-35b-a3b": { - "input_cost_per_token": 3.125e-07, + "input_cost_per_token": 1.625e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", - "output_cost_per_token": 1.25e-06, + "output_cost_per_token": 1.3e-06, "source": "https://openrouter.ai/qwen/qwen3.5-35b-a3b", "supports_function_calling": true, "supports_reasoning": true, @@ -44390,7 +44995,7 @@ "deprecation_date": "2026-06-04", "input_cost_per_token": 2e-06, "litellm_provider": "together_ai", - "max_input_tokens": 256000, + "max_input_tokens": 262144, "mode": "chat", "output_cost_per_token": 2e-06, "source": "https://api.together.ai/v1/models", @@ -44460,7 +45065,7 @@ "supports_parallel_function_calling": true, "supports_reasoning": true, "supports_tool_choice": true, - "max_input_tokens": 128000, + "max_input_tokens": 131072, "max_output_tokens": 16384 }, "together_ai/meta-llama/Llama-3.2-3B-Instruct-Turbo": { @@ -44643,7 +45248,7 @@ "deprecation_date": "2026-04-02", "input_cost_per_token": 2e-07, "litellm_provider": "together_ai", - "max_input_tokens": 128000, + "max_input_tokens": 131072, "mode": "chat", "output_cost_per_token": 1.1e-06, "source": "https://api.together.ai/v1/models", @@ -44870,6 +45475,20 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "together_ai/deepseek-ai/DeepSeek-V4.1-Flash": { + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://api.together.xyz/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "together_ai/deepseek-ai/DeepSeek-V4-Pro": { "deprecation_date": "2026-08-27", "cache_read_input_token_cost": 2e-07, @@ -45199,9 +45818,10 @@ "output_cost_per_token": 1.25e-05, "supports_function_calling": true, "supports_pdf_input": true, - "supports_prompt_caching": false, + "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 6.25e-07 }, "us.amazon.nova-pro-v1:0": { "cache_read_input_token_cost": 2e-07, @@ -49453,6 +50073,7 @@ "output_cost_per_token": 5e-07, "source": "https://docs.x.ai/developers/models", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, @@ -49469,6 +50090,7 @@ "output_cost_per_token": 5e-07, "source": "https://docs.x.ai/developers/models", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -49489,6 +50111,7 @@ "output_cost_per_token_above_200k_tokens": 5e-06, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, @@ -49508,6 +50131,7 @@ "output_cost_per_token_above_200k_tokens": 5e-06, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -49528,6 +50152,7 @@ "output_cost_per_token_above_200k_tokens": 5e-06, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -49547,6 +50172,7 @@ "output_cost_per_token_above_200k_tokens": 1.2e-05, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -56015,6 +56641,73 @@ "input_cost_per_second": 8.33333333333e-05, "supports_response_schema": false }, + "gemini-3.8-live": { + "input_cost_per_audio_token": 3e-06, + "input_cost_per_image_token": 1e-06, + "input_cost_per_token": 7.5e-07, + "input_cost_per_video_per_second": 3.3333333333333335e-05, + "litellm_provider": "gemini", + "max_input_tokens": 131072, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "realtime", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 4.5e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_vision": true, + "supports_web_search": true, + "gemini_audio_only_live": true + }, + "gemini-3.8-live-extended-thinking": { + "input_cost_per_audio_token": 3e-06, + "input_cost_per_image_token": 1e-06, + "input_cost_per_token": 7.5e-07, + "input_cost_per_video_per_second": 3.3333333333333335e-05, + "litellm_provider": "gemini", + "max_input_tokens": 131072, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "realtime", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 4.5e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_vision": true, + "supports_web_search": true, + "gemini_audio_only_live": true, + "supports_reasoning": true + }, "gemini/gemini-2.5-flash-native-audio-latest": { "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 5e-07, @@ -56180,18 +56873,17 @@ "supports_web_search": false }, "gemini-flash-latest": { - "cache_read_input_token_cost": 3e-08, - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_token": 7.5e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", - "output_cost_per_reasoning_token": 2.5e-06, - "output_cost_per_token": 2.5e-06, + "output_cost_per_reasoning_token": 3.75e-06, + "output_cost_per_token": 3.75e-06, "rpm": 100000, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -56220,25 +56912,37 @@ "supports_web_search": true, "tpm": 8000000, "search_context_cost_per_query": { - "search_context_size_low": 0.035, - "search_context_size_medium": 0.035, - "search_context_size_high": 0.035 + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 }, - "google_maps_grounding_cost_per_query": 0.025 + "google_maps_grounding_cost_per_query": 0.014, + "cache_read_input_token_cost_flex": 3.75e-08, + "cache_read_input_token_cost_priority": 1.35e-07, + "input_cost_per_token_batches": 3.75e-07, + "input_cost_per_token_flex": 3.75e-07, + "input_cost_per_token_priority": 1.35e-06, + "output_cost_per_token_batches": 1.875e-06, + "output_cost_per_token_flex": 1.875e-06, + "output_cost_per_token_priority": 6.75e-06, + "prompt_cache_min_tokens": 4096, + "supports_audio_input": true, + "supports_native_streaming": true, + "supports_video_input": true, + "web_search_billing_unit": "per_query" }, "gemini-flash-lite-latest": { - "cache_read_input_token_cost": 1e-08, - "input_cost_per_audio_token": 3e-07, - "input_cost_per_token": 1e-07, + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 3e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", - "output_cost_per_reasoning_token": 4e-07, - "output_cost_per_token": 4e-07, + "output_cost_per_reasoning_token": 2.5e-06, + "output_cost_per_token": 2.5e-06, "rpm": 15, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-lite", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -56267,29 +56971,42 @@ "supports_web_search": true, "tpm": 250000, "search_context_cost_per_query": { - "search_context_size_low": 0.035, - "search_context_size_medium": 0.035, - "search_context_size_high": 0.035 + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 }, - "google_maps_grounding_cost_per_query": 0.025 + "google_maps_grounding_cost_per_query": 0.014, + "cache_read_input_token_cost_flex": 2e-08, + "cache_read_input_token_cost_priority": 5e-08, + "input_cost_per_token_batches": 1.5e-07, + "input_cost_per_token_flex": 1.5e-07, + "input_cost_per_token_priority": 5.4e-07, + "output_cost_per_token_batches": 1.25e-06, + "output_cost_per_token_flex": 1.25e-06, + "output_cost_per_token_priority": 4.5e-06, + "supports_audio_input": true, + "supports_native_streaming": true, + "supports_video_input": true, + "web_search_billing_unit": "per_query" }, "gemini-pro-latest": { - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, "rpm": 2000, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", - "/v1/completions" + "/v1/completions", + "/v1/batch" ], "supported_modalities": [ "text", @@ -56313,29 +57030,45 @@ "supports_web_search": true, "tpm": 800000, "search_context_cost_per_query": { - "search_context_size_low": 0.035, - "search_context_size_medium": 0.035, - "search_context_size_high": 0.035 + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 }, - "google_maps_grounding_cost_per_query": 0.025 + "google_maps_grounding_cost_per_query": 0.014, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, + "cache_read_input_token_cost_priority": 3.6e-07, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, + "input_cost_per_token_batches": 1e-06, + "input_cost_per_token_priority": 3.6e-06, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, + "output_cost_per_token_batches": 6e-06, + "output_cost_per_token_priority": 2.16e-05, + "prompt_cache_min_tokens": 4096, + "supports_native_streaming": true, + "supports_url_context": true, + "web_search_billing_unit": "per_query", + "cache_read_input_token_cost_flex": 2e-07, + "input_cost_per_token_flex": 1e-06, + "output_cost_per_token_flex": 6e-06 }, "gemini/gemini-pro-latest": { - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, "rpm": 2000, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", - "/v1/completions" + "/v1/completions", + "/v1/batch" ], "supported_modalities": [ "text", @@ -56359,11 +57092,26 @@ "supports_web_search": true, "tpm": 800000, "search_context_cost_per_query": { - "search_context_size_low": 0.035, - "search_context_size_medium": 0.035, - "search_context_size_high": 0.035 + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 }, - "google_maps_grounding_cost_per_query": 0.025 + "google_maps_grounding_cost_per_query": 0.014, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, + "cache_read_input_token_cost_priority": 3.6e-07, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, + "input_cost_per_token_batches": 1e-06, + "input_cost_per_token_priority": 3.6e-06, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, + "output_cost_per_token_batches": 6e-06, + "output_cost_per_token_priority": 2.16e-05, + "prompt_cache_min_tokens": 4096, + "supports_native_streaming": true, + "supports_url_context": true, + "web_search_billing_unit": "per_query", + "cache_read_input_token_cost_flex": 2e-07, + "input_cost_per_token_flex": 1e-06, + "output_cost_per_token_flex": 6e-06 }, "gemini-exp-1206": { "cache_read_input_token_cost": 3e-08, @@ -57263,6 +58011,38 @@ } ] }, + "volcengine/doubao-seed-2-1-pro-260628": { + "cache_read_input_token_cost": 1.725e-07, + "input_cost_per_token": 8.625e-07, + "litellm_provider": "volcengine", + "max_input_tokens": 256000, + "max_output_tokens": 256000, + "max_tokens": 256000, + "mode": "chat", + "output_cost_per_token": 4.3125e-06, + "source": "https://www.volcengine.com/docs/82379/1544106", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": false, + "supports_vision": true + }, + "volcengine/doubao-seed-2-1-turbo-260628": { + "cache_read_input_token_cost": 8.625e-08, + "input_cost_per_token": 4.3125e-07, + "litellm_provider": "volcengine", + "max_input_tokens": 256000, + "max_output_tokens": 256000, + "max_tokens": 256000, + "mode": "chat", + "output_cost_per_token": 2.15625e-06, + "source": "https://www.volcengine.com/docs/82379/1544106", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": false, + "supports_vision": true + }, "volcengine/doubao-seed-2-0-lite-260215": { "litellm_provider": "volcengine", "max_input_tokens": 256000, @@ -58989,7 +59769,8 @@ "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_response_schema": true }, "fireworks_ai/accounts/fireworks/models/kimi-k3": { "cache_read_input_token_cost": 3e-07, @@ -59065,7 +59846,8 @@ "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_response_schema": true }, "fireworks_ai/glm-5p2-fast": { "cache_read_input_token_cost": 2.1e-07, @@ -62056,6 +62838,36 @@ "supports_tool_choice": true, "supports_vision": false }, + "fireworks_ai/accounts/fireworks/routers/glm-5p3-fast": { + "cache_read_input_token_cost": 3.9e-07, + "input_cost_per_token": 2.1e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 6.6e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "fireworks_ai/glm-5p3-fast": { + "cache_read_input_token_cost": 3.9e-07, + "input_cost_per_token": 2.1e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 6.6e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, "fireworks_ai/accounts/fireworks/models/glm-5p3-flash": { "cache_read_input_token_cost": 3e-08, "cache_read_input_token_cost_priority": 3.75e-08, @@ -64225,9 +65037,9 @@ "supports_prompt_caching": true }, "openrouter/z-ai/glm-5.3-flash": { - "input_cost_per_token": 1.5e-07, - "output_cost_per_token": 5e-07, - "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 9e-08, + "output_cost_per_token": 3e-07, + "cache_read_input_token_cost": 1.8e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, "max_output_tokens": 131072, @@ -64277,9 +65089,9 @@ "supports_prompt_caching": true }, "openrouter/qwen/qwen3.8-27b": { - "input_cost_per_token": 4.2e-07, - "output_cost_per_token": 3e-06, - "cache_read_input_token_cost": 8.5e-08, + "input_cost_per_token": 2.14e-07, + "output_cost_per_token": 2.55e-06, + "cache_read_input_token_cost": 1.5e-07, "litellm_provider": "openrouter", "max_input_tokens": 1000000, "max_output_tokens": 131072, @@ -64362,9 +65174,9 @@ "supports_prompt_caching": true }, "openrouter/deepseek/deepseek-v4-flash-0731": { - "input_cost_per_token": 6.5e-08, - "output_cost_per_token": 1.8e-07, - "cache_read_input_token_cost": 1.6e-08, + "input_cost_per_token": 6e-08, + "output_cost_per_token": 1.2e-07, + "cache_read_input_token_cost": 1.2e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, "max_output_tokens": 943718, @@ -64432,9 +65244,9 @@ "supports_vision": false }, "openrouter/moonshotai/kimi-k3": { - "input_cost_per_token": 2.1e-06, - "output_cost_per_token": 1.053e-05, - "cache_read_input_token_cost": 2.35e-07, + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "cache_read_input_token_cost": 3e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 943718, @@ -64531,9 +65343,9 @@ "supports_prompt_caching": true }, "openrouter/z-ai/glm-5.2": { - "input_cost_per_token": 6e-07, - "output_cost_per_token": 2e-06, - "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.4e-06, + "output_cost_per_token": 4.4e-06, + "cache_read_input_token_cost": 1.4e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 131072, @@ -64564,9 +65376,9 @@ "supports_vision": false }, "openrouter/moonshotai/kimi-k2.7-code": { - "input_cost_per_token": 7.1e-07, - "output_cost_per_token": 3.5e-06, - "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 7.062e-07, + "output_cost_per_token": 3.21e-06, + "cache_read_input_token_cost": 1.8e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 235929, @@ -64868,8 +65680,8 @@ "supports_prompt_caching": true }, "openrouter/google/gemma-4-26b-a4b-it": { - "input_cost_per_token": 4.2e-08, - "output_cost_per_token": 2.2e-07, + "input_cost_per_token": 9e-08, + "output_cost_per_token": 3e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 16384, @@ -65416,8 +66228,8 @@ "supports_vision": true }, "openrouter/qwen/qwen3-vl-30b-a3b-instruct": { - "input_cost_per_token": 1.5e-07, - "output_cost_per_token": 6e-07, + "input_cost_per_token": 1.3e-07, + "output_cost_per_token": 5.2e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 16384, @@ -65677,8 +66489,8 @@ "supports_vision": false }, "openrouter/qwen/qwen3-30b-a3b-instruct-2507": { - "input_cost_per_token": 9e-08, - "output_cost_per_token": 3e-07, + "input_cost_per_token": 4.815e-08, + "output_cost_per_token": 1.9305e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 32000, @@ -65737,7 +66549,7 @@ "supports_vision": false }, "openrouter/minimax/minimax-m1": { - "input_cost_per_token": 5.5e-07, + "input_cost_per_token": 4e-07, "output_cost_per_token": 2.2e-06, "litellm_provider": "openrouter", "max_input_tokens": 1000000, @@ -65876,8 +66688,8 @@ "supports_vision": false }, "openrouter/qwen/qwen3-14b": { - "input_cost_per_token": 2.275e-07, - "output_cost_per_token": 9.1e-07, + "input_cost_per_token": 1.2e-07, + "output_cost_per_token": 2.4e-07, "litellm_provider": "openrouter", "max_input_tokens": 131072, "max_output_tokens": 16384, @@ -65939,8 +66751,8 @@ "supports_prompt_caching": true }, "openrouter/meta-llama/llama-4-maverick": { - "input_cost_per_token": 2e-07, - "output_cost_per_token": 6.96e-07, + "input_cost_per_token": 1.875e-07, + "output_cost_per_token": 6.525e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 115200, @@ -66361,14 +67173,6 @@ "supports_response_schema": true, "supports_vision": false }, - "fireworks_ai/accounts/fireworks/routers/glm-5p3-fast": { - "cache_read_input_token_cost": 3.9e-07, - "input_cost_per_token": 2.1e-06, - "litellm_provider": "fireworks_ai", - "mode": "chat", - "output_cost_per_token": 6.6e-06, - "source": "https://api.fireworks.ai/v1/serverless/models" - }, "together_ai/arcee-ai/trinity-mini": { "input_cost_per_token": 4.5e-08, "litellm_provider": "together_ai", @@ -66408,14 +67212,6 @@ "output_cost_per_token": 1.6e-06, "source": "https://api.together.ai/v1/models" }, - "together_ai/deepseek-ai/DeepSeek-V4.1-Flash": { - "cache_read_input_token_cost": 6e-09, - "input_cost_per_token": 3e-07, - "litellm_provider": "together_ai", - "mode": "chat", - "output_cost_per_token": 1.2e-06, - "source": "https://api.together.ai/v1/models" - }, "vertex_ai/gemini-2.5-flash-native-audio": { "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 5e-07, diff --git a/litellm/models/verification_token.py b/litellm/models/verification_token.py index 06ff877a41a..89048c56a9f 100644 --- a/litellm/models/verification_token.py +++ b/litellm/models/verification_token.py @@ -18,6 +18,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase): key_name: str | None = None key_alias: str | None = None spend: float = 0.0 + total_spend: float = 0.0 max_budget: float | None = None expires: str | datetime | None = None models: list = [] @@ -69,6 +70,7 @@ class LiteLLM_DeletedVerificationToken(LiteLLM_VerificationToken): """Audit record for deleted keys; mirrors the token plus deletion metadata.""" id: str | None = None + organization_id: str | None = None deleted_at: datetime | None = None deleted_by: str | None = None deleted_by_api_key: str | None = None diff --git a/litellm/ocr/input.py b/litellm/ocr/input.py deleted file mode 100644 index bcb448371c4..00000000000 --- a/litellm/ocr/input.py +++ /dev/null @@ -1,112 +0,0 @@ -from collections.abc import Mapping -from os import PathLike -from typing import Final, Literal, Protocol, cast # noqa: TID251 # native callables are validated when loaded - -from typing_extensions import NotRequired, ReadOnly, TypedDict - -from litellm.rust_bridge.bindings import NativeBinding -from litellm.rust_bridge.configuration import rust_ocr_enabled - - -class FileReader(Protocol): - def read(self) -> bytes | str: ... - - -class FileDocument(TypedDict): - type: ReadOnly[Literal["file"]] - file: ReadOnly[bytes | PathLike[str] | FileReader] - mime_type: ReadOnly[NotRequired[str]] - - -class NativeFileDocument(Protocol): - def __call__(self, document: Mapping[str, object]) -> dict[str, str]: ... - - -class NativeUploadDocument(Protocol): - def __call__(self, file_content: bytes, file_name: str | None, content_type: str | None) -> dict[str, str]: ... - - -class NativeMimeType(Protocol): - def __call__(self, file_name: str) -> str: ... - - -_FILE_DOCUMENT: Final = NativeBinding( - "_ocr_file_document", - validate=lambda value: ( - cast( # cast-ok: native export owns the callable signature - NativeFileDocument, value - ) - if callable(value) - else None - ), -) -_UPLOAD_DOCUMENT: Final = NativeBinding( - "_ocr_upload_document", - validate=lambda value: ( - cast( # cast-ok: native export owns the callable signature - NativeUploadDocument, value - ) - if callable(value) - else None - ), -) -_MAX_FILE_BYTES: Final = NativeBinding( - "_OCR_MAX_FILE_BYTES", validate=lambda value: value if isinstance(value, int) and value > 0 else None -) -_MIME_TYPE: Final = NativeBinding( - "_ocr_mime_type", - validate=lambda value: ( - cast( # cast-ok: native export owns the callable signature - NativeMimeType, value - ) - if callable(value) - else None - ), -) -_PYTHON_MAX_FILE_BYTES: Final = 50 * 1024 * 1024 - - -def get_mime_type(file_path: str) -> str: - native: Final = _MIME_TYPE.load() if rust_ocr_enabled() else None - if native is None: - from litellm.ocr import legacy - - return legacy.get_mime_type(file_path) - return native(file_path) - - -def get_max_file_bytes() -> int: - limit: Final = _MAX_FILE_BYTES.load() if rust_ocr_enabled() else None - if limit is None: - return _PYTHON_MAX_FILE_BYTES - return limit - - -def convert_file_document_to_url_document(document: FileDocument) -> dict[str, str]: - native: Final = _FILE_DOCUMENT.load() if rust_ocr_enabled() else None - if native is None: - from litellm.ocr import legacy - - return legacy.convert_file_document_to_url_document(document) - return native(document) - - -def convert_upload_to_url_document( - file_content: bytes, filename: str | None, content_type: str | None -) -> dict[str, str]: - native: Final = _UPLOAD_DOCUMENT.load() if rust_ocr_enabled() else None - if native is None: - from litellm.ocr import legacy - - if len(file_content) > _PYTHON_MAX_FILE_BYTES: - raise ValueError("OCR file exceeds the size limit") - content_mime: Final = content_type.split(";")[0].strip() if content_type else None - mime_type: Final = ( - legacy.get_mime_type(filename) - if filename and (not content_mime or content_mime == "application/octet-stream") - else content_mime or "application/octet-stream" - ) - return legacy.convert_file_document_to_url_document( - {"type": "file", "file": file_content, "mime_type": mime_type} - ) - return native(file_content, filename, content_type) diff --git a/litellm/ocr/legacy.py b/litellm/ocr/legacy.py index a742be274b3..f0cf6cc82cc 100644 --- a/litellm/ocr/legacy.py +++ b/litellm/ocr/legacy.py @@ -11,7 +11,7 @@ from collections.abc import Coroutine, Mapping from dataclasses import dataclass from io import IOBase from types import MappingProxyType -from typing import Final, cast # noqa: TID251 # adapters preserve the legacy untyped contracts +from typing import Final, Protocol, cast # noqa: TID251 # adapters preserve the legacy untyped contracts import httpx @@ -26,7 +26,6 @@ from litellm.llms.base_llm.ocr.transformation import ( parse_ocr_request_format, ) from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler -from litellm.ocr.input import FileReader from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import CustomPricingLiteLLMParams from litellm.utils import ProviderConfigManager, client @@ -34,6 +33,10 @@ from litellm.utils import ProviderConfigManager, client base_llm_http_handler: Final = BaseLLMHTTPHandler() +class FileReader(Protocol): + def read(self) -> bytes | str: ... + + @dataclass(frozen=True, slots=True) class _PreparedOCRRequest: model: str diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 382c5d6aae4..c6371c0c33f 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -5,7 +5,7 @@ import httpx from litellm.llms.base_llm.ocr.transformation import OCRResponse from litellm.ocr import legacy -from litellm.ocr.input import convert_file_document_to_url_document, get_mime_type +from litellm.ocr.legacy import convert_file_document_to_url_document, get_mime_type from litellm.rust_bridge.bindings import native_exception_types from litellm.rust_bridge.configuration import rust_ocr_enabled from litellm.rust_bridge.ocr import LiteLLMOcrRequest diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 63d0bfcc5b8..14e3635f079 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2004,6 +2004,13 @@ RouterSettingsDict = Annotated[ class NewTeamRequest(TeamBase): router_settings: RouterSettingsDict | None = None model_aliases: dict | None = None + model_max_budget: GenericBudgetConfigType | None = Field( + default=None, + description=( + "Max budget per model for every key on the team, overridable per key " + "(e.g. {'gpt-4o': {'max_budget': 10, 'budget_duration': '1d'}})" + ), + ) tags: list | None = None guardrails: list[str] | None = None policies: list[str] | None = None @@ -2105,6 +2112,13 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase): access_group_ids: list[str] | None = None budget_limits: list[BudgetLimitEntry] | None = None # multiple concurrent budget windows default_team_member_models: list[str] | None = None # default allowed_models seeded onto new team members + model_max_budget: GenericBudgetConfigType | None = Field( + default=None, + description=( + "Max budget per model for every key on the team, overridable per key " + "(e.g. {'gpt-4o': {'max_budget': 10, 'budget_duration': '1d'}})" + ), + ) class PatchTeamRequest(UpdateTeamRequest): @@ -3032,6 +3046,7 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken): team_tpd_limit: int | None = None team_max_budget: float | None = None team_soft_budget: float | None = None + team_model_max_budget: dict[str, object] | None = None team_models: list = [] team_blocked: bool = False soft_budget: float | None = None @@ -3710,6 +3725,7 @@ class AllCallbacks(LiteLLMPydanticObjectBase): "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION_NAME", + "S3_LOG_PROMPTS_ONLY", ], ) @@ -4462,6 +4478,7 @@ class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable): # Parent org's model ceiling, reported only to callers who can manage the team. # None = no org or not a manager; [] or ["all-proxy-models"] = no ceiling. organization_models: list[str] | None = None + model_max_budget_usage: Mapping[str, Mapping[str, object]] | None = None class TeamInfoResponseObject(TypedDict): diff --git a/litellm/proxy/auth/fallback_budget.py b/litellm/proxy/auth/fallback_budget.py new file mode 100644 index 00000000000..e356f8acc7d --- /dev/null +++ b/litellm/proxy/auth/fallback_budget.py @@ -0,0 +1,166 @@ +""" +Enforce the caller's budget against router fallback targets. + +Budget is checked once, during auth, against the *requested* model group. A zero-cost group takes +`_is_model_cost_zero`'s bypass and waives every budget check; the router then picks a fallback +target after auth, inside `run_async_fallback`, and nothing re-checks budget on the group that +actually bills. So a free model with a paid fallback spends without a gate. + +This predicate is injected into the router to re-check budget for each fallback target before it is +attempted, mirroring `fallback_model_access.py`. It deliberately leaves the primary attempt alone: +a zero-cost model is never blocked by budget, and only the paid fallback is refused. On by default; +set `general_settings.enforce_fallback_budget: false` to restore the unguarded behaviour. + +Scope: the key's and the user's `max_budget`. Not covered yet, and each needs a read-only evaluation +path before it can be: team, team-member, end-user, org, global and per-model budgets, whose +auth-path functions enforce rather than report (they raise), so reusing them would fire threshold +alerts and take spend reservations for a target that is then skipped; and the key's rolling +`budget_limits` windows, whose accumulated spend lives only in per-window counters +(`spend:key:{token}:window:{budget_duration}`), so enforcing them means more counter reads on the +fallback path rather than reusing state auth already loaded. + +Two known limitations of that narrow scope, both shared with `fallback_model_access.py`: + +* This reads the spend counter, it does not reserve against it. Requests already in flight all + observe the same pre-billing figure, so a cap can be crossed by roughly the number of concurrent + fallbacks times their cost. Auth-time enforcement avoids this by pre-filling the counter through + `reserve_budget_for_request`, which the zero-cost bypass skips. Turning the soft cap into a hard + one means reserving per fallback attempt and reconciling on completion. +* A request that reaches the router without `metadata["user_api_key_auth"]` is not restricted. + Only `add_litellm_data_to_request` populates that key, so endpoints that assemble metadata by + hand (for example `/queue/chat/completions`) fall through as unauthenticated. +""" + +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from typing import Final + +from pydantic import BaseModel, ValidationError + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.auth_checks import ( + _is_model_cost_zero, # pyright: ignore[reportPrivateUsage] # the zero-cost predicate the auth-time budget checks use; no public equivalent +) +from litellm.router import Router + + +class _RequestMetadata(BaseModel): + user_api_key_auth: UserAPIKeyAuth | None = None + + +class _FallbackBudgetSettings(BaseModel): + enforce_fallback_budget: bool = True + + +def _token_in_metadata(metadata: object) -> UserAPIKeyAuth | None: + try: + return _RequestMetadata.model_validate(metadata).user_api_key_auth + except ValidationError: + return None + + +def _user_api_key_auth_from_request(request_kwargs: Mapping[str, object]) -> UserAPIKeyAuth | None: + return next( + ( + token + for field in ("metadata", "litellm_metadata") + if (token := _token_in_metadata(request_kwargs.get(field))) is not None + ), + None, + ) + + +def _enforced_by_general_settings() -> bool: + from litellm.proxy.proxy_server import general_settings + + return _FallbackBudgetSettings.model_validate(general_settings).enforce_fallback_budget + + +def _applies_user_budget_to_team_keys() -> bool: + from litellm.proxy.proxy_server import general_settings + + return general_settings.get("apply_user_budget_to_team_keys") is True + + +async def _counter_spend(counter_key: str, fallback_spend: float, max_budget: float) -> float: + """ + Read a spend counter the same way the auth-time budget checks do. + + `max_budget` is not advisory: it makes `get_current_spend` re-check the counter against the + authoritative recorded spend before admitting. A counter restored from an older Redis snapshot + reads as a hit rather than a clean miss, so without this the reseed path never runs and a + stale-low counter would keep admitting paid fallbacks past the cap. + """ + from litellm.proxy.proxy_server import get_current_spend + + return await get_current_spend( + counter_key=counter_key, + fallback_spend=fallback_spend, + max_budget=max_budget, + ) + + +async def is_token_within_budget_for_model(*, model: str, valid_token: UserAPIKeyAuth, llm_router: Router) -> bool: + """ + True when the key and the user behind it can still pay for `model`. + + A zero-cost fallback target is always allowed: refusing it would deny a request on spend some + other model accrued, which is the same reasoning behind the auth-time bypass. + """ + if _is_model_cost_zero(model=model, llm_router=llm_router): + return True + + key_budget: Final = valid_token.max_budget + if key_budget is not None and valid_token.token is not None: + key_spend: Final = await _counter_spend( + counter_key=f"spend:key:{valid_token.token}", + fallback_spend=valid_token.spend or 0.0, + max_budget=key_budget, + ) + if key_spend >= key_budget: + return False + + # Mirrors `_PROXY_MaxBudgetLimiter`: a team key does not carry the key owner's personal budget + # unless the proxy opts in, so the personal cap must not gate the fallback either. + user_budget: Final = valid_token.user_max_budget + if ( + user_budget is not None + and valid_token.user_id is not None + and (valid_token.team_id is None or _applies_user_budget_to_team_keys()) + ): + user_spend: Final = await _counter_spend( + counter_key=f"spend:user:{valid_token.user_id}", + fallback_spend=valid_token.user_spend or 0.0, + max_budget=user_budget, + ) + if user_spend >= user_budget: + return False + + return True + + +@dataclass(frozen=True, slots=True) +class RouterFallbackBudgetCheck: + """ + `FallbackBudgetCheck` for the proxy's router: while `is_enforced()` is true, a paid fallback + target is attempted only when the caller is still within budget. Requests that carry no key + (for example internal health checks) are not restricted. + """ + + is_enforced: Callable[[], bool] + + async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], llm_router: Router) -> bool: + if not self.is_enforced(): + return True + valid_token: Final = _user_api_key_auth_from_request(request_kwargs) + if valid_token is None: + return True + try: + return await is_token_within_budget_for_model(model=model, valid_token=valid_token, llm_router=llm_router) + except Exception as e: # noqa: BLE001 # fail closed: a spend lookup failure must not bill the caller + verbose_proxy_logger.warning("Skipping fallback to model=%s: budget lookup failed: %s", model, e) + return False + + +router_fallback_budget_check: Final = RouterFallbackBudgetCheck(is_enforced=_enforced_by_general_settings) diff --git a/litellm/proxy/auth/team_grants.py b/litellm/proxy/auth/team_grants.py index 0421659c331..2029ee342ae 100644 --- a/litellm/proxy/auth/team_grants.py +++ b/litellm/proxy/auth/team_grants.py @@ -59,6 +59,7 @@ class TeamGrants(TypedDict, total=False): team_tpd_limit: ReadOnly[int | None] team_max_budget: ReadOnly[float | None] team_soft_budget: ReadOnly[float | None] + team_model_max_budget: ReadOnly[dict[str, object] | None] team_spend: ReadOnly[float | None] team_models: ReadOnly[Sequence[str]] team_blocked: ReadOnly[bool] @@ -101,6 +102,7 @@ def team_grants( team_tpd_limit=team_object.tpd_limit, team_max_budget=team_object.max_budget, team_soft_budget=team_object.soft_budget, + team_model_max_budget=team_object.model_max_budget, team_spend=team_object.spend, team_models=tuple(team_object.models), team_blocked=team_object.blocked, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 6757c0c594d..4cbd4213463 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -304,6 +304,16 @@ class _UserModelBudgetLimiter(Protocol): ) -> bool: ... +class _TeamModelBudgetLimiter(Protocol): + async def is_team_within_model_budget( + self, + team_id: str, + team_model_max_budget: Mapping[str, object], + key_model_max_budget: Mapping[str, object] | None, + model: str, + ) -> bool: ... + + class _TokenTeamModels(Protocol): @property def team_models(self) -> list[str]: ... @@ -374,6 +384,25 @@ async def _check_user_model_budget( ) +async def _check_team_model_budget( + valid_token: UserAPIKeyAuth, + model_max_budget_limiter: _TeamModelBudgetLimiter, + models: list[str], +) -> None: + """Enforce the team's `model_max_budget` for every requested model the key does not override.""" + team_model_max_budget: Final = valid_token.team_model_max_budget + if valid_token.team_id is None or not team_model_max_budget: + return + key_model_max_budget: Final[Mapping[str, object] | None] = valid_token.model_max_budget + for model_name in models: + await model_max_budget_limiter.is_team_within_model_budget( + team_id=valid_token.team_id, + team_model_max_budget=team_model_max_budget, + key_model_max_budget=key_model_max_budget, + model=model_name, + ) + + async def _check_key_model_budget_with_fallback( valid_token: UserAPIKeyAuth, model_max_budget_limiter: _KeyModelBudgetLimiter, @@ -2376,6 +2405,7 @@ async def _user_api_key_auth_builder( team_id=valid_token.team_id, max_budget=valid_token.team_max_budget, soft_budget=valid_token.team_soft_budget, + model_max_budget=valid_token.team_model_max_budget, spend=valid_token.team_spend, tpm_limit=valid_token.team_tpm_limit, rpm_limit=valid_token.team_rpm_limit, @@ -2530,6 +2560,7 @@ def _team_obj_from_token(valid_token: UserAPIKeyAuth) -> LiteLLM_TeamTableCached team_id=valid_token.team_id, max_budget=valid_token.team_max_budget, soft_budget=valid_token.team_soft_budget, + model_max_budget=valid_token.team_model_max_budget, spend=valid_token.team_spend, tpm_limit=valid_token.team_tpm_limit, rpm_limit=valid_token.team_rpm_limit, @@ -2606,6 +2637,7 @@ async def _run_centralized_common_checks( litellm_proxy_admin_name, llm_router, master_key, + model_max_budget_limiter, prisma_client, proxy_logging_obj, user_api_key_cache, @@ -2874,6 +2906,21 @@ async def _run_centralized_common_checks( finally: release_spend_counter_batch() + if not skip_budget_checks: + await _check_team_model_budget( + valid_token=user_api_key_auth_obj, + model_max_budget_limiter=model_max_budget_limiter, + models=_get_model_names_for_budget_checks( + model=_get_model_from_request_context( + request_data=request_data, + route=route, + request=request, + llm_router=llm_router, + team_id=user_api_key_auth_obj.team_id, + ) + ), + ) + await _reserve_budget_after_common_checks( user_api_key_auth_obj=user_api_key_auth_obj, request=request, diff --git a/litellm/proxy/db/create_views.py b/litellm/proxy/db/create_views.py index d3f3de730ab..f7131091c0b 100644 --- a/litellm/proxy/db/create_views.py +++ b/litellm/proxy/db/create_views.py @@ -78,6 +78,7 @@ async def create_missing_views(db: SupportsRawQueries) -> None: v.*, t.spend AS team_spend, t.max_budget AS team_max_budget, + t.model_max_budget AS team_model_max_budget, t.tpm_limit AS team_tpm_limit, t.rpm_limit AS team_rpm_limit, t.tpd_limit AS team_tpd_limit, diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index a90d1351fd7..599ae90bcae 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -18,6 +18,8 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, overload from urllib.parse import quote, unquote +from typing_extensions import ReadOnly, TypedDict + import litellm from litellm._logging import verbose_proxy_logger from litellm.caching import RedisCache @@ -109,6 +111,10 @@ def _batch_cost_row_to_write(payload: SpendLogsPayload, disable_spend_logs: bool return MappingProxyType({field: value for field, value in payload.items() if field in _BATCH_COST_CLAIM_FIELDS}) +class _SpendIncrement(TypedDict): + increment: ReadOnly[float] + + class _SpendBatch(Protocol): litellm_usertable: BatchTable litellm_verificationtoken: BatchTable @@ -1615,10 +1621,12 @@ class DBSpendUpdateWriter: async with transaction.batch_() as batcher: # Sort by token for consistent lock ordering across pods to prevent deadlocks. for token, response_cost in sorted(key_list_transactions.items()): + spend_increment: _SpendIncrement = {"increment": response_cost} batcher.litellm_verificationtoken.update_many( # 'update_many' prevents error from being raised if no row exists where={"token": token}, data={ - "spend": {"increment": response_cost}, + "spend": spend_increment, + "total_spend": spend_increment, "last_active": datetime.now(timezone.utc), }, ) diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index afb9997f2e6..6874d7aa73e 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -37,6 +37,7 @@ from litellm.types.guardrails import ( ApplyGuardrailResponse, BaseLitellmParams, BedrockGuardrailConfigModel, + BedrockGuardrailStreamingParams, Guardrail, GuardrailEventHooks, GuardrailInfoResponse, @@ -1959,7 +1960,10 @@ async def get_provider_specific_params(): ``` """ # Get fields from the models - bedrock_fields: Final = _get_fields_from_model(BedrockGuardrailConfigModel) + bedrock_fields: Final = { + **_get_fields_from_model(BedrockGuardrailConfigModel), + **_get_fields_from_model(BedrockGuardrailStreamingParams), + } presidio_fields: Final = _get_fields_from_model(PresidioPresidioConfigModelUserInterface) lakera_v2_fields: Final = _get_fields_from_model(LakeraV2GuardrailConfigModel) tool_permission_fields: Final = _get_fields_from_model(ToolPermissionGuardrailConfigModel) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 2c407d91a48..434c52ca6f3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -248,6 +248,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): streaming_buffer_until_moderated: bool | None = None, streaming_sampling_rate: int | None = None, streaming_end_of_stream_only: bool | None = None, + streaming_buffer_release_on_scan: bool | None = None, **kwargs, ): self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) @@ -258,6 +259,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): "streaming_buffer_until_moderated": streaming_buffer_until_moderated, "streaming_sampling_rate": streaming_sampling_rate, "streaming_end_of_stream_only": streaming_end_of_stream_only, + "streaming_buffer_release_on_scan": streaming_buffer_release_on_scan, } ) ) @@ -321,13 +323,18 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): self.streaming_buffer_until_moderated = streaming_params.streaming_buffer_until_moderated self.streaming_sampling_rate = streaming_params.streaming_sampling_rate self.streaming_end_of_stream_only = streaming_params.streaming_end_of_stream_only + self.streaming_buffer_release_on_scan = streaming_params.streaming_buffer_release_on_scan def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None: super().update_in_memory_litellm_params(litellm_params) self._set_streaming_params(BedrockGuardrailStreamingParams.from_extras(litellm_params.model_extra)) def _streams_incrementally(self) -> bool: - return not self.streaming_buffer_until_moderated and not self.mask_response_content + if self.mask_response_content: + return False + if not self.streaming_buffer_until_moderated: + return True + return self.streaming_buffer_release_on_scan and not self.streaming_end_of_stream_only @classmethod def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: diff --git a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/__init__.py index c88e6e97a96..59f02817e5f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/__init__.py @@ -23,6 +23,8 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" event_hook=litellm_params.mode, default_on=litellm_params.default_on, fail_on_error=litellm_params.fail_on_error, + streaming_buffer_until_moderated=streaming_params.streaming_buffer_until_moderated, + streaming_buffer_release_on_scan=streaming_params.streaming_buffer_release_on_scan, streaming_end_of_stream_only=streaming_params.streaming_end_of_stream_only, streaming_sampling_rate=streaming_params.streaming_sampling_rate, ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py index 8fed1f906e5..924bbd2bc1a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py @@ -260,6 +260,8 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): api_key: str | None = None, api_base: str | None = None, fail_on_error: bool | None = True, + streaming_buffer_until_moderated: bool | None = None, + streaming_buffer_release_on_scan: bool | None = None, streaming_end_of_stream_only: bool | None = None, streaming_sampling_rate: int | None = None, async_handler: AsyncHTTPHandler | None = None, @@ -287,6 +289,8 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): CrowdStrikeAIDRGuardrailConfigModelOptionalParams( streaming_end_of_stream_only=streaming_end_of_stream_only, streaming_sampling_rate=streaming_sampling_rate, + streaming_buffer_until_moderated=streaming_buffer_until_moderated, + streaming_buffer_release_on_scan=streaming_buffer_release_on_scan, ) ) @@ -310,6 +314,8 @@ class CrowdStrikeAIDRHandler(CustomGuardrail): ) def _set_streaming_params(self, streaming_params: CrowdStrikeAIDRGuardrailConfigModelOptionalParams) -> None: + self.streaming_buffer_until_moderated: bool = streaming_params.streaming_buffer_until_moderated or False + self.streaming_buffer_release_on_scan: bool = streaming_params.streaming_buffer_release_on_scan or False self.streaming_end_of_stream_only: bool = streaming_params.streaming_end_of_stream_only or False self.streaming_sampling_rate: int = streaming_params.streaming_sampling_rate or 5 diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index ee5cd7c4cb8..7f51d733d4c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -956,6 +956,7 @@ class UnifiedLLMGuardrails(CustomLogger): buffer_until_moderated: bool = _streaming_flag( "streaming_buffer_until_moderated", buffer_until_moderated_default ) + release_on_scan: Final[bool] = _streaming_flag("streaming_buffer_release_on_scan", False) if ( buffer_until_moderated @@ -970,9 +971,7 @@ class UnifiedLLMGuardrails(CustomLogger): ) buffer_until_moderated = False - # Buffering can only moderate the assembled response, so it always - # defers to end-of-stream. - if buffer_until_moderated: + if buffer_until_moderated and not release_on_scan: end_of_stream_only = True if guardrail_to_apply is None: @@ -1026,12 +1025,14 @@ class UnifiedLLMGuardrails(CustomLogger): chunk_counter = 0 responses_so_far: Final[list[object]] = [] responses_yielded: Final[list[object]] = [] + withheld_items: Final[list[object]] = [] # mutable-ok: streaming window must be released incrementally pending_end_of_stream_items: Final[list[object]] = [] # Whether any real response chunk has been forwarded to the client. # Drives how a block terminates the stream: continue the in-progress # message (True) vs emit a standalone block message (False, buffered). chunks_yielded = False last_scan_key: StreamingScanKey | None = None # rebind-ok: replaced after every scan round + tool_calls_in_flight = False # rebind-ok: tracks the latest scan key's unscanned tool calls async for item in response: chunk_counter += 1 @@ -1069,21 +1070,37 @@ class UnifiedLLMGuardrails(CustomLogger): chunks_yielded = True responses_yielded.append(item) yield item + else: + withheld_items.append(item) continue # Process chunk based on sampling rate + if buffer_until_moderated: + withheld_items.append(item) if chunk_counter % sampling_rate == 0: endpoint_translation = mappings[CallTypes(call_type)]() scan_key = endpoint_translation.get_streaming_scan_key(responses_so_far) + if scan_key is not None: + tool_calls_in_flight = scan_key.tool_calls_in_flight + hold_window = buffer_until_moderated and (scan_key is None or tool_calls_in_flight) if _is_redundant_scan(scan_key, last_scan_key): verbose_proxy_logger.debug( "Skipping streaming chunk %s for guardrail %s: nothing new to scan since the last round", chunk_counter, guardrail_to_apply.guardrail_name, ) - chunks_yielded = True - responses_yielded.append(item) - yield item + if buffer_until_moderated: + if hold_window: + continue + for withheld_item in withheld_items: + chunks_yielded = True + responses_yielded.append(withheld_item) + yield withheld_item + withheld_items.clear() + else: + chunks_yielded = True + responses_yielded.append(item) + yield item continue verbose_proxy_logger.debug( @@ -1093,13 +1110,9 @@ class UnifiedLLMGuardrails(CustomLogger): guardrail_to_apply.guardrail_name, ) - # Deep-copy the current chunk before guardrail processing. - # process_output_streaming_response modifies responses_so_far - # in-place: it puts the combined guardrailed text in the first - # chunk and clears all subsequent chunks to "". Without this - # copy, yielding processed_items[-1] would yield an empty - # string, permanently losing this chunk's content. - original_item = copy.deepcopy(item) + original_items = ( + tuple(copy.deepcopy(withheld_items)) if buffer_until_moderated else (copy.deepcopy(item),) + ) try: await endpoint_translation.process_output_streaming_response( @@ -1144,13 +1157,24 @@ class UnifiedLLMGuardrails(CustomLogger): return if scan_key is not None: last_scan_key = scan_key - chunks_yielded = True - responses_yielded.append(original_item) - yield original_item + if hold_window: + verbose_proxy_logger.debug( + "Holding %s buffered chunks for guardrail %s: this round could not scan the whole window", + len(withheld_items), + guardrail_to_apply.guardrail_name, + ) + withheld_items[:] = original_items + continue + for original_item in original_items: + chunks_yielded = True + responses_yielded.append(original_item) + yield original_item + withheld_items.clear() else: - chunks_yielded = True - responses_yielded.append(item) - yield item + if not buffer_until_moderated: + chunks_yielded = True + responses_yielded.append(item) + yield item # Stream has ended - do final processing with all collected chunks if call_type is not None and CallTypes(call_type) in mappings: @@ -1162,14 +1186,13 @@ class UnifiedLLMGuardrails(CustomLogger): endpoint_translation = mappings[CallTypes(call_type)]() - # When buffering, snapshot the original chunks before moderation. - # A shallow copy suffices: end-of-stream - # process_output_streaming_response builds a separate assembled - # response (it does not mutate the individual chunks in place), and - # the chunks themselves are replayed verbatim -- so we only need to - # preserve the list, not clone every chunk (deepcopy would double - # peak memory for large responses). - buffered_items: Final = list(responses_so_far) if buffer_until_moderated else None + buffered_items: Final = ( + tuple(copy.deepcopy(withheld_items)) + if buffer_until_moderated and release_on_scan and not end_of_stream_only + else tuple(withheld_items) + if buffer_until_moderated + else None + ) end_scan_key: Final = endpoint_translation.get_streaming_scan_key(responses_so_far) if _is_redundant_scan(end_scan_key, last_scan_key): verbose_proxy_logger.debug( diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index b7ab215a2cd..356eb7c96c6 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -44,6 +44,7 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail): streaming_buffer_until_moderated=streaming_params.streaming_buffer_until_moderated, streaming_sampling_rate=streaming_params.streaming_sampling_rate, streaming_end_of_stream_only=streaming_params.streaming_end_of_stream_only, + streaming_buffer_release_on_scan=streaming_params.streaming_buffer_release_on_scan, ) litellm.logging_callback_manager.add_litellm_callback(_bedrock_callback) return _bedrock_callback diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index efaaab277a9..bbfc7325f40 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -19,12 +19,14 @@ from litellm.types.utils import BudgetConfig, StandardLoggingPayload VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX: Final = "virtual_key_spend" END_USER_SPEND_CACHE_KEY_PREFIX: Final = "end_user_model_spend" USER_SPEND_CACHE_KEY_PREFIX: Final = "user_model_spend" +TEAM_SPEND_CACHE_KEY_PREFIX: Final = "team_model_spend" _SPEND_CACHE_KEY_PREFIXES: Final = MappingProxyType( { Litellm_EntityType.KEY: VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX, Litellm_EntityType.USER: USER_SPEND_CACHE_KEY_PREFIX, Litellm_EntityType.END_USER: END_USER_SPEND_CACHE_KEY_PREFIX, + Litellm_EntityType.TEAM: TEAM_SPEND_CACHE_KEY_PREFIX, } ) @@ -37,6 +39,7 @@ _BUDGET_START_TIME_KEY_PREFIXES: Final = MappingProxyType( Litellm_EntityType.KEY: "virtual_key_budget_start_time", Litellm_EntityType.USER: "user_model_budget_start_time", Litellm_EntityType.END_USER: "end_user_budget_start_time", + Litellm_EntityType.TEAM: "team_model_budget_start_time", } ) @@ -139,6 +142,18 @@ def resolve_model_budget(model: str, model_max_budget: Mapping[str, object]) -> return None +def team_model_budget_applies(model: str, key_model_max_budget: Mapping[str, object] | None) -> bool: + """A key entry that spend-gates `model` overrides the team cap: it is then gated on and billed to the key alone.""" + if not key_model_max_budget: + return True + resolved: Final = resolve_model_budget(model=model, model_max_budget=key_model_max_budget) + return resolved is None or not _spend_gated(resolved.budget_config) + + +def _spend_gated(budget_config: BudgetConfig) -> bool: + return budget_config.max_budget is not None and budget_config.max_budget >= 0 + + def _budget_model_candidates(model: str) -> tuple[str, ...]: """Names a budget may be configured under for a request on `model`, most specific first. @@ -346,6 +361,30 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): exceeded_message=f"LiteLLM End User: {end_user_id}, exceeded budget for model={model}", ) + async def is_team_within_model_budget( + self, + team_id: str, + team_model_max_budget: Mapping[str, object], + key_model_max_budget: Mapping[str, object] | None, + model: str, + ) -> bool: + """ + Check if the team is within the model budget, unless the key's own + `model_max_budget` overrides it for `model` + + Raises: + BudgetExceededError: If the team has exceeded the model budget + """ + if not team_model_budget_applies(model=model, key_model_max_budget=key_model_max_budget): + return True + return await self._is_entity_within_model_budget( + entity_type=Litellm_EntityType.TEAM, + entity_id=team_id, + model_max_budget=team_model_max_budget, + model=model, + exceeded_message=f"LiteLLM Team: {team_id}, exceeded budget for model={model}", + ) + async def _is_entity_within_model_budget( self, entity_type: Litellm_EntityType, @@ -456,11 +495,26 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): return response_cost: Final[float] = standard_logging_payload.get("response_cost", 0) + key_model_max_budget: Final = _metadata.get("user_api_key_model_max_budget") entity_budgets: Final = ( ( Litellm_EntityType.KEY, payload_metadata.get("user_api_key_hash"), - _metadata.get("user_api_key_model_max_budget"), + key_model_max_budget, + ), + ( + Litellm_EntityType.TEAM, + payload_metadata.get("user_api_key_team_id"), + ( + _metadata.get("user_api_key_team_model_max_budget") + if team_model_budget_applies( + model=model, + key_model_max_budget=( + key_model_max_budget if isinstance(key_model_max_budget, Mapping) else None + ), + ) + else None + ), ), ( Litellm_EntityType.USER, @@ -478,7 +532,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): if not resolved_budgets: verbose_proxy_logger.debug( "Not running _PROXY_VirtualKeyModelMaxBudgetLimiter.async_log_success_event: " - "no key, user or end-user model_max_budget covers model=%s", + "no key, team, user or end-user model_max_budget covers model=%s", model, ) return diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index f72720b4726..89826694bb6 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -531,6 +531,7 @@ class RequestRateLimiterStash: owner_litellm_call_id: str | None = None rate_limit_response: RateLimitResponse | None = None parallel_slot: ParallelSlotAcquisition | None = None + parallel_slot_release_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False, compare=False) reserved_tokens: int = 0 reserved_model: str | None = None reserved_scopes: frozenset[tuple[str, str]] = field(default_factory=frozenset) @@ -1620,6 +1621,20 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): statuses.append(self._gauge_status(gauge, in_flight + 1, "OK")) return RateLimitResponse(overall_code="OK", statuses=statuses) + async def _release_stashed_parallel_slot( + self, + stash: RequestRateLimiterStash | None, + parent_otel_span: Span | None, + ) -> None: + if stash is None: + return + async with stash.parallel_slot_release_lock: + acquisition: Final = stash.parallel_slot + if acquisition is None: + return + await self._release_parallel_request_slots(acquisition, parent_otel_span) + stash.parallel_slot = None # rebind-ok: marks this request's slot as released + async def _release_parallel_request_slots( self, acquisition: ParallelSlotAcquisition, @@ -3379,13 +3394,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): parent_otel_span=user_api_key_dict.parent_otel_span, ) stash.reservation_released = True - acquisition: Final = stash.parallel_slot - if acquisition is not None: - await self._release_parallel_request_slots( - acquisition=acquisition, - parent_otel_span=user_api_key_dict.parent_otel_span, - ) - stash.parallel_slot = None + await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) self._handle_rate_limit_error( response=io_response, descriptors=descriptors, @@ -3700,13 +3709,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) if tpm_response["overall_code"] == "OVER_LIMIT": - acquisition: Final = stash.parallel_slot - if acquisition is not None: - await self._release_parallel_request_slots( - acquisition=acquisition, - parent_otel_span=user_api_key_dict.parent_otel_span, - ) - stash.parallel_slot = None + await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) self._handle_rate_limit_error( response=tpm_response, descriptors=descriptors, @@ -4524,13 +4527,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): verbose_proxy_logger.debug("INSIDE parallel request limiter ASYNC SUCCESS LOGGING") stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(kwargs)) - acquisition: Final = stash.parallel_slot if stash is not None else None - if stash is not None and acquisition is not None: - await self._release_parallel_request_slots( - acquisition=acquisition, - parent_otel_span=litellm_parent_otel_span, - ) - stash.parallel_slot = None + await self._release_stashed_parallel_slot(stash, litellm_parent_otel_span) pipeline_operations: Final = self._build_success_event_pipeline_operations( kwargs=kwargs, @@ -4650,13 +4647,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): pipeline_operations: Final[list[RedisPipelineIncrementOperation]] = [] stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(kwargs)) - acquisition: Final = stash.parallel_slot if stash is not None else None - if stash is not None and acquisition is not None: - await self._release_parallel_request_slots( - acquisition=acquisition, - parent_otel_span=litellm_parent_otel_span, - ) - stash.parallel_slot = None + await self._release_stashed_parallel_slot(stash, litellm_parent_otel_span) # Skip the reservation refund if async_post_call_failure_hook # already released it (proxy-level rejection that also bubbles up @@ -4764,23 +4755,23 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): object's current max_parallel_requests configuration, which can change mid-request) decides whether there is anything to release. """ - stash: Final = get_request_stash() - if stash is None or stash.parallel_slot is None: - return - - await self._release_parallel_request_slots( - acquisition=stash.parallel_slot, - parent_otel_span=None, - ) - stash.parallel_slot = None + await self._release_stashed_parallel_slot(get_request_stash(), None) async def async_post_call_success_hook(self, data: dict, user_api_key_dict: UserAPIKeyAuth, response): """ - Post-call hook to update rate limit headers in the response. + Release completed-request slots and update rate limit headers in the response. """ try: - stash: Final = get_request_stash() - litellm_proxy_rate_limit_response: Final = stash.rate_limit_response if stash is not None else None + slot_stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(data)) + await self._release_stashed_parallel_slot(slot_stash, user_api_key_dict.parent_otel_span) + except Exception as e: + verbose_proxy_logger.exception("Error releasing parallel request slot in post-call hook: %s", e) + + try: + header_stash: Final = get_request_stash() + litellm_proxy_rate_limit_response: Final = ( + header_stash.rate_limit_response if header_stash is not None else None + ) if litellm_proxy_rate_limit_response is not None and response_has_hidden_params(response): additional_headers: Final = ensure_response_additional_headers(response) @@ -4848,12 +4839,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): stash: Final = get_request_stash() if stash is None: return - if stash.parallel_slot is not None: - await self._release_parallel_request_slots( - acquisition=stash.parallel_slot, - parent_otel_span=user_api_key_dict.parent_otel_span, - ) - stash.parallel_slot = None + await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) if stash.batch_enqueued_reservation is not None: await self.batch_enqueued_token_store.refund( diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 18f4280c01c..b03f1e4348c 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -25,6 +25,7 @@ from litellm.constants import ( LITELLM_PROXY_MASTER_KEY_ALIAS, OTEL_SERVICE_NAME_METADATA_KEYS, PRE_CALL_EXECUTED_GUARDRAILS_KEY, + ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, ROUTING_REQUEST_TAGS_METADATA_KEY, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY, @@ -369,7 +370,13 @@ _CLIENT_PRICING_METADATA_FIELDS: Final = frozenset({"model_info", "standard_logg # and read by spend logs as fact; a client value has no legitimate meaning and no # key or team setting keeps it, so the strip is never gated. _ROUTER_RESERVED_METADATA_FIELDS: Final = frozenset( - {"attempted_fallbacks", "original_model_group", "request_retry_count", CLIENT_OUTPUT_CEILING_METADATA_KEY} + { + "attempted_fallbacks", + "original_model_group", + "request_retry_count", + CLIENT_OUTPUT_CEILING_METADATA_KEY, + ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, + } ) _ALLOW_CLIENT_PRICING_OVERRIDE_METADATA_KEY: Final = "allow_client_pricing_override" @@ -2327,6 +2334,7 @@ async def add_litellm_data_to_request( # Team spend, budget - used by prometheus.py data[_metadata_variable_name]["user_api_key_team_max_budget"] = user_api_key_dict.team_max_budget data[_metadata_variable_name]["user_api_key_team_spend"] = user_api_key_dict.team_spend + data[_metadata_variable_name]["user_api_key_team_model_max_budget"] = user_api_key_dict.team_model_max_budget data[_metadata_variable_name]["user_api_key_request_route"] = user_api_key_dict.request_route # API Key spend, budget - used by prometheus.py diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index 98155ad6839..973311608ed 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -55,6 +55,7 @@ def validate_budget_duration(budget_duration: str | None, status_code: int = 400 from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache from litellm.proxy._types import ( + CommonProxyErrors, KeyRequestBase, LiteLLM_ManagementEndpoint_MetadataFields, LiteLLM_ManagementEndpoint_MetadataFields_Premium, @@ -73,12 +74,62 @@ from litellm.proxy._types import ( # noqa: F401 re-exported from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.utils import _premium_user_check from litellm.repositories.team_repository import TeamRepository +from litellm.types.utils import BudgetConfig if TYPE_CHECKING: from litellm.proxy._types import NewProjectRequest, UpdateProjectRequest from litellm.proxy.utils import PrismaClient, ProxyLogging +def validate_team_model_max_budget( + model_max_budget: Mapping[str, BudgetConfig] | None, + premium_user: bool, +) -> None: + """Reject a team `model_max_budget` the limiter could not enforce (no duration, bad cap, tpm/rpm limits).""" + if not model_max_budget: + return + if premium_user is not True: + raise HTTPException( + status_code=403, + detail={ + "error": f"Setting model_max_budget on a team is an enterprise feature. {CommonProxyErrors.not_premium_user.value}" + }, + ) + for model_name, budget_config in model_max_budget.items(): + if not model_name.strip(): + raise HTTPException( + status_code=400, + detail={"error": "model_max_budget keys must be non-empty model names"}, + ) + max_budget = budget_config.max_budget + if max_budget is None or not math.isfinite(max_budget) or max_budget < 0: + raise HTTPException( + status_code=400, + detail={ + "error": ( + f"model_max_budget[{model_name!r}].max_budget must be a non-negative finite number. " + f"Received: {max_budget}" + ) + }, + ) + if budget_config.budget_duration is None: + raise HTTPException( + status_code=400, + detail={"error": f"model_max_budget[{model_name!r}] requires a budget_duration, e.g. '1d' or '30d'"}, + ) + validate_budget_duration(budget_config.budget_duration) + if budget_config.tpm_limit is not None or budget_config.rpm_limit is not None: + raise HTTPException( + status_code=400, + detail={ + "error": ( + f"model_max_budget[{model_name!r}] tpm_limit/rpm_limit are not enforced on a team; " + "set per-model rate limits on the key instead" + ) + }, + ) + + def require_caller_user_id_for_non_admin( user_api_key_dict: UserAPIKeyAuth, ) -> str: diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index ee8ae66ea11..802a7c3e469 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -4166,7 +4166,10 @@ async def info_key_fn( Returns: - key: str - The key that was looked up, echoed back as it was passed in - - info: dict - The key's row, minus the hashed token + - info: dict - The key's row, minus the hashed token. Deleted keys are served from the + LiteLLM_DeletedVerificationToken archive and carry deleted_at / deleted_by + - status: "active" | "expired" | "revoked" | "deleted" - Derived from blocked, expires and + whether the row came from the archive - key_alias: str | None - User-friendly key alias - spend: float - Amount spent by the key. When budget_duration is set this covers only the current budget window, not the key's lifetime @@ -4220,10 +4223,15 @@ async def info_key_fn( hashed_key: str | None = key if key is not None: hashed_key = _hash_token_if_needed(token=key) - key_info = await _prisma_table(VerificationTokenRepository(prisma_client)).find_unique( + live_key_info: Final = await _prisma_table(VerificationTokenRepository(prisma_client)).find_unique( where={"token": hashed_key}, include={"litellm_budget_table": True}, ) + key_info: Final = ( + live_key_info + if live_key_info is not None + else await _find_deleted_key_info(prisma_client=prisma_client, hashed_key=hashed_key) + ) if key_info is None: raise ProxyException( message="Key not found in database", @@ -4231,7 +4239,6 @@ async def info_key_fn( param="key", code=status.HTTP_404_NOT_FOUND, ) - if ( await _can_user_query_key_info( user_api_key_dict=user_api_key_dict, @@ -4245,38 +4252,46 @@ async def info_key_fn( detail=f"You are not allowed to access this key's info. Your role={user_api_key_dict.user_role}", ) ## REMOVE HASHED TOKEN INFO BEFORE RETURNING ## - try: - key_info = key_info.model_dump() - except Exception: - # if using pydantic v1 - key_info = key_info.dict() # pyright: ignore[reportDeprecated] # deliberate pydantic v1 fallback - key_token_hash: Final[str | None] = key_info.pop("token") + key_info_dict: Final = key_info.model_dump() + key_token_hash: Final[str | None] = key_info_dict.pop("token") + key_info_dict["status"] = ( + "deleted" if live_key_info is None else _derive_key_status(key_info_dict, now=datetime.now(timezone.utc)) + ) - model_max_budget = key_info.get("model_max_budget") or {} - budget_table: Final = key_info.get("litellm_budget_table") or {} + model_max_budget = key_info_dict.get("model_max_budget") or {} + budget_table: Final = key_info_dict.get("litellm_budget_table") or {} if not model_max_budget and isinstance(budget_table, dict): model_max_budget = budget_table.get("model_max_budget") or {} if model_max_budget and key_token_hash: - key_info["model_max_budget_usage"] = await _build_model_max_budget_usage( + key_info_dict["model_max_budget_usage"] = await _build_model_max_budget_usage( api_key_hash=key_token_hash, model_max_budget=model_max_budget, user_api_key_cache=model_max_budget_limiter.dual_cache, ) budget_limits_usage: Final = await _build_budget_limits_usage( - budget_limits=key_info.get("budget_limits"), + budget_limits=key_info_dict.get("budget_limits"), api_key_hash=key_token_hash, ) if budget_limits_usage is not None: - key_info["budget_limits_usage"] = budget_limits_usage + key_info_dict["budget_limits_usage"] = budget_limits_usage - # Attach object_permission if object_permission_id is set - key_info = await attach_object_permission_to_dict(key_info, prisma_client) - - return {"key": key, "info": key_info} + return {"key": key, "info": await attach_object_permission_to_dict(key_info_dict, prisma_client)} except Exception as e: raise handle_exception_on_proxy(e) +async def _find_deleted_key_info( + prisma_client: PrismaClient, hashed_key: str | None +) -> LiteLLM_DeletedVerificationToken | None: + archived_row: Final = await _deleted_verification_token_table(prisma_client).find_first( + where={"token": hashed_key}, + order={"deleted_at": "desc"}, + ) + if archived_row is None: + return None + return LiteLLM_DeletedVerificationToken.model_validate(archived_row.model_dump()) + + def _check_model_access_group(models: list[str] | None, llm_router: Router | None, premium_user: bool) -> Literal[True]: """ if is_model_access_group is True + is_wildcard_route is True, check if user is a premium user @@ -6216,6 +6231,24 @@ async def get_member_team_ids( VALID_EXPIRES_FILTER_VALUES: Final = frozenset({"active", "expired"}) +KeyStatus = Literal["active", "expired", "revoked", "deleted"] +VALID_STATUS_FILTER_VALUES: Final[frozenset[KeyStatus]] = frozenset({"active", "expired", "revoked", "deleted"}) + + +class _KeyStatusSource(BaseModel): + blocked: bool | None = None + expires: datetime | None = None + + +def _derive_key_status(row: Mapping[str, object], now: datetime) -> KeyStatus: + source: Final = _KeyStatusSource.model_validate(row) + if source.blocked is True: + return "revoked" + if source.expires is None: + return "active" + expires_utc: Final = source.expires if source.expires.tzinfo else source.expires.replace(tzinfo=timezone.utc) + return "expired" if expires_utc < now else "active" + @router.get( "/key/list", @@ -6252,7 +6285,10 @@ async def list_keys( ), sort_order: str = Query(default="desc", description="Sort order ('asc' or 'desc')"), expand: list[str] | None = Query(None, description="Expand related objects (e.g. 'user')"), - status: str | None = Query(None, description="Filter by status (e.g. 'deleted')"), + status: str | None = Query( + None, + description="Filter by status: 'active' (not blocked, not expired), 'expired' (not blocked, past expiry), 'revoked' (blocked) or 'deleted' (archived keys). Omit to return live keys regardless of status.", + ), project_id: str | None = Query(None, description="Filter keys by project ID"), access_group_id: str | None = Query(None, description="Filter keys by access group ID"), agent_id: str | None = Query(None, description="Filter keys by agent ID"), @@ -6270,7 +6306,9 @@ async def list_keys( Parameters: expand: Optional[List[str]] - Expand related objects (e.g. 'user' to include user information) - status: Optional[str] - Filter by status. Currently supports "deleted" to query deleted keys. + status: Optional[str] - Filter by status: "active", "expired", "revoked" (blocked) or "deleted". + "deleted" reads the LiteLLM_DeletedVerificationToken archive; the other values partition the + live key table, so every live key matches exactly one of them. Returns: { @@ -6292,11 +6330,10 @@ async def list_keys( verbose_proxy_logger.error("Database not connected") raise Exception("Database not connected") - # Validate status parameter - if status is not None and status != "deleted": + if status is not None and status not in VALID_STATUS_FILTER_VALUES: raise HTTPException( status_code=400, - detail={"error": "Invalid status value. Currently only 'deleted' is supported."}, + detail={"error": "Invalid status value. Supported: 'active', 'expired', 'revoked', 'deleted'."}, ) if isinstance(expires, str) and expires not in VALID_EXPIRES_FILTER_VALUES: @@ -6608,6 +6645,18 @@ def _build_expires_where_clause(expires_filter: str, now: datetime) -> dict[str, return {"OR": [{"expires": None}, {"expires": {"gte": now}}]} +def _not_blocked_where_clause() -> dict[str, object]: + return {"OR": [{"blocked": None}, {"blocked": False}]} + + +def _build_status_where_clause(status_filter: str | None, now: datetime) -> dict[str, object] | None: + if status_filter == "revoked": + return {"blocked": True} + if status_filter in ("expired", "active"): + return {"AND": [_not_blocked_where_clause(), _build_expires_where_clause(status_filter, now)]} + return None + + def _build_key_search_where(search: str) -> KeySearchWhere: search_where: Final[KeySearchWhere] = { "OR": ( @@ -6635,6 +6684,7 @@ def _build_key_filter_conditions( use_key_alias_substring_matching: bool = False, expires_filter: str | None = None, search: str | None = None, + status_filter: str | None = None, ) -> Mapping[str, object]: """Build filter conditions for key listing. @@ -6724,6 +6774,8 @@ def _build_key_filter_conditions( # Apply team_id, project_id and access_group_id as global AND filters so they # narrow results across all visibility conditions (own keys, team keys, etc.) + now: Final = datetime.now(timezone.utc) + status_where: Final = _build_status_where_clause(status_filter, now) global_filters: Final[tuple[Mapping[str, object], ...]] = ( *( ( @@ -6741,10 +6793,11 @@ def _build_key_filter_conditions( *(({"access_group_ids": {"hasSome": [access_group_id]}},) if access_group_id else ()), *(({"agent_id": agent_id},) if agent_id and isinstance(agent_id, str) else ()), *( - (_build_expires_where_clause(expires_filter, datetime.now(timezone.utc)),) + (_build_expires_where_clause(expires_filter, now),) if expires_filter is not None and expires_filter in VALID_EXPIRES_FILTER_VALUES else () ), + *((status_where,) if status_where is not None else ()), ) combined_where: Final[Mapping[str, object]] = {"AND": [where, *global_filters]} if global_filters else where verbose_proxy_logger.debug("Filter conditions: %s", combined_where) @@ -6817,6 +6870,7 @@ async def _list_key_helper( use_key_alias_substring_matching=use_key_alias_substring_matching, expires_filter=expires_filter, search=search, + status_filter=status, ) # Calculate skip for pagination diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index e719d6d761a..23b02c22f24 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -22,7 +22,7 @@ from typing import TYPE_CHECKING, Annotated, Final, NamedTuple, NoReturn, Protoc import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status -from pydantic import BaseModel, JsonValue +from pydantic import BaseModel, JsonValue, ValidationError from typing_extensions import ReadOnly, TypedDict import litellm @@ -38,6 +38,7 @@ from litellm.proxy._types import ( DeleteTeamRequest, LiteLLM_AuditLogs, LiteLLM_DeletedTeamTable, + Litellm_EntityType, LiteLLM_ManagementEndpoint_MetadataFields, LiteLLM_ManagementEndpoint_MetadataFields_Premium, LiteLLM_ModelTable, @@ -95,6 +96,10 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars from litellm.proxy.common_utils.json_merge_patch import apply_json_merge_patch from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.hooks.model_max_budget_limiter import ( + build_model_max_budget_usage, + resolve_model_budget, +) from litellm.proxy.management_endpoints.common_daily_activity import ( get_daily_activity_aggregated, ) @@ -108,6 +113,7 @@ from litellm.proxy.management_endpoints.common_utils import ( _upsert_budget_and_membership, _user_has_admin_view, validate_budget_duration, + validate_team_model_max_budget, ) from litellm.proxy.management_endpoints.organization_endpoints import ( add_member_to_organization, @@ -177,6 +183,7 @@ from litellm.types.proxy.management_endpoints.team_endpoints import ( TeamUserSpendRow, UpdateTeamMemberPermissionsRequest, ) +from litellm.types.utils import BudgetConfig if TYPE_CHECKING: from prisma import Prisma @@ -1170,6 +1177,62 @@ def _check_team_budget_update_authority( ) +def _existing_model_cap(raw_budget_config: object) -> BudgetConfig | None: + try: + return BudgetConfig.model_validate(raw_budget_config) + except ValidationError: + return None + + +def _check_team_model_budget_update_authority( + data: UpdateTeamRequest, + user_api_key_dict: UserAPIKeyAuth, + existing_model_max_budget: Mapping[str, object] | None, +) -> None: + """Like `_check_team_budget_update_authority`: only a proxy admin may raise, re-window or drop a per-model cap.""" + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: + return + if "model_max_budget" not in data.model_fields_set or not existing_model_max_budget: + return + requested: Final[Mapping[str, BudgetConfig]] = data.model_max_budget or {} + for model_name, raw_existing in existing_model_max_budget.items(): + existing = _existing_model_cap(raw_existing) + if existing is None or existing.max_budget is None or model_name in requested: + continue + raise HTTPException( + status_code=403, + detail={ + "error": ( + f"Only a proxy admin can remove a team's model_max_budget for {model_name!r}. " + f"Current max_budget={existing.max_budget}." + ) + }, + ) + for model_name, proposed in requested.items(): + governing = resolve_model_budget(model=model_name, model_max_budget=existing_model_max_budget) + if governing is None: + continue + cap = governing.budget_config + if cap.max_budget is None: + continue + if ( + proposed.max_budget is None + or proposed.max_budget > cap.max_budget + or proposed.budget_duration != cap.budget_duration + ): + raise HTTPException( + status_code=403, + detail={ + "error": ( + f"Only a proxy admin can raise a team's model_max_budget for {model_name!r} or change its " + f"budget_duration. Current max_budget={cap.max_budget} per {cap.budget_duration} " + f"(entry {governing.budget_model!r}), requested={proposed.max_budget} per " + f"{proposed.budget_duration}." + ) + }, + ) + + def _should_auto_add_team_creator( user_api_key_dict: UserAPIKeyAuth, general_settings: Mapping[str, object], @@ -1230,6 +1293,7 @@ async def new_team( - prompts: Optional[List[str]] - List of prompts that the team is allowed to use. - organization_id: Optional[str] - The organization id of the team. Default is None. Create via `/organization/new`. - model_aliases: Optional[dict] - Model aliases for the team. [Docs](https://docs.litellm.ai/docs/proxy/team_based_routing#create-team-with-model-alias) + - model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}} - guardrails: Optional[List[str]] - Guardrails for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails) - policies: Optional[List[str]] - Policies for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails/guardrail_policies) - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key. @@ -1291,6 +1355,7 @@ async def new_team( general_settings, litellm_proxy_admin_name, llm_router, + premium_user, prisma_client, user_api_key_cache, ) @@ -1321,6 +1386,7 @@ async def new_team( validate_budget_duration(data.budget_duration) validate_budget_duration(data.team_member_budget_duration) + validate_team_model_max_budget(model_max_budget=data.model_max_budget, premium_user=premium_user) if data.soft_budget is not None: if data.max_budget is not None: @@ -1980,6 +2046,7 @@ async def update_team( - tags: Optional[List[str]] - Tags for [tracking spend](https://litellm.vercel.app/docs/proxy/enterprise#tracking-spend-for-custom-tags) and/or doing [tag-based routing](https://litellm.vercel.app/docs/proxy/tag_routing). - organization_id: Optional[str] - The organization id of the team. Default is None. Create via `/organization/new`. - model_aliases: Optional[dict] - Model aliases for the team. [Docs](https://docs.litellm.ai/docs/proxy/team_based_routing#create-team-with-model-alias) + - model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}} - guardrails: Optional[List[str]] - Guardrails for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails) - policies: Optional[List[str]] - Policies for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails/guardrail_policies) - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key. @@ -2031,6 +2098,7 @@ async def update_team( from litellm.proxy.proxy_server import ( litellm_proxy_admin_name, llm_router, + premium_user, prisma_client, proxy_logging_obj, user_api_key_cache, @@ -2069,6 +2137,7 @@ async def update_team( validate_budget_duration(data.budget_duration) validate_budget_duration(data.team_member_budget_duration) + validate_team_model_max_budget(model_max_budget=data.model_max_budget, premium_user=premium_user) existing_team_row = await _raw_team_db(TeamRepository(prisma_client)).find_unique( where={"team_id": data.team_id} @@ -2204,8 +2273,15 @@ async def update_team( user_api_key_dict=user_api_key_dict, existing_team_max_budget=existing_team_row.max_budget, ) + _check_team_model_budget_update_authority( + data=data, + user_api_key_dict=user_api_key_dict, + existing_model_max_budget=existing_team_row.model_max_budget, + ) updated_kv = data.json(exclude_unset=True) + if "model_max_budget" in updated_kv and updated_kv["model_max_budget"] is None: + updated_kv["model_max_budget"] = {} # Drop server-owned metadata keys from caller input so they can only # be written by the same code path that creates the underlying rows. @@ -4473,7 +4549,7 @@ async def team_info( ``` """ from litellm.proxy._types import TeamInfoResponseObjectTeamTable - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import model_max_budget_limiter, prisma_client try: if prisma_client is None: @@ -4573,6 +4649,12 @@ async def team_info( update={ # mutable-ok: pydantic update payload "members_with_roles": hydrated_members, "organization_models": organization_models, + "model_max_budget_usage": await build_model_max_budget_usage( + entity_type=Litellm_EntityType.TEAM, + entity_id=team_id, + model_max_budget=resolved_team_info.model_max_budget, + cache=model_max_budget_limiter.dual_cache, + ), } ) diff --git a/litellm/proxy/ocr_endpoints/endpoints.py b/litellm/proxy/ocr_endpoints/endpoints.py index 53ebbe91b54..dde3d5ceb50 100644 --- a/litellm/proxy/ocr_endpoints/endpoints.py +++ b/litellm/proxy/ocr_endpoints/endpoints.py @@ -15,12 +15,13 @@ from litellm.llms.base_llm.ocr.transformation import ( OCRResponse, parse_ocr_request_format, ) -from litellm.ocr.input import convert_upload_to_url_document, get_max_file_bytes +from litellm.ocr.legacy import convert_file_document_to_url_document, get_mime_type from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing router: Final = APIRouter() +_MAX_FILE_BYTES: Final = 50 * 1024 * 1024 def _build_document_from_upload( @@ -28,7 +29,15 @@ def _build_document_from_upload( filename: str | None, content_type: str | None, ) -> dict[str, str]: - return convert_upload_to_url_document(file_content, filename, content_type) + supplied_mime: Final = content_type.split(";")[0].strip() if content_type else None + mime_type: Final = ( + get_mime_type(filename) + if filename and (not supplied_mime or supplied_mime == "application/octet-stream") + else supplied_mime + ) + return convert_file_document_to_url_document( + {"type": "file", "file": file_content, "mime_type": mime_type or "application/octet-stream"} + ) def _with_request_format(data: Mapping[str, Any], request: Request) -> Mapping[str, Any]: @@ -103,9 +112,11 @@ async def _parse_multipart_form(request: Request) -> dict[str, Any]: # Seek to start in case the file was already partially read by middleware await uploaded_file.seek(0) - file_content: Final = await uploaded_file.read(get_max_file_bytes() + 1) + file_content: Final = await uploaded_file.read(_MAX_FILE_BYTES + 1) if not file_content: raise ValueError("Uploaded file is empty") + if len(file_content) > _MAX_FILE_BYTES: + raise ValueError("OCR file exceeds the size limit") document: Final = _build_document_from_upload( file_content=file_content, diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 686544d352c..de6f9cb7647 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -609,6 +609,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): # merely shares the name. if not request_dispatched_to_pass_through_endpoint(request): _metadata["user_api_key_model_max_budget"] = user_api_key_dict.model_max_budget + _metadata["user_api_key_team_model_max_budget"] = user_api_key_dict.team_model_max_budget _metadata["user_api_key_user_model_max_budget"] = user_api_key_dict.user_model_max_budget _metadata["user_api_key_end_user_model_max_budget"] = user_api_key_dict.end_user_model_max_budget _metadata.update( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 23bb8b6225b..1f31597c010 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -323,6 +323,7 @@ from litellm.proxy.auth.auth_utils import ( log_once_if_budget_reservation_disabled, warn_once_if_custom_auth_skips_common_checks, ) +from litellm.proxy.auth.fallback_budget import router_fallback_budget_check from litellm.proxy.auth.fallback_model_access import router_fallback_access_check from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.litellm_license import AUTO_ROUTER_LICENSE_REMEDY, LicenseCheck @@ -6161,6 +6162,7 @@ class ProxyConfig: ), ignore_invalid_deployments=True, # don't raise an error if a deployment is invalid fallback_access_check=router_fallback_access_check, + fallback_budget_check=router_fallback_budget_check, auto_router_capability_limit=_license_check.auto_router_capability_limit, ) @@ -6622,6 +6624,7 @@ class ProxyConfig: search_tools=search_tools, ignore_invalid_deployments=True, fallback_access_check=router_fallback_access_check, + fallback_budget_check=router_fallback_budget_check, auto_router_capability_limit=_license_check.auto_router_capability_limit, ) verbose_proxy_logger.debug("updated llm_router: %s", llm_router) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index d2375903c47..139fb031671 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -426,6 +426,7 @@ model LiteLLM_VerificationToken { key_alias String? soft_budget_cooldown Boolean @default(false) // key-level state on if budget alerts need to be cooled down spend Float @default(0.0) + total_spend Float @default(0.0) expires DateTime? models String[] aliases Json @default("{}") @@ -528,6 +529,7 @@ model LiteLLM_DeletedVerificationToken { key_alias String? soft_budget_cooldown Boolean @default(false) spend Float @default(0.0) + total_spend Float @default(0.0) expires DateTime? models String[] aliases Json @default("{}") diff --git a/litellm/proxy/spend_tracking/carried_budget_state.py b/litellm/proxy/spend_tracking/carried_budget_state.py index efd3a78d211..da8bf60ebda 100644 --- a/litellm/proxy/spend_tracking/carried_budget_state.py +++ b/litellm/proxy/spend_tracking/carried_budget_state.py @@ -25,6 +25,7 @@ def carry_team_and_user_budget_state( budget_reset_at=team_object.budget_reset_at, max_budget=team_object.max_budget, ) + valid_token.team_model_max_budget = team_object.model_max_budget # rebind-ok: caller keeps this object if user_object is not None: valid_token.user_budget_snapshot = UserBudgetSnapshot( # rebind-ok: same object the caller keeps using budget_reset_at=user_object.budget_reset_at, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 215fb143f7b..99cd78e0b58 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -4340,6 +4340,7 @@ class PrismaClient: v.*, t.spend AS team_spend, t.max_budget AS team_max_budget, + t.model_max_budget AS team_model_max_budget, t.tpm_limit AS team_tpm_limit, t.rpm_limit AS team_rpm_limit, t.tpd_limit AS team_tpd_limit @@ -4779,6 +4780,7 @@ class PrismaClient: t.spend AS team_spend, t.max_budget AS team_max_budget, t.soft_budget AS team_soft_budget, + t.model_max_budget AS team_model_max_budget, t.tpm_limit AS team_tpm_limit, t.rpm_limit AS team_rpm_limit, t.tpd_limit AS team_tpd_limit, diff --git a/litellm/router.py b/litellm/router.py index 789fc81d8d3..5f5522e9fd4 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -63,6 +63,7 @@ from litellm.constants import ( DEFAULT_MAX_LRU_CACHE_SIZE, INTERNAL_CALL_ORIGIN_METADATA_KEY, OUTPUT_TOKEN_CEILING_PARAMS, + ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, ROUTING_REQUEST_TAGS_METADATA_KEY, RUNTIME_UPDATABLE_ROUTER_SETTINGS, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, @@ -132,7 +133,7 @@ from litellm.router_utils.add_retry_fallback_headers import ( get_hidden_params_dict, prepare_response_for_header_attachment, replace_complexity_router_headers, - response_in_flight_token_count, + response_total_token_count, ) from litellm.router_utils.auto_router_model_naming import ( AUTO_ROUTER_MODEL_PREFIX, @@ -215,6 +216,8 @@ from litellm.router_utils.reasoning_effort_capability import ( resolve_supported_reasoning_efforts, ) from litellm.router_utils.router_callbacks.track_deployment_metrics import ( + find_deployment_metadata, + get_counted_usage_tokens, increment_deployment_failures_for_current_minute, increment_deployment_successes_for_current_minute, ) @@ -240,6 +243,7 @@ from litellm.types.router import ( DeploymentModelListingInfo, DeploymentTypedDict, FallbackAccessCheck, + FallbackBudgetCheck, GuardrailTypedDict, LiteLLM_Params, MockRouterTestingParams, @@ -777,6 +781,7 @@ class Router: background_health_check_model_groups: Sequence[str] | None = None, enable_weighted_failover: bool = False, fallback_access_check: FallbackAccessCheck | None = None, + fallback_budget_check: FallbackBudgetCheck | None = None, auto_router_capability_limit: AutoRouterCapabilityLimit | None = None, ) -> None: """ @@ -815,6 +820,7 @@ class Router: ignore_invalid_deployments (bool): Ignores invalid deployments, and continues with other deployments. Default is to raise an error. enable_weighted_failover (bool): When True and the routing strategy is "simple-shuffle", a retryable failure on one deployment causes the request to re-pick (weighted) across the other deployments in the same model group before any cross-group fallback runs. Bounded by `max_fallbacks`. Async-only: currently honored by `router.acompletion()` and other async entrypoints. The sync `router.completion()` path falls back to the regular fallback flow. Defaults to False. fallback_access_check (Optional[FallbackAccessCheck]): Awaited before each cross-model-group fallback attempt on the async path; a fallback target it rejects is skipped. Defaults to None (every configured fallback is attempted). + fallback_budget_check (Optional[FallbackBudgetCheck]): Awaited before each cross-model-group fallback attempt on the async path; a fallback target it rejects as over budget is skipped. Defaults to None (budget is not re-checked on fallback). Returns: Router: An instance of the litellm.Router class. @@ -856,6 +862,7 @@ class Router: self.ignore_invalid_deployments = ignore_invalid_deployments self.auto_router_capability_limit = auto_router_capability_limit self.fallback_access_check: Final = fallback_access_check + self.fallback_budget_check: Final = fallback_budget_check self.debug_level = debug_level self.enable_pre_call_checks = enable_pre_call_checks self.enable_tag_filtering = enable_tag_filtering @@ -7937,6 +7944,7 @@ class Router: response = original_function(*args, **kwargs) if coroutine_checker.is_async_callable(response) or inspect.isawaitable(response): response = await response + await self.increment_deployment_usage_for_response(response=response, request_kwargs=kwargs) ## PROCESS RESPONSE HEADERS response = await self.set_response_headers(response=response, model_group=model_group, request_kwargs=kwargs) @@ -8153,8 +8161,6 @@ class Router: """ Track remaining tpm/rpm quota for model in model_list """ - from litellm.types.caching import RedisPipelineIncrementOperation - try: # WS session wrappers fire with result=None; per-turn costs tracked by inner calls. if kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime"): @@ -8162,114 +8168,135 @@ class Router: standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) if standard_logging_object is None: raise ValueError("standard_logging_object is None") - if kwargs["litellm_params"].get("metadata") is None: - pass - else: - deployment_name: Final = kwargs["litellm_params"]["metadata"].get( - "deployment", None - ) # stable name - works for wildcard routes as well - # Get model_group and id from kwargs like the sync version does - model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None) - model_info: Final = kwargs["litellm_params"].get("model_info", {}) or {} - id = model_info.get("id", None) - if model_group is None or id is None: - return - elif isinstance(id, int): - id = str(id) + litellm_params: Final = kwargs["litellm_params"] + metadata: Final = litellm_params.get("metadata") + if metadata is None: + return + model_group: Final = metadata.get("model_group", None) + model_info: Final = litellm_params.get("model_info", {}) or {} + deployment_id: Final = model_info.get("id", None) + if model_group is None or deployment_id is None or self.get_deployment(model_id=str(deployment_id)) is None: + return - ## get deployment info - deployment_info: Final = self.get_deployment(model_id=id) + # Always track deployment successes for cooldown logic, regardless of TPM/RPM limits + increment_deployment_successes_for_current_minute( + litellm_router_instance=self, + deployment_id=str(deployment_id), + ) - if deployment_info is None: - return - else: - deployment_model_info: Final = self.get_router_model_info( - deployment=deployment_info, - received_model_name=model_group, - ) - # get tpm/rpm from deployment info - tpm: Final = deployment_info.get("tpm", None) - rpm: Final = deployment_info.get("rpm", None) - - ## check tpm/rpm in litellm_params - tpm_litellm_params: Final = deployment_info.litellm_params.tpm - rpm_litellm_params: Final = deployment_info.litellm_params.rpm - - ## check tpm/rpm in model_info - tpm_model_info: Final = deployment_model_info.get("tpm", None) - rpm_model_info: Final = deployment_model_info.get("rpm", None) - - # Always track deployment successes for cooldown logic, regardless of TPM/RPM limits - increment_deployment_successes_for_current_minute( - litellm_router_instance=self, - deployment_id=id, - ) - - deployment_dict = deployment_info if isinstance(deployment_info, dict) else deployment_info.model_dump() - has_io_token_limits: Final = deployment_has_io_token_limits(deployment_dict) - - ## Nothing to track only when neither tpm/rpm nor itpm/otpm limits are - ## set. IO deployments still record TPM/RPM usage here so TPM-aware - ## routing strategies see their real load in mixed model groups; their - ## itpm/otpm enforcement runs separately in ModelRateLimitingCheck. - if ( - tpm is None - and rpm is None - and tpm_litellm_params is None - and rpm_litellm_params is None - and tpm_model_info is None - and rpm_model_info is None - and not has_io_token_limits - ): - return - - parent_otel_span: Final = _get_parent_otel_span_from_kwargs(kwargs) - total_tokens: Final[float] = standard_logging_object.get("total_tokens", 0) - - # ------------ - # Setup values - # ------------ - dt: Final = get_utc_datetime() - current_minute: Final = dt.strftime("%H-%M") # use the same timezone regardless of system clock - - tpm_key = RouterCacheEnum.TPM.value.format(id=id, current_minute=current_minute, model=deployment_name) - # ------------ - # Update usage - # ------------ - # update cache - pipeline_operations: Final[list[RedisPipelineIncrementOperation]] = [] - - ## TPM - pipeline_operations.append( - RedisPipelineIncrementOperation( - key=tpm_key, - increment_value=total_tokens, - ttl=RoutingArgs.ttl.value, - ) - ) - - ## RPM - rpm_key = RouterCacheEnum.RPM.value.format(id=id, current_minute=current_minute, model=deployment_name) - pipeline_operations.append( - RedisPipelineIncrementOperation( - key=rpm_key, - increment_value=1, - ttl=RoutingArgs.ttl.value, - ) - ) - - await self.cache.async_increment_cache_pipeline( - increment_list=pipeline_operations, - parent_otel_span=parent_otel_span, - ) - - return tpm_key + total_tokens: Final[float] = standard_logging_object.get("total_tokens", 0) + counted_tokens: Final = get_counted_usage_tokens(litellm_params) + deployment_name: Final = metadata.get("deployment", None) + return await self._increment_deployment_usage( + deployment_id=str(deployment_id), + deployment_name=deployment_name if isinstance(deployment_name, str) else None, + model_group=model_group, + total_tokens=total_tokens if counted_tokens is None else max(0, total_tokens - counted_tokens), + rpm_increment=1 if counted_tokens is None else 0, + parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs), + ) except Exception as e: verbose_router_logger.debug( "litellm.router.Router::deployment_callback_on_success(): Exception occured - %s", e ) + async def increment_deployment_usage_for_response( + self, + response: object, + request_kwargs: dict[str, object], + ) -> None: + if response is None: + return + try: + deployment_metadata: Final = find_deployment_metadata(request_kwargs) + model_group: Final = request_kwargs.get("model") + if deployment_metadata is None or not isinstance(model_group, str): + return + model_info: Final = deployment_metadata["model_info"] + deployment_id: Final = model_info.get("id") if isinstance(model_info, dict) else None + if deployment_id is None: + return + total_tokens: Final = response_total_token_count(response) + deployment_name: Final = deployment_metadata.get("deployment") + deployment_metadata[ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY] = total_tokens + try: + await self._increment_deployment_usage( + deployment_id=str(deployment_id), + deployment_name=deployment_name if isinstance(deployment_name, str) else None, + model_group=model_group, + total_tokens=total_tokens, + rpm_increment=1, + parent_otel_span=_get_parent_otel_span_from_kwargs(request_kwargs), + ) + except Exception: + deployment_metadata.pop(ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, None) + raise + except Exception as e: + verbose_router_logger.debug( + "litellm.router.Router::increment_deployment_usage_for_response(): Exception occured - %s", e + ) + + async def _increment_deployment_usage( + self, + *, + deployment_id: str, + deployment_name: str | None, + model_group: str, + total_tokens: float, + rpm_increment: int, + parent_otel_span: Span | None, + ) -> str | None: + from litellm.types.caching import RedisPipelineIncrementOperation + + deployment_info: Final = self.get_deployment(model_id=deployment_id) + if deployment_info is None: + return None + deployment_model_info: Final = self.get_router_model_info( + deployment=deployment_info, + received_model_name=model_group, + ) + configured_limits: Final = ( + deployment_info.get("tpm", None), + deployment_info.get("rpm", None), + deployment_info.litellm_params.tpm, + deployment_info.litellm_params.rpm, + deployment_model_info.get("tpm", None), + deployment_model_info.get("rpm", None), + ) + ## Nothing to track only when neither tpm/rpm nor itpm/otpm limits are + ## set. IO deployments still record TPM/RPM usage here so TPM-aware + ## routing strategies see their real load in mixed model groups; their + ## itpm/otpm enforcement runs separately in ModelRateLimitingCheck. + if all(limit is None for limit in configured_limits) and not deployment_has_io_token_limits( + deployment_info.model_dump() + ): + return None + if total_tokens <= 0 and rpm_increment <= 0: + return None + + current_minute: Final = get_utc_datetime().strftime("%H-%M") # use the same timezone regardless of system clock + tpm_key: Final = RouterCacheEnum.TPM.value.format( + id=deployment_id, current_minute=current_minute, model=deployment_name + ) + rpm_key: Final = RouterCacheEnum.RPM.value.format( + id=deployment_id, current_minute=current_minute, model=deployment_name + ) + pipeline_operations: Final[list[RedisPipelineIncrementOperation]] = [ + RedisPipelineIncrementOperation(key=key, increment_value=increment_value, ttl=RoutingArgs.ttl.value) + for key, increment_value in ((tpm_key, total_tokens), (rpm_key, rpm_increment)) + ] + post_increment_values: Final = await self.cache.async_increment_cache_pipeline( + increment_list=pipeline_operations, + parent_otel_span=parent_otel_span, + ) + if post_increment_values is not None and self.cache.redis_cache is not None: + for operation, value in zip(pipeline_operations, post_increment_values): + await self.cache.async_set_cache( + operation["key"], int(value), local_only=True, ttl=RoutingArgs.ttl.value + ) + return tpm_key + def sync_deployment_callback_on_success( self, kwargs, # kwargs to completion @@ -11205,15 +11232,7 @@ class Router: if model_group is not None: remaining_usage: Final = await self.get_remaining_model_group_usage(model_group) - # get_remaining_model_group_usage reads the router's TPM/RPM counter, - # which is incremented post-response by deployment_callback_on_success. - # Replay the in-flight increment for TPM/RPM only (LIT-2719); ITPM/OTPM - # counters are incremented at reservation time and must not be adjusted. - apply_remaining_usage_headers( - additional_headers, - remaining_usage, - response_in_flight_token_count(response), - ) + apply_remaining_usage_headers(additional_headers, remaining_usage) return response def _build_model_name_index(self, model_list: list) -> None: diff --git a/litellm/router_utils/add_retry_fallback_headers.py b/litellm/router_utils/add_retry_fallback_headers.py index bc88feef7d2..cbca5880b52 100644 --- a/litellm/router_utils/add_retry_fallback_headers.py +++ b/litellm/router_utils/add_retry_fallback_headers.py @@ -151,7 +151,7 @@ def apply_quality_router_decision_headers( additional_headers[header] = str(decision[field]) -def response_in_flight_token_count(response: object) -> int: +def response_total_token_count(response: object) -> int: usage: Final = response.get("usage") if isinstance(response, dict) else getattr(response, "usage", None) if usage is None: return 0 @@ -166,15 +166,10 @@ def response_in_flight_token_count(response: object) -> int: def apply_remaining_usage_headers( additional_headers: dict[str, object], remaining_usage: dict[str, int], - in_flight_tokens: int, ) -> None: - in_flight_delta: Final = { - "x-ratelimit-remaining-tokens": in_flight_tokens, - "x-ratelimit-remaining-requests": 1, - } for header, value in remaining_usage.items(): if value is not None and header not in additional_headers: - additional_headers[header] = value - in_flight_delta.get(header, 0) + additional_headers[header] = value def _normalize_hidden_params(hidden_params: object) -> dict[str, object]: diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 94164d0ea0c..d0abaed4d3a 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -421,6 +421,25 @@ async def _is_fallback_target_authorized( return False +async def _is_fallback_target_within_budget( + litellm_router: LitellmRouter, + fallback_entry: str | Mapping[str, object], + original_model_group: str, + kwargs: Mapping[str, object], +) -> bool: + budget_check: Final = litellm_router.fallback_budget_check + target: Final = _get_fallback_target_model_group(fallback_entry) + if budget_check is None or target is None or target == original_model_group: + return True + if await budget_check(model=target, request_kwargs=kwargs, llm_router=litellm_router): + return True + verbose_router_logger.info( + "Skipping fallback to model_group = %s: caller is over budget", + mask_sensitive_structure(fallback_entry), + ) + return False + + def references_provider_scoped_resource(kwargs: Mapping[str, object]) -> bool: """ True when a file, batch, or fine-tuning job operation names an id that only exists @@ -528,6 +547,8 @@ async def run_async_fallback( continue if not await _is_fallback_target_authorized(litellm_router, mg, original_model_group, kwargs): continue + if not await _is_fallback_target_within_budget(litellm_router, mg, original_model_group, kwargs): + continue attempt_key = fallback_attempt_key(mg) if attempt_key is not None: if attempt_key in attempted: diff --git a/litellm/router_utils/router_callbacks/track_deployment_metrics.py b/litellm/router_utils/router_callbacks/track_deployment_metrics.py index 01893e925bc..6b422e98ec8 100644 --- a/litellm/router_utils/router_callbacks/track_deployment_metrics.py +++ b/litellm/router_utils/router_callbacks/track_deployment_metrics.py @@ -9,8 +9,11 @@ get_deployment_failures_for_current_minute get_deployment_successes_for_current_minute """ +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final +from litellm.constants import ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY + if TYPE_CHECKING: from litellm.router import Router as _Router @@ -18,6 +21,26 @@ if TYPE_CHECKING: else: LitellmRouter = Any +_METADATA_CHANNELS: Final = ("litellm_metadata", "metadata") + + +def find_deployment_metadata(kwargs: Mapping[str, object]) -> dict[str, object] | None: + buckets: Final = (kwargs.get(channel) for channel in _METADATA_CHANNELS) + return next((bucket for bucket in buckets if isinstance(bucket, dict) and "model_info" in bucket), None) + + +def get_counted_usage_tokens(litellm_params: Mapping[str, object]) -> int | None: + buckets: Final = (litellm_params.get(channel) for channel in _METADATA_CHANNELS) + counted: Final = next( + ( + bucket[ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY] + for bucket in buckets + if isinstance(bucket, dict) and ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY in bucket + ), + None, + ) + return counted if isinstance(counted, int) and not isinstance(counted, bool) else None + def increment_deployment_successes_for_current_minute( litellm_router_instance: LitellmRouter, diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index e62c85f4599..32b20bb7931 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -33,15 +33,6 @@ def aocr( timeout_seconds: float | None = None, ) -> Future[dict[str, object]]: ... -_OCR_MAX_FILE_BYTES: int - -def _ocr_upload_document( - file_content: bytes, - file_name: str | None = None, - content_type: str | None = None, -) -> dict[str, str]: ... -def _ocr_file_document(document: Mapping[str, object]) -> dict[str, str]: ... -def _ocr_mime_type(file_name: str) -> str: ... def _ocr_lifecycle( request: LiteLLMOcrRequest, args: tuple[object, ...], @@ -139,15 +130,11 @@ class TokenCounter: def gil_stats() -> dict[str, int]: ... __all__ = [ - "_OCR_MAX_FILE_BYTES", "ResponsesWebSocketConnection", "RustBridgeDeclined", "RustUpstreamError", "TokenCounter", - "_ocr_file_document", "_ocr_lifecycle", - "_ocr_mime_type", - "_ocr_upload_document", "achat_completions", "amessages", "aocr", diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 92fe41ba717..6346e13f3ba 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -682,6 +682,12 @@ class BedrockGuardrailStreamingParams(BaseModel): "and the scan result lands in guardrail_information; a flagged response still ends the " "stream with a block message (disable_exception_on_block=true) or an error frame.", ) + streaming_buffer_release_on_scan: bool = Field( + default=False, + description="When buffering, scan the accumulated response every streaming_sampling_rate chunks " + "and release the withheld chunks once the scan passes, instead of holding everything to end of stream. " + "Flagged content is never released. Ignored when streaming_end_of_stream_only is true.", + ) @classmethod def from_extras(cls, extras: Mapping[str, object] | None) -> "BedrockGuardrailStreamingParams": diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/crowdstrike_aidr.py b/litellm/types/proxy/guardrails/guardrail_hooks/crowdstrike_aidr.py index 6beca030a3a..df1caab6af6 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/crowdstrike_aidr.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/crowdstrike_aidr.py @@ -4,6 +4,14 @@ from .base import GuardrailConfigModel class CrowdStrikeAIDRGuardrailConfigModelOptionalParams(BaseModel): + streaming_buffer_until_moderated: bool | None = Field( + default=None, + description="When True, withhold streamed chunks until moderation passes. Defaults to False when unset.", + ) + streaming_buffer_release_on_scan: bool | None = Field( + default=None, + description="When buffering, release withheld chunks after each passing scan. Defaults to False when unset.", + ) streaming_end_of_stream_only: bool | None = Field( default=None, description="If False (default when unset), post_call scans the accumulated streamed response every " diff --git a/litellm/types/router.py b/litellm/types/router.py index 7c3e4d6943f..584d2494db4 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -963,6 +963,19 @@ class FallbackAccessCheck(Protocol): async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], llm_router: "Router") -> bool: ... +class FallbackBudgetCheck(Protocol): + """ + Decides whether the caller behind `request_kwargs` is still within budget for fallback `model`. + + Budget is enforced once during auth, against the *requested* model group. A fallback target is + chosen later, inside the router, so a zero-cost group that falls back to a priced one bills + without any budget gate. The router runs this before every cross-model-group fallback attempt + and skips targets it rejects, leaving the free attempt itself untouched. + """ + + async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], llm_router: "Router") -> bool: ... + + class AutoRouterCapabilityLimit(Protocol): """ Resolves how many complexity routers may claim each licensed capability right now; None means unlimited. diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 3abf80c74b7..92e5b1c4ff7 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -352,6 +352,17 @@ "supports_function_calling": true, "supports_pdf_input": true }, + "writer.palmyra-vision-7b": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 4096, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-writer-palmyra-vision-7b.html", + "supports_vision": true + }, "amazon.nova-lite-v1:0": { "cache_read_input_token_cost": 1.5e-08, "input_cost_per_token": 6e-08, @@ -3796,6 +3807,53 @@ "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false }, + "azure_ai/gpt-5.5-2026-04-24": { + "deprecation_date": "2027-10-26", + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "cache_read_input_token_cost_priority": 1e-06, + "cache_read_input_token_cost_above_272k_tokens_priority": 2e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, + "input_cost_per_token_priority": 1e-05, + "input_cost_per_token_above_272k_tokens_priority": 2e-05, + "litellm_provider": "azure_ai", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "output_cost_per_token_priority": 6e-05, + "output_cost_per_token_above_272k_tokens_priority": 9e-05, + "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false + }, "azure_ai/gpt-5.4": { "deprecation_date": "2027-09-02", "cache_read_input_token_cost": 2.5e-07, @@ -5311,6 +5369,7 @@ "supports_vision": true }, "azure/gpt-4o": { + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 1.25e-06, "input_cost_per_token": 2.5e-06, "litellm_provider": "azure", @@ -7510,6 +7569,70 @@ "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false }, + "azure/gpt-5.6-sol-2026-07-09": { + "cache_creation_input_token_cost": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 1e-05, + "cache_creation_input_token_cost_above_272k_tokens_flex": 5e-06, + "cache_creation_input_token_cost_priority": 1e-05, + "cache_creation_input_token_cost_above_272k_tokens_priority": 2e-05, + "cache_creation_input_token_cost_flex": 2.5e-06, + "cache_read_input_token_cost": 4e-07, + "cache_read_input_token_cost_above_272k_tokens": 8e-07, + "cache_read_input_token_cost_above_272k_tokens_flex": 4e-07, + "cache_read_input_token_cost_priority": 8e-07, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.6e-06, + "cache_read_input_token_cost_flex": 2e-07, + "deprecation_date": "2028-01-11", + "input_cost_per_token": 4e-06, + "input_cost_per_token_above_272k_tokens": 8e-06, + "input_cost_per_token_above_272k_tokens_flex": 4e-06, + "input_cost_per_token_priority": 8e-06, + "input_cost_per_token_above_272k_tokens_priority": 1.6e-05, + "input_cost_per_token_flex": 2e-06, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2e-05, + "output_cost_per_token_above_272k_tokens": 3e-05, + "output_cost_per_token_above_272k_tokens_flex": 1.5e-05, + "output_cost_per_token_priority": 4e-05, + "output_cost_per_token_above_272k_tokens_priority": 6e-05, + "output_cost_per_token_flex": 1e-05, + "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://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false + }, "azure/gpt-5.6-terra": { "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, @@ -7574,6 +7697,70 @@ "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false }, + "azure/gpt-5.6-terra-2026-07-09": { + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06, + "cache_creation_input_token_cost_priority": 5e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-05, + "cache_creation_input_token_cost_flex": 1.25e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_272k_tokens": 4e-07, + "cache_read_input_token_cost_above_272k_tokens_flex": 2e-07, + "cache_read_input_token_cost_priority": 4e-07, + "cache_read_input_token_cost_above_272k_tokens_priority": 8e-07, + "cache_read_input_token_cost_flex": 1e-07, + "deprecation_date": "2028-01-11", + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_272k_tokens": 4e-06, + "input_cost_per_token_above_272k_tokens_flex": 2e-06, + "input_cost_per_token_priority": 4e-06, + "input_cost_per_token_above_272k_tokens_priority": 8e-06, + "input_cost_per_token_flex": 1e-06, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_272k_tokens": 1.8e-05, + "output_cost_per_token_above_272k_tokens_flex": 9e-06, + "output_cost_per_token_priority": 2.4e-05, + "output_cost_per_token_above_272k_tokens_priority": 3.6e-05, + "output_cost_per_token_flex": 6e-06, + "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://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false + }, "azure/gpt-5.6-luna": { "cache_creation_input_token_cost": 2.5e-07, "cache_creation_input_token_cost_above_272k_tokens": 5e-07, @@ -7638,6 +7825,70 @@ "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false }, + "azure/gpt-5.6-luna-2026-07-09": { + "cache_creation_input_token_cost": 2.5e-07, + "cache_creation_input_token_cost_above_272k_tokens": 5e-07, + "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-07, + "cache_creation_input_token_cost_priority": 5e-07, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1e-06, + "cache_creation_input_token_cost_flex": 1.25e-07, + "cache_read_input_token_cost": 2e-08, + "cache_read_input_token_cost_above_272k_tokens": 4e-08, + "cache_read_input_token_cost_above_272k_tokens_flex": 2e-08, + "cache_read_input_token_cost_priority": 4e-08, + "cache_read_input_token_cost_above_272k_tokens_priority": 8e-08, + "cache_read_input_token_cost_flex": 1e-08, + "deprecation_date": "2028-01-11", + "input_cost_per_token": 2e-07, + "input_cost_per_token_above_272k_tokens": 4e-07, + "input_cost_per_token_above_272k_tokens_flex": 2e-07, + "input_cost_per_token_priority": 4e-07, + "input_cost_per_token_above_272k_tokens_priority": 8e-07, + "input_cost_per_token_flex": 1e-07, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "output_cost_per_token_above_272k_tokens": 1.8e-06, + "output_cost_per_token_above_272k_tokens_flex": 9e-07, + "output_cost_per_token_priority": 2.4e-06, + "output_cost_per_token_above_272k_tokens_priority": 3.6e-06, + "output_cost_per_token_flex": 6e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models" + }, "azure/gpt-6-astra": { "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, @@ -7686,6 +7937,54 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure/gpt-6-astra-2026-09-03": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, + "cache_read_input_token_cost": 1e-06, + "cache_read_input_token_cost_above_272k_tokens": 2e-06, + "input_cost_per_token": 1e-05, + "input_cost_per_token_above_272k_tokens": 2e-05, + "litellm_provider": "azure", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "output_cost_per_token_above_272k_tokens": 7.5e-05, + "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://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": false, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "azure/gpt-chat-latest": { "cache_read_input_token_cost": 5e-07, "deprecation_date": "2026-12-02", @@ -8502,6 +8801,63 @@ "output_cost_per_token_batches": 1.5e-05, "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-5.5-2026-04-24": { + "deprecation_date": "2027-10-26", + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "cache_read_input_token_cost_priority": 1.25e-06, + "cache_read_input_token_cost_above_272k_tokens_priority": 2e-06, + "cache_read_input_token_cost_flex": 2.5e-07, + "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, + "input_cost_per_token_above_272k_tokens_flex": 5e-06, + "input_cost_per_token_priority": 1.25e-05, + "input_cost_per_token_above_272k_tokens_priority": 2e-05, + "input_cost_per_token_flex": 2.5e-06, + "litellm_provider": "azure", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 3e-05, + "output_cost_per_token_batches": 1.5e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, + "output_cost_per_token_priority": 7.5e-05, + "output_cost_per_token_above_272k_tokens_priority": 9e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models" + }, "azure/us/gpt-5.5-2026-04-23": { "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, @@ -8550,6 +8906,57 @@ "output_cost_per_token_batches": 1.65e-05, "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/us/gpt-5.5-2026-04-24": { + "deprecation_date": "2027-10-26", + "cache_read_input_token_cost": 5.5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, + "cache_read_input_token_cost_priority": 1.375e-06, + "input_cost_per_token": 5.5e-06, + "input_cost_per_token_batches": 2.75e-06, + "input_cost_per_token_above_272k_tokens": 1.1e-05, + "input_cost_per_token_priority": 1.375e-05, + "output_cost_per_token": 3.3e-05, + "output_cost_per_token_batches": 1.65e-05, + "output_cost_per_token_above_272k_tokens": 4.95e-05, + "output_cost_per_token_priority": 8.25e-05, + "litellm_provider": "azure", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models" + }, "azure/eu/gpt-5.5-2026-04-23": { "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, @@ -8598,6 +9005,57 @@ "output_cost_per_token_batches": 1.65e-05, "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/eu/gpt-5.5-2026-04-24": { + "deprecation_date": "2027-10-26", + "cache_read_input_token_cost": 5.5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, + "cache_read_input_token_cost_priority": 1.375e-06, + "input_cost_per_token": 5.5e-06, + "input_cost_per_token_batches": 2.75e-06, + "input_cost_per_token_above_272k_tokens": 1.1e-05, + "input_cost_per_token_priority": 1.375e-05, + "output_cost_per_token": 3.3e-05, + "output_cost_per_token_batches": 1.65e-05, + "output_cost_per_token_above_272k_tokens": 4.95e-05, + "output_cost_per_token_priority": 8.25e-05, + "litellm_provider": "azure", + "max_input_tokens": 1050000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/batch", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_none_reasoning_effort": true, + "supports_xhigh_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models" + }, "azure/gpt-5.5-pro": { "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, @@ -9055,6 +9513,36 @@ "supports_vision": true, "supports_pdf_input": true }, + "azure/gpt-image-2.5-flare": { + "deprecation_date": "2027-09-09", + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_image_token": 8e-06, + "litellm_provider": "azure", + "mode": "image_generation", + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "azure/gpt-image-2.5-sunburst": { + "deprecation_date": "2027-09-09", + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_token": 5e-06, + "input_cost_per_image_token": 8e-06, + "litellm_provider": "azure", + "mode": "image_generation", + "output_cost_per_image_token": 3e-05, + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, "azure/gpt-image-2-2026-04-21": { "cache_read_input_token_cost": 1.25e-06, "deprecation_date": "2027-10-21", @@ -11371,6 +11859,7 @@ "output_cost_per_token": 2.5e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, @@ -11389,6 +11878,7 @@ "output_cost_per_token": 2.5e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, @@ -13904,7 +14394,7 @@ "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "anthropic", - "max_input_tokens": 200000, + "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", @@ -13939,7 +14429,7 @@ "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "anthropic", - "max_input_tokens": 200000, + "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", @@ -13960,7 +14450,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "source": "https://docs.anthropic.com/en/docs/about-claude/pricing" }, "claude-sonnet-5": { "deprecation_date": "2027-06-30", @@ -15089,6 +15580,21 @@ "supports_function_calling": true, "supports_tool_choice": true }, + "command-a-plus-05-2026": { + "input_cost_per_token": 0.0, + "litellm_provider": "cohere_chat", + "max_input_tokens": 128000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 0.0, + "source": "https://docs.cohere.com/docs/command-a-plus", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "command-light": { "input_cost_per_token": 3e-07, "litellm_provider": "cohere_chat", @@ -26884,13 +27390,71 @@ "supports_image_size": false }, "gemini/gemini-flash-latest": { + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_token": 7.5e-07, + "litellm_provider": "gemini", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 3.75e-06, + "output_cost_per_token": 3.75e-06, + "rpm": 15, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_output": false, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_vision": true, + "supports_web_search": true, + "tpm": 250000, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "google_maps_grounding_cost_per_query": 0.014, + "cache_read_input_token_cost_flex": 3.75e-08, + "cache_read_input_token_cost_priority": 1.35e-07, + "input_cost_per_token_batches": 3.75e-07, + "input_cost_per_token_flex": 3.75e-07, + "input_cost_per_token_priority": 1.35e-06, + "output_cost_per_token_batches": 1.875e-06, + "output_cost_per_token_flex": 1.875e-06, + "output_cost_per_token_priority": 6.75e-06, + "prompt_cache_min_tokens": 4096, + "supports_audio_input": true, + "supports_native_streaming": true, + "supports_video_input": true, + "web_search_billing_unit": "per_query" + }, + "gemini/gemini-flash-lite-latest": { "cache_read_input_token_cost": 3e-08, - "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", "output_cost_per_reasoning_token": 2.5e-06, "output_cost_per_token": 2.5e-06, @@ -26924,58 +27488,23 @@ "supports_web_search": true, "tpm": 250000, "search_context_cost_per_query": { - "search_context_size_low": 0.035, - "search_context_size_medium": 0.035, - "search_context_size_high": 0.035 + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 }, - "google_maps_grounding_cost_per_query": 0.025 - }, - "gemini/gemini-flash-lite-latest": { - "cache_read_input_token_cost": 1e-08, - "input_cost_per_audio_token": 3e-07, - "input_cost_per_token": 1e-07, - "litellm_provider": "gemini", - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, - "mode": "chat", - "output_cost_per_reasoning_token": 4e-07, - "output_cost_per_token": 4e-07, - "rpm": 15, - "source": "https://ai.google.dev/gemini-api/docs/pricing", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions", - "/v1/batch" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_audio_output": false, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_reasoning": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_vision": true, - "supports_web_search": true, - "tpm": 250000, - "search_context_cost_per_query": { - "search_context_size_low": 0.035, - "search_context_size_medium": 0.035, - "search_context_size_high": 0.035 - }, - "google_maps_grounding_cost_per_query": 0.025 + "google_maps_grounding_cost_per_query": 0.014, + "cache_read_input_token_cost_flex": 2e-08, + "cache_read_input_token_cost_priority": 5e-08, + "input_cost_per_token_batches": 1.5e-07, + "input_cost_per_token_flex": 1.5e-07, + "input_cost_per_token_priority": 5.4e-07, + "output_cost_per_token_batches": 1.25e-06, + "output_cost_per_token_flex": 1.25e-06, + "output_cost_per_token_priority": 4.5e-06, + "supports_audio_input": true, + "supports_native_streaming": true, + "supports_video_input": true, + "web_search_billing_unit": "per_query" }, "gemini/gemini-2.5-flash-lite-preview-06-17": { "deprecation_date": "2025-11-18", @@ -32153,7 +32682,7 @@ "input_cost_per_token": 4e-06, "input_cost_per_token_above_272k_tokens": 8e-06, "litellm_provider": "openai", - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -32192,7 +32721,7 @@ "input_cost_per_token": 4e-06, "input_cost_per_token_above_272k_tokens": 8e-06, "litellm_provider": "openai", - "max_input_tokens": 1050000, + "max_input_tokens": 922000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", @@ -34977,6 +35506,7 @@ "supports_tool_choice": true }, "inception/mercury-2.5": { + "cache_read_input_token_cost": 2e-08, "input_cost_per_token": 2e-07, "litellm_provider": "inception", "max_input_tokens": 260000, @@ -34986,6 +35516,7 @@ "output_cost_per_token": 7.5e-07, "source": "https://docs.inceptionlabs.ai/get-started/models", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true @@ -36525,6 +37056,57 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "mistral/zai-glm-5-3": { + "cache_read_input_token_cost": 1.4e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "mistral", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://docs.mistral.ai/models/zai-glm-5-3", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "mistral/zai-glm-5": { + "cache_read_input_token_cost": 1.4e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "mistral", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://docs.mistral.ai/models/zai-glm-5-3", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "mistral/zai-glm-latest": { + "cache_read_input_token_cost": 1.4e-07, + "input_cost_per_token": 1.4e-06, + "litellm_provider": "mistral", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://docs.mistral.ai/models/zai-glm-5-3", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "mistral/glm-5-2": { "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.4e-06, @@ -38151,6 +38733,15 @@ "supports_reasoning": true, "source": "https://tokenfactory.nebius.com/models/catalog/text2text/deepseek-ai%2FDeepSeek-V4-Pro" }, + "nebius/deepseek-ai/DeepSeek-V4-Pro-0813": { + "input_cost_per_token": 1.32e-06, + "litellm_provider": "nebius", + "mode": "chat", + "output_cost_per_token": 3.96e-06, + "source": "https://tokenfactory.nebius.com/models/catalog/text2text/deepseek-ai%2FDeepSeek-V4-Pro-0813", + "supports_function_calling": true, + "supports_reasoning": true + }, "nebius/MiniMaxAI/MiniMax-M2.5": { "max_tokens": 196608, "max_input_tokens": 196608, @@ -38417,6 +39008,17 @@ "supports_reasoning": true, "source": "https://tokenfactory.nebius.com/models/catalog/text2text/zai-org%2FGLM-5.2" }, + "nebius/zai-org/GLM-5.3": { + "input_cost_per_token": 1.4e-06, + "litellm_provider": "nebius", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://tokenfactory.nebius.com/models/catalog/text2text/zai-org%2FGLM-5.3", + "supports_function_calling": true, + "supports_reasoning": true + }, "nebius/zai-org/GLM-5.3-Flash": { "max_tokens": 1024000, "max_input_tokens": 1024000, @@ -40250,15 +40852,16 @@ "supports_tool_choice": true }, "openrouter/deepseek/deepseek-chat": { - "input_cost_per_token": 3.2e-07, + "input_cost_per_token": 2.574e-07, "litellm_provider": "openrouter", "max_input_tokens": 65536, "max_output_tokens": 8192, "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 8.9e-07, + "output_cost_per_token": 1.0287e-06, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "source": "https://openrouter.ai/api/v1/models" }, "openrouter/deepseek/deepseek-chat-v3-0324": { "input_cost_per_token": 2.5e-07, @@ -40351,21 +40954,21 @@ "supports_tool_choice": true }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 8.59908e-07, + "input_cost_per_token": 1.6e-06, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.719816e-06, + "output_cost_per_token": 3.2e-06, "source": "https://openrouter.ai/deepseek/deepseek-v4-pro", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.1659e-08 + "cache_read_input_token_cost": 1.35e-07 }, "openrouter/deepseek/deepseek-v4.1-flash": { "input_cost_per_token": 1.5e-07, @@ -40660,12 +41263,13 @@ "supports_vision": true }, "openrouter/gryphe/mythomax-l2-13b": { - "input_cost_per_token": 6e-08, + "input_cost_per_token": 8e-08, "litellm_provider": "openrouter", "max_tokens": 8192, "mode": "chat", - "output_cost_per_token": 6e-08, - "supports_tool_choice": true + "output_cost_per_token": 1.1e-07, + "supports_tool_choice": true, + "source": "https://openrouter.ai/api/v1/models" }, "openrouter/mancer/weaver": { "input_cost_per_token": 4e-07, @@ -40801,14 +41405,15 @@ "max_output_tokens": 131072 }, "openrouter/mistralai/mistral-small-3.2-24b-instruct": { - "input_cost_per_token": 7.5e-08, + "input_cost_per_token": 9.375e-08, "litellm_provider": "openrouter", "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 2e-07, + "output_cost_per_token": 2.5e-07, "supports_tool_choice": true, "max_input_tokens": 128000, - "max_output_tokens": 128000 + "max_output_tokens": 128000, + "source": "https://openrouter.ai/api/v1/models" }, "openrouter/mistralai/mixtral-8x22b-instruct": { "input_cost_per_token": 2e-06, @@ -41329,13 +41934,13 @@ "supports_tool_choice": true }, "openrouter/qwen/qwen3-235b-a22b-2507": { - "input_cost_per_token": 2.2e-07, + "input_cost_per_token": 8.75e-08, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 262144, "max_tokens": 262144, "mode": "chat", - "output_cost_per_token": 8.8e-07, + "output_cost_per_token": 3.5e-07, "source": "https://openrouter.ai/qwen/qwen3-235b-a22b-2507", "supports_function_calling": true, "supports_tool_choice": true @@ -41368,13 +41973,13 @@ "supports_vision": true }, "openrouter/qwen/qwen3.5-35b-a3b": { - "input_cost_per_token": 3.125e-07, + "input_cost_per_token": 1.625e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", - "output_cost_per_token": 1.25e-06, + "output_cost_per_token": 1.3e-06, "source": "https://openrouter.ai/qwen/qwen3.5-35b-a3b", "supports_function_calling": true, "supports_reasoning": true, @@ -44390,7 +44995,7 @@ "deprecation_date": "2026-06-04", "input_cost_per_token": 2e-06, "litellm_provider": "together_ai", - "max_input_tokens": 256000, + "max_input_tokens": 262144, "mode": "chat", "output_cost_per_token": 2e-06, "source": "https://api.together.ai/v1/models", @@ -44460,7 +45065,7 @@ "supports_parallel_function_calling": true, "supports_reasoning": true, "supports_tool_choice": true, - "max_input_tokens": 128000, + "max_input_tokens": 131072, "max_output_tokens": 16384 }, "together_ai/meta-llama/Llama-3.2-3B-Instruct-Turbo": { @@ -44643,7 +45248,7 @@ "deprecation_date": "2026-04-02", "input_cost_per_token": 2e-07, "litellm_provider": "together_ai", - "max_input_tokens": 128000, + "max_input_tokens": 131072, "mode": "chat", "output_cost_per_token": 1.1e-06, "source": "https://api.together.ai/v1/models", @@ -44870,6 +45475,20 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "together_ai/deepseek-ai/DeepSeek-V4.1-Flash": { + "cache_read_input_token_cost": 6e-09, + "input_cost_per_token": 3e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://api.together.xyz/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "together_ai/deepseek-ai/DeepSeek-V4-Pro": { "deprecation_date": "2026-08-27", "cache_read_input_token_cost": 2e-07, @@ -45199,9 +45818,10 @@ "output_cost_per_token": 1.25e-05, "supports_function_calling": true, "supports_pdf_input": true, - "supports_prompt_caching": false, + "supports_prompt_caching": true, "supports_response_schema": true, - "supports_vision": true + "supports_vision": true, + "cache_read_input_token_cost": 6.25e-07 }, "us.amazon.nova-pro-v1:0": { "cache_read_input_token_cost": 2e-07, @@ -49453,6 +50073,7 @@ "output_cost_per_token": 5e-07, "source": "https://docs.x.ai/developers/models", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, @@ -49469,6 +50090,7 @@ "output_cost_per_token": 5e-07, "source": "https://docs.x.ai/developers/models", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -49489,6 +50111,7 @@ "output_cost_per_token_above_200k_tokens": 5e-06, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, @@ -49508,6 +50131,7 @@ "output_cost_per_token_above_200k_tokens": 5e-06, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -49528,6 +50152,7 @@ "output_cost_per_token_above_200k_tokens": 5e-06, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -49547,6 +50172,7 @@ "output_cost_per_token_above_200k_tokens": 1.2e-05, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -56015,6 +56641,73 @@ "input_cost_per_second": 8.33333333333e-05, "supports_response_schema": false }, + "gemini-3.8-live": { + "input_cost_per_audio_token": 3e-06, + "input_cost_per_image_token": 1e-06, + "input_cost_per_token": 7.5e-07, + "input_cost_per_video_per_second": 3.3333333333333335e-05, + "litellm_provider": "gemini", + "max_input_tokens": 131072, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "realtime", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 4.5e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_vision": true, + "supports_web_search": true, + "gemini_audio_only_live": true + }, + "gemini-3.8-live-extended-thinking": { + "input_cost_per_audio_token": 3e-06, + "input_cost_per_image_token": 1e-06, + "input_cost_per_token": 7.5e-07, + "input_cost_per_video_per_second": 3.3333333333333335e-05, + "litellm_provider": "gemini", + "max_input_tokens": 131072, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "realtime", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 4.5e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_vision": true, + "supports_web_search": true, + "gemini_audio_only_live": true, + "supports_reasoning": true + }, "gemini/gemini-2.5-flash-native-audio-latest": { "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 5e-07, @@ -56180,18 +56873,17 @@ "supports_web_search": false }, "gemini-flash-latest": { - "cache_read_input_token_cost": 3e-08, - "input_cost_per_audio_token": 1e-06, - "input_cost_per_token": 3e-07, + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_token": 7.5e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", - "output_cost_per_reasoning_token": 2.5e-06, - "output_cost_per_token": 2.5e-06, + "output_cost_per_reasoning_token": 3.75e-06, + "output_cost_per_token": 3.75e-06, "rpm": 100000, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-preview", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -56220,25 +56912,37 @@ "supports_web_search": true, "tpm": 8000000, "search_context_cost_per_query": { - "search_context_size_low": 0.035, - "search_context_size_medium": 0.035, - "search_context_size_high": 0.035 + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 }, - "google_maps_grounding_cost_per_query": 0.025 + "google_maps_grounding_cost_per_query": 0.014, + "cache_read_input_token_cost_flex": 3.75e-08, + "cache_read_input_token_cost_priority": 1.35e-07, + "input_cost_per_token_batches": 3.75e-07, + "input_cost_per_token_flex": 3.75e-07, + "input_cost_per_token_priority": 1.35e-06, + "output_cost_per_token_batches": 1.875e-06, + "output_cost_per_token_flex": 1.875e-06, + "output_cost_per_token_priority": 6.75e-06, + "prompt_cache_min_tokens": 4096, + "supports_audio_input": true, + "supports_native_streaming": true, + "supports_video_input": true, + "web_search_billing_unit": "per_query" }, "gemini-flash-lite-latest": { - "cache_read_input_token_cost": 1e-08, - "input_cost_per_audio_token": 3e-07, - "input_cost_per_token": 1e-07, + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 3e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", - "output_cost_per_reasoning_token": 4e-07, - "output_cost_per_token": 4e-07, + "output_cost_per_reasoning_token": 2.5e-06, + "output_cost_per_token": 2.5e-06, "rpm": 15, - "source": "https://ai.google.dev/gemini-api/docs/models#gemini-2.5-flash-lite", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -56267,29 +56971,42 @@ "supports_web_search": true, "tpm": 250000, "search_context_cost_per_query": { - "search_context_size_low": 0.035, - "search_context_size_medium": 0.035, - "search_context_size_high": 0.035 + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 }, - "google_maps_grounding_cost_per_query": 0.025 + "google_maps_grounding_cost_per_query": 0.014, + "cache_read_input_token_cost_flex": 2e-08, + "cache_read_input_token_cost_priority": 5e-08, + "input_cost_per_token_batches": 1.5e-07, + "input_cost_per_token_flex": 1.5e-07, + "input_cost_per_token_priority": 5.4e-07, + "output_cost_per_token_batches": 1.25e-06, + "output_cost_per_token_flex": 1.25e-06, + "output_cost_per_token_priority": 4.5e-06, + "supports_audio_input": true, + "supports_native_streaming": true, + "supports_video_input": true, + "web_search_billing_unit": "per_query" }, "gemini-pro-latest": { - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, "rpm": 2000, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", - "/v1/completions" + "/v1/completions", + "/v1/batch" ], "supported_modalities": [ "text", @@ -56313,29 +57030,45 @@ "supports_web_search": true, "tpm": 800000, "search_context_cost_per_query": { - "search_context_size_low": 0.035, - "search_context_size_medium": 0.035, - "search_context_size_high": 0.035 + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 }, - "google_maps_grounding_cost_per_query": 0.025 + "google_maps_grounding_cost_per_query": 0.014, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, + "cache_read_input_token_cost_priority": 3.6e-07, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, + "input_cost_per_token_batches": 1e-06, + "input_cost_per_token_priority": 3.6e-06, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, + "output_cost_per_token_batches": 6e-06, + "output_cost_per_token_priority": 2.16e-05, + "prompt_cache_min_tokens": 4096, + "supports_native_streaming": true, + "supports_url_context": true, + "web_search_billing_unit": "per_query", + "cache_read_input_token_cost_flex": 2e-07, + "input_cost_per_token_flex": 1e-06, + "output_cost_per_token_flex": 6e-06 }, "gemini/gemini-pro-latest": { - "cache_read_input_token_cost": 1.25e-07, - "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "input_cost_per_token": 1.25e-06, - "input_cost_per_token_above_200k_tokens": 2.5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, "litellm_provider": "gemini", "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_tokens": 65535, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", - "output_cost_per_token": 1e-05, - "output_cost_per_token_above_200k_tokens": 1.5e-05, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, "rpm": 2000, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", - "/v1/completions" + "/v1/completions", + "/v1/batch" ], "supported_modalities": [ "text", @@ -56359,11 +57092,26 @@ "supports_web_search": true, "tpm": 800000, "search_context_cost_per_query": { - "search_context_size_low": 0.035, - "search_context_size_medium": 0.035, - "search_context_size_high": 0.035 + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 }, - "google_maps_grounding_cost_per_query": 0.025 + "google_maps_grounding_cost_per_query": 0.014, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, + "cache_read_input_token_cost_priority": 3.6e-07, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, + "input_cost_per_token_batches": 1e-06, + "input_cost_per_token_priority": 3.6e-06, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, + "output_cost_per_token_batches": 6e-06, + "output_cost_per_token_priority": 2.16e-05, + "prompt_cache_min_tokens": 4096, + "supports_native_streaming": true, + "supports_url_context": true, + "web_search_billing_unit": "per_query", + "cache_read_input_token_cost_flex": 2e-07, + "input_cost_per_token_flex": 1e-06, + "output_cost_per_token_flex": 6e-06 }, "gemini-exp-1206": { "cache_read_input_token_cost": 3e-08, @@ -57263,6 +58011,38 @@ } ] }, + "volcengine/doubao-seed-2-1-pro-260628": { + "cache_read_input_token_cost": 1.725e-07, + "input_cost_per_token": 8.625e-07, + "litellm_provider": "volcengine", + "max_input_tokens": 256000, + "max_output_tokens": 256000, + "max_tokens": 256000, + "mode": "chat", + "output_cost_per_token": 4.3125e-06, + "source": "https://www.volcengine.com/docs/82379/1544106", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": false, + "supports_vision": true + }, + "volcengine/doubao-seed-2-1-turbo-260628": { + "cache_read_input_token_cost": 8.625e-08, + "input_cost_per_token": 4.3125e-07, + "litellm_provider": "volcengine", + "max_input_tokens": 256000, + "max_output_tokens": 256000, + "max_tokens": 256000, + "mode": "chat", + "output_cost_per_token": 2.15625e-06, + "source": "https://www.volcengine.com/docs/82379/1544106", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": false, + "supports_vision": true + }, "volcengine/doubao-seed-2-0-lite-260215": { "litellm_provider": "volcengine", "max_input_tokens": 256000, @@ -58989,7 +59769,8 @@ "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_response_schema": true }, "fireworks_ai/accounts/fireworks/models/kimi-k3": { "cache_read_input_token_cost": 3e-07, @@ -59065,7 +59846,8 @@ "source": "https://api.fireworks.ai/v1/serverless/models", "supports_function_calling": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "supports_response_schema": true }, "fireworks_ai/glm-5p2-fast": { "cache_read_input_token_cost": 2.1e-07, @@ -62056,6 +62838,36 @@ "supports_tool_choice": true, "supports_vision": false }, + "fireworks_ai/accounts/fireworks/routers/glm-5p3-fast": { + "cache_read_input_token_cost": 3.9e-07, + "input_cost_per_token": 2.1e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 6.6e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "fireworks_ai/glm-5p3-fast": { + "cache_read_input_token_cost": 3.9e-07, + "input_cost_per_token": 2.1e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 6.6e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false + }, "fireworks_ai/accounts/fireworks/models/glm-5p3-flash": { "cache_read_input_token_cost": 3e-08, "cache_read_input_token_cost_priority": 3.75e-08, @@ -64225,9 +65037,9 @@ "supports_prompt_caching": true }, "openrouter/z-ai/glm-5.3-flash": { - "input_cost_per_token": 1.5e-07, - "output_cost_per_token": 5e-07, - "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 9e-08, + "output_cost_per_token": 3e-07, + "cache_read_input_token_cost": 1.8e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, "max_output_tokens": 131072, @@ -64277,9 +65089,9 @@ "supports_prompt_caching": true }, "openrouter/qwen/qwen3.8-27b": { - "input_cost_per_token": 4.2e-07, - "output_cost_per_token": 3e-06, - "cache_read_input_token_cost": 8.5e-08, + "input_cost_per_token": 2.14e-07, + "output_cost_per_token": 2.55e-06, + "cache_read_input_token_cost": 1.5e-07, "litellm_provider": "openrouter", "max_input_tokens": 1000000, "max_output_tokens": 131072, @@ -64362,9 +65174,9 @@ "supports_prompt_caching": true }, "openrouter/deepseek/deepseek-v4-flash-0731": { - "input_cost_per_token": 6.5e-08, - "output_cost_per_token": 1.8e-07, - "cache_read_input_token_cost": 1.6e-08, + "input_cost_per_token": 6e-08, + "output_cost_per_token": 1.2e-07, + "cache_read_input_token_cost": 1.2e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, "max_output_tokens": 943718, @@ -64432,9 +65244,9 @@ "supports_vision": false }, "openrouter/moonshotai/kimi-k3": { - "input_cost_per_token": 2.1e-06, - "output_cost_per_token": 1.053e-05, - "cache_read_input_token_cost": 2.35e-07, + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "cache_read_input_token_cost": 3e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 943718, @@ -64531,9 +65343,9 @@ "supports_prompt_caching": true }, "openrouter/z-ai/glm-5.2": { - "input_cost_per_token": 6e-07, - "output_cost_per_token": 2e-06, - "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.4e-06, + "output_cost_per_token": 4.4e-06, + "cache_read_input_token_cost": 1.4e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 131072, @@ -64564,9 +65376,9 @@ "supports_vision": false }, "openrouter/moonshotai/kimi-k2.7-code": { - "input_cost_per_token": 7.1e-07, - "output_cost_per_token": 3.5e-06, - "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 7.062e-07, + "output_cost_per_token": 3.21e-06, + "cache_read_input_token_cost": 1.8e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 235929, @@ -64868,8 +65680,8 @@ "supports_prompt_caching": true }, "openrouter/google/gemma-4-26b-a4b-it": { - "input_cost_per_token": 4.2e-08, - "output_cost_per_token": 2.2e-07, + "input_cost_per_token": 9e-08, + "output_cost_per_token": 3e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 16384, @@ -65416,8 +66228,8 @@ "supports_vision": true }, "openrouter/qwen/qwen3-vl-30b-a3b-instruct": { - "input_cost_per_token": 1.5e-07, - "output_cost_per_token": 6e-07, + "input_cost_per_token": 1.3e-07, + "output_cost_per_token": 5.2e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 16384, @@ -65677,8 +66489,8 @@ "supports_vision": false }, "openrouter/qwen/qwen3-30b-a3b-instruct-2507": { - "input_cost_per_token": 9e-08, - "output_cost_per_token": 3e-07, + "input_cost_per_token": 4.815e-08, + "output_cost_per_token": 1.9305e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 32000, @@ -65737,7 +66549,7 @@ "supports_vision": false }, "openrouter/minimax/minimax-m1": { - "input_cost_per_token": 5.5e-07, + "input_cost_per_token": 4e-07, "output_cost_per_token": 2.2e-06, "litellm_provider": "openrouter", "max_input_tokens": 1000000, @@ -65876,8 +66688,8 @@ "supports_vision": false }, "openrouter/qwen/qwen3-14b": { - "input_cost_per_token": 2.275e-07, - "output_cost_per_token": 9.1e-07, + "input_cost_per_token": 1.2e-07, + "output_cost_per_token": 2.4e-07, "litellm_provider": "openrouter", "max_input_tokens": 131072, "max_output_tokens": 16384, @@ -65939,8 +66751,8 @@ "supports_prompt_caching": true }, "openrouter/meta-llama/llama-4-maverick": { - "input_cost_per_token": 2e-07, - "output_cost_per_token": 6.96e-07, + "input_cost_per_token": 1.875e-07, + "output_cost_per_token": 6.525e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 115200, @@ -66361,14 +67173,6 @@ "supports_response_schema": true, "supports_vision": false }, - "fireworks_ai/accounts/fireworks/routers/glm-5p3-fast": { - "cache_read_input_token_cost": 3.9e-07, - "input_cost_per_token": 2.1e-06, - "litellm_provider": "fireworks_ai", - "mode": "chat", - "output_cost_per_token": 6.6e-06, - "source": "https://api.fireworks.ai/v1/serverless/models" - }, "together_ai/arcee-ai/trinity-mini": { "input_cost_per_token": 4.5e-08, "litellm_provider": "together_ai", @@ -66408,14 +67212,6 @@ "output_cost_per_token": 1.6e-06, "source": "https://api.together.ai/v1/models" }, - "together_ai/deepseek-ai/DeepSeek-V4.1-Flash": { - "cache_read_input_token_cost": 6e-09, - "input_cost_per_token": 3e-07, - "litellm_provider": "together_ai", - "mode": "chat", - "output_cost_per_token": 1.2e-06, - "source": "https://api.together.ai/v1/models" - }, "vertex_ai/gemini-2.5-flash-native-audio": { "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 5e-07, diff --git a/schema.prisma b/schema.prisma index d2375903c47..139fb031671 100644 --- a/schema.prisma +++ b/schema.prisma @@ -426,6 +426,7 @@ model LiteLLM_VerificationToken { key_alias String? soft_budget_cooldown Boolean @default(false) // key-level state on if budget alerts need to be cooled down spend Float @default(0.0) + total_spend Float @default(0.0) expires DateTime? models String[] aliases Json @default("{}") @@ -528,6 +529,7 @@ model LiteLLM_DeletedVerificationToken { key_alias String? soft_budget_cooldown Boolean @default(false) spend Float @default(0.0) + total_spend Float @default(0.0) expires DateTime? models String[] aliases Json @default("{}") diff --git a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py index 096efc33aaf..efe41e1da9a 100644 --- a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py +++ b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py @@ -587,6 +587,8 @@ def _success_kwargs( response_cost=0.5, key_hash=None, key_model_max_budget=None, + team_id=None, + team_model_max_budget=None, user_id=None, user_model_max_budget=None, end_user_id=None, @@ -600,6 +602,7 @@ def _success_kwargs( "end_user": end_user_id, "metadata": { "user_api_key_hash": key_hash, + "user_api_key_team_id": team_id, "user_api_key_user_id": user_id, "user_api_key_end_user_id": end_user_id, }, @@ -607,6 +610,7 @@ def _success_kwargs( "litellm_params": { "metadata": { "user_api_key_model_max_budget": key_model_max_budget, + "user_api_key_team_model_max_budget": team_model_max_budget, "user_api_key_user_model_max_budget": user_model_max_budget, "user_api_key_end_user_model_max_budget": end_user_model_max_budget, }, @@ -1417,3 +1421,266 @@ async def test_spend_logged_on_one_replica_is_enforced_and_reported_on_another() replica_c = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache(redis_cache=shared_redis)) with pytest.raises(litellm.BudgetExceededError): await replica_c.is_key_within_model_budget(user_api_key, "gpt-4") + + +def _log_success(limiter, **kwargs): + return limiter.async_log_success_event( + _success_kwargs(**kwargs), response_obj=None, start_time=None, end_time=None + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "request_model", + ["gpt-4", "openai/gpt-4"], + ids=["bare_model", "provider_prefixed_model"], +) +async def test_team_model_budget_is_shared_by_every_key_without_an_override(request_model): + """ + Two keys on the same team, neither carrying a matching key-level entry, + charge one team counter and are both refused once it is spent. + """ + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + team_model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1d"}} + check = lambda: limiter.is_team_within_model_budget( + team_id="team-1", + team_model_max_budget=team_model_max_budget, + key_model_max_budget=None, + model=request_model, + ) + + assert await check() is True + await _log_success( + limiter, + model_group=request_model, + response_cost=0.6, + key_hash="vk-a", + team_id="team-1", + team_model_max_budget=team_model_max_budget, + ) + assert await check() is True + await _log_success( + limiter, + model_group=request_model, + response_cost=0.6, + key_hash="vk-b", + team_id="team-1", + team_model_max_budget=team_model_max_budget, + ) + + assert await dual_cache.async_get_cache(key="team_model_spend:team-1:gpt-4:1d") == pytest.approx(1.2) + with pytest.raises(litellm.BudgetExceededError) as exc: + await check() + assert exc.value.entity_type == Litellm_EntityType.TEAM.value + assert await build_model_max_budget_usage( + entity_type=Litellm_EntityType.TEAM, + entity_id="team-1", + model_max_budget=team_model_max_budget, + cache=dual_cache, + ) == {"gpt-4": {"current_spend": pytest.approx(1.2), "budget_limit": 1.0, "time_period": "1d"}} + + +@pytest.mark.asyncio +async def test_key_override_replaces_the_team_cap_for_that_model(): + """ + A key with its own entry for the model is gated on the key counter alone: + the exhausted team counter does not block it, and its spend never lands on + the team counter. + """ + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + team_model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1d"}} + key_model_max_budget = {"gpt-4": {"budget_limit": 5.0, "time_period": "1d"}} + await dual_cache.async_set_cache(key="team_model_spend:team-1:gpt-4:1d", value=9.0) + + assert ( + await limiter.is_team_within_model_budget( + team_id="team-1", + team_model_max_budget=team_model_max_budget, + key_model_max_budget=key_model_max_budget, + model="openai/gpt-4", + ) + is True + ) + + await _log_success( + limiter, + model_group="openai/gpt-4", + response_cost=2.0, + key_hash="vk-override", + key_model_max_budget=key_model_max_budget, + team_id="team-1", + team_model_max_budget=team_model_max_budget, + ) + + assert await dual_cache.async_get_cache(key="team_model_spend:team-1:gpt-4:1d") == 9.0 + assert await dual_cache.async_get_cache(key="virtual_key_spend:vk-override:gpt-4:1d") == 2.0 + + +@pytest.mark.asyncio +async def test_key_entry_for_another_model_does_not_lift_the_team_cap(): + """A key override only covers the model it names; other models stay on the team counter.""" + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + team_model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1d"}} + key_model_max_budget = {"claude-3": {"budget_limit": 5.0, "time_period": "1d"}} + + await _log_success( + limiter, + model_group="gpt-4", + response_cost=1.5, + key_hash="vk-other", + key_model_max_budget=key_model_max_budget, + team_id="team-1", + team_model_max_budget=team_model_max_budget, + ) + + assert await dual_cache.async_get_cache(key="team_model_spend:team-1:gpt-4:1d") == 1.5 + with pytest.raises(litellm.BudgetExceededError): + await limiter.is_team_within_model_budget( + team_id="team-1", + team_model_max_budget=team_model_max_budget, + key_model_max_budget=key_model_max_budget, + model="gpt-4", + ) + + +@pytest.mark.asyncio +async def test_team_budget_leaves_unconfigured_models_alone(): + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + team_model_max_budget = {"gpt-4": {"budget_limit": 0.0, "time_period": "1d"}} + + assert ( + await limiter.is_team_within_model_budget( + team_id="team-1", + team_model_max_budget=team_model_max_budget, + key_model_max_budget=None, + model="claude-3", + ) + is True + ) + with patch.object(limiter, "_increment_spend_for_key", new_callable=AsyncMock) as mock_increment: + await _log_success( + limiter, + model_group="claude-3", + team_id="team-1", + team_model_max_budget=team_model_max_budget, + ) + mock_increment.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_team_counters_are_isolated_by_team_model_and_window(): + """Same model on two teams, and two models with different windows on one team, never share a counter.""" + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + team_model_max_budget = { + "gpt-4": {"budget_limit": 10.0, "time_period": "1d"}, + "claude-3": {"budget_limit": 10.0, "time_period": "30d"}, + } + + for team_id, model in (("team-1", "gpt-4"), ("team-2", "gpt-4"), ("team-1", "claude-3")): + await _log_success( + limiter, + model_group=model, + response_cost=1.0, + team_id=team_id, + team_model_max_budget=team_model_max_budget, + ) + + assert await dual_cache.async_get_cache(key="team_model_spend:team-1:gpt-4:1d") == 1.0 + assert await dual_cache.async_get_cache(key="team_model_spend:team-2:gpt-4:1d") == 1.0 + assert await dual_cache.async_get_cache(key="team_model_spend:team-1:claude-3:30d") == 1.0 + assert await dual_cache.async_get_cache(key="team_model_budget_start_time:team-1:claude-3:30d") is not None + + +@pytest.mark.asyncio +async def test_malformed_team_entry_is_skipped_and_its_sibling_still_enforced(): + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) + team_model_max_budget = { + "gpt-4": {"budget_limit": "not-a-number", "time_period": "1d"}, + "claude-3": {"budget_limit": 0.0, "time_period": "1d"}, + } + + assert ( + await limiter.is_team_within_model_budget( + team_id="team-1", + team_model_max_budget=team_model_max_budget, + key_model_max_budget=None, + model="gpt-4", + ) + is True + ) + with pytest.raises(litellm.BudgetExceededError): + await limiter.is_team_within_model_budget( + team_id="team-1", + team_model_max_budget=team_model_max_budget, + key_model_max_budget=None, + model="claude-3", + ) + + +@pytest.mark.asyncio +async def test_malformed_key_entry_does_not_count_as_an_override(): + """A key entry the limiter cannot enforce must not also switch the team cap off.""" + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + team_model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1d"}} + key_model_max_budget = {"gpt-4": {"budget_limit": "not-a-number", "time_period": "1d"}} + + await _log_success( + limiter, + model_group="gpt-4", + response_cost=1.5, + key_hash="vk-bad", + key_model_max_budget=key_model_max_budget, + team_id="team-1", + team_model_max_budget=team_model_max_budget, + ) + + assert await dual_cache.async_get_cache(key="team_model_spend:team-1:gpt-4:1d") == 1.5 + with pytest.raises(litellm.BudgetExceededError): + await limiter.is_team_within_model_budget( + team_id="team-1", + team_model_max_budget=team_model_max_budget, + key_model_max_budget=key_model_max_budget, + model="gpt-4", + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "key_entry", + [ + {"time_period": "1d", "tpm_limit": 100}, + {"time_period": "1d", "rpm_limit": 10}, + {"budget_limit": -1.0, "time_period": "1d"}, + ], +) +async def test_key_entry_without_a_spend_cap_does_not_lift_the_team_cap(key_entry): + """A key row that only rate-limits the model, or has no enforceable cap, leaves the team cap in force.""" + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + team_model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1d"}} + key_model_max_budget = {"gpt-4": key_entry} + + await _log_success( + limiter, + model_group="openai/gpt-4", + response_cost=1.5, + key_hash="vk-rate-limited", + key_model_max_budget=key_model_max_budget, + team_id="team-1", + team_model_max_budget=team_model_max_budget, + ) + + assert await dual_cache.async_get_cache(key="team_model_spend:team-1:gpt-4:1d") == 1.5 + with pytest.raises(litellm.BudgetExceededError): + await limiter.is_team_within_model_budget( + team_id="team-1", + team_model_max_budget=team_model_max_budget, + key_model_max_budget=key_model_max_budget, + model="openai/gpt-4", + ) diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index 14d86743557..bb389693311 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -1,3 +1,4 @@ +import asyncio import json import os import traceback @@ -10,9 +11,14 @@ import pytest import litellm from unittest.mock import patch, MagicMock, AsyncMock from create_mock_standard_logging_payload import create_standard_logging_payload -from litellm.types.utils import StandardLoggingPayload +from litellm.types.utils import ModelResponse, StandardLoggingPayload +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.types.caching import RedisPipelineIncrementOperation +from litellm.router_utils.router_callbacks.track_deployment_metrics import get_deployment_successes_for_current_minute from litellm.types.router import Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo -from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS +from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY @pytest.fixture @@ -928,18 +934,7 @@ async def test_set_response_headers(model_list): @pytest.mark.asyncio -async def test_set_response_headers_subtracts_in_flight_delta(model_list): - """ - LIT-2719: router-derived `x-ratelimit-remaining-*` headers must be - post-decrement (match OpenAI/Anthropic vendor semantics) so the proxy's - HTTP response headers and the prometheus gauges that read them stay - comparable across providers. - - Router's TPM/RPM counter is incremented post-response by - `deployment_callback_on_success`, so `get_remaining_model_group_usage` - sees pre-decrement values. `set_response_headers` must replay the - in-flight increment before writing the headers. - """ +async def test_set_response_headers_passes_through_post_increment_counters(model_list): from pydantic import BaseModel class _Usage(BaseModel): @@ -952,49 +947,10 @@ async def test_set_response_headers_subtracts_in_flight_delta(model_list): router = Router(model_list=model_list) router.get_remaining_model_group_usage = AsyncMock( return_value={ - "x-ratelimit-remaining-tokens": 1000, + "x-ratelimit-remaining-tokens": 958, "x-ratelimit-limit-tokens": 1000, - "x-ratelimit-remaining-requests": 100, + "x-ratelimit-remaining-requests": 99, "x-ratelimit-limit-requests": 100, - } - ) - - resp = _Resp() - resp._hidden_params = {} - await router.set_response_headers(response=resp, model_group="gpt-3.5-turbo") - - headers = resp._hidden_params["additional_headers"] - assert headers["x-ratelimit-remaining-tokens"] == 958 - assert headers["x-ratelimit-remaining-requests"] == 99 - # Limit headers pass through unmodified. - assert headers["x-ratelimit-limit-tokens"] == 1000 - assert headers["x-ratelimit-limit-requests"] == 100 - - -@pytest.mark.asyncio -async def test_set_response_headers_in_flight_delta_only_adjusts_tpm_rpm(model_list): - """ - The in-flight replay applies only to the post-incremented TPM/RPM counters - (`x-ratelimit-remaining-tokens` / `-requests`). The ITPM/OTPM counters are - incremented at reservation time (pre-call), so the input/output token - headers already reflect this request and must pass through untouched. - """ - from pydantic import BaseModel - - class _Usage(BaseModel): - total_tokens: int = 30 - prompt_tokens: int = 20 - completion_tokens: int = 10 - - class _Resp(BaseModel): - usage: _Usage = _Usage() - _hidden_params: dict = {} - - router = Router(model_list=model_list) - router.get_remaining_model_group_usage = AsyncMock( - return_value={ - "x-ratelimit-remaining-tokens": 1000, - "x-ratelimit-remaining-requests": 100, "x-ratelimit-remaining-input-tokens": 1000, "x-ratelimit-remaining-output-tokens": 500, } @@ -1005,14 +961,336 @@ async def test_set_response_headers_in_flight_delta_only_adjusts_tpm_rpm(model_l await router.set_response_headers(response=resp, model_group="gpt-3.5-turbo") headers = resp._hidden_params["additional_headers"] - # TPM/RPM headers replay the in-flight increment... - assert headers["x-ratelimit-remaining-tokens"] == 970 + assert headers["x-ratelimit-remaining-tokens"] == 958 assert headers["x-ratelimit-remaining-requests"] == 99 - # ...but the reservation-based input/output headers pass through unchanged. + assert headers["x-ratelimit-limit-tokens"] == 1000 + assert headers["x-ratelimit-limit-requests"] == 100 assert headers["x-ratelimit-remaining-input-tokens"] == 1000 assert headers["x-ratelimit-remaining-output-tokens"] == 500 +def _rpm_tpm_router(model_id: str) -> Router: + return Router( + model_list=[ + { + "model_name": "gpt-5-mini", + "litellm_params": {"model": "gpt-5-mini", "api_key": "sk-fake", "tpm": 1000, "rpm": 100}, + "model_info": {"id": model_id}, + } + ] + ) + + +def _ratelimit_headers(response: ModelResponse | CustomStreamWrapper) -> dict[str, int]: + return {k: v for k, v in response._hidden_params["additional_headers"].items() if k.startswith("x-ratelimit-")} + + +@pytest.mark.asyncio +async def test_acompletion_headers_read_post_increment_counter_and_count_once(): + router = _rpm_tpm_router("lit-3058-async") + + response = await router.acompletion( + model="gpt-5-mini", messages=[{"role": "user", "content": "hi"}], mock_response="pong" + ) + total_tokens = response.usage.total_tokens + assert total_tokens > 0 + + headers = _ratelimit_headers(response) + assert headers["x-ratelimit-remaining-tokens"] == 1000 - total_tokens + assert headers["x-ratelimit-remaining-requests"] == 99 + assert await router.get_model_group_usage("gpt-5-mini") == (total_tokens, 1) + + await asyncio.sleep(0.5) + assert await router.get_model_group_usage("gpt-5-mini") == (total_tokens, 1) + + +@pytest.mark.asyncio +async def test_acompletion_wildcard_route_headers_and_counter_use_resolved_deployment_name(): + router = Router( + model_list=[ + { + "model_name": "openai/*", + "litellm_params": {"model": "openai/*", "api_key": "sk-fake", "tpm": 1000, "rpm": 100}, + "model_info": {"id": "lit-3058-wildcard"}, + } + ] + ) + + response = await router.acompletion( + model="openai/gpt-5-mini", messages=[{"role": "user", "content": "hi"}], mock_response="pong" + ) + total_tokens = response.usage.total_tokens + + headers = _ratelimit_headers(response) + assert headers["x-ratelimit-remaining-tokens"] == 1000 - total_tokens + assert headers["x-ratelimit-remaining-requests"] == 99 + assert await router.get_model_group_usage("openai/gpt-5-mini") == (total_tokens, 1) + + +@pytest.mark.asyncio +async def test_acompletion_stream_counts_request_before_headers_and_tokens_once_on_completion(): + router = _rpm_tpm_router("lit-3058-stream") + + stream = await router.acompletion( + model="gpt-5-mini", + messages=[{"role": "user", "content": "hi"}], + mock_response="pong pong pong", + stream=True, + stream_options={"include_usage": True}, + ) + headers = _ratelimit_headers(stream) + assert headers["x-ratelimit-remaining-tokens"] == 1000 + assert headers["x-ratelimit-remaining-requests"] == 99 + assert await router.get_model_group_usage("gpt-5-mini") == (0, 1) + + chunks = [chunk async for chunk in stream] + total_tokens = chunks[-1].usage.total_tokens + assert total_tokens > 0 + + await asyncio.sleep(0.5) + assert await router.get_model_group_usage("gpt-5-mini") == (total_tokens, 1) + + +@pytest.mark.asyncio +async def test_deployment_callback_on_success_adds_only_uncounted_tokens(): + import time + + router = _rpm_tpm_router("lit-3058-callback") + standard_logging_payload = create_standard_logging_payload() + standard_logging_payload["total_tokens"] = 100 + kwargs = { + "litellm_params": { + "metadata": { + "deployment": "gpt-5-mini", + "model_group": "gpt-5-mini", + ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: 60, + }, + "model_info": {"id": "lit-3058-callback"}, + }, + "standard_logging_object": standard_logging_payload, + } + + tpm_key = await router.deployment_callback_on_success( + kwargs=kwargs, + completion_response=litellm.ModelResponse(model="gpt-5-mini", usage={"total_tokens": 100}), + start_time=time.time(), + end_time=time.time(), + ) + + assert tpm_key is not None + assert await router.get_model_group_usage("gpt-5-mini") == (40, 0) + + +class _GatedIncrementCache(DualCache): + def __init__(self) -> None: + super().__init__(in_memory_cache=InMemoryCache()) + self.first_increment_started = asyncio.Event() + self.release_first_increment = asyncio.Event() + self.increment_calls = 0 + + async def async_increment_cache_pipeline( + self, + increment_list: list[RedisPipelineIncrementOperation], + local_only: bool = False, + parent_otel_span: object = None, + **kwargs: object, + ) -> list[float] | None: + self.increment_calls += 1 + if self.increment_calls == 1: + self.first_increment_started.set() + await self.release_first_increment.wait() + return await super().async_increment_cache_pipeline( + increment_list, local_only=local_only, parent_otel_span=parent_otel_span, **kwargs + ) + + +@pytest.mark.asyncio +async def test_success_callback_running_during_pre_header_increment_does_not_double_count(): + router = _rpm_tpm_router("lit-3058-race") + cache = _GatedIncrementCache() + router.cache = cache + + request = asyncio.ensure_future( + router.acompletion(model="gpt-5-mini", messages=[{"role": "user", "content": "hi"}], mock_response="pong") + ) + await asyncio.wait_for(cache.first_increment_started.wait(), timeout=5) + for _ in range(50): + if get_deployment_successes_for_current_minute(router, "lit-3058-race") == 1: + break + await asyncio.sleep(0.1) + assert get_deployment_successes_for_current_minute(router, "lit-3058-race") == 1 + assert cache.increment_calls == 1 + + cache.release_first_increment.set() + response = await request + + assert await router.get_model_group_usage("gpt-5-mini") == (response.usage.total_tokens, 1) + + +class _UnavailableIncrementCache(DualCache): + def __init__(self) -> None: + super().__init__(in_memory_cache=InMemoryCache()) + self.first_increment_started = asyncio.Event() + self.release_first_increment = asyncio.Event() + self.increment_calls = 0 + + async def async_increment_cache_pipeline( + self, + increment_list: list[RedisPipelineIncrementOperation], + local_only: bool = False, + parent_otel_span: object = None, + **kwargs: object, + ) -> list[float] | None: + self.increment_calls += 1 + if self.increment_calls == 1: + self.first_increment_started.set() + await self.release_first_increment.wait() + raise RuntimeError("cache unavailable") + + +@pytest.mark.asyncio +async def test_callback_observing_stamp_before_pre_header_increment_fails_leaves_no_stamp_behind(): + router = _rpm_tpm_router("lit-3058-fail") + cache = _UnavailableIncrementCache() + router.cache = cache + metadata: dict[str, object] = {} + + request = asyncio.ensure_future( + router.acompletion( + model="gpt-5-mini", messages=[{"role": "user", "content": "hi"}], mock_response="pong", metadata=metadata + ) + ) + await asyncio.wait_for(cache.first_increment_started.wait(), timeout=5) + assert metadata[ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY] == 30 + for _ in range(50): + if get_deployment_successes_for_current_minute(router, "lit-3058-fail") == 1: + break + await asyncio.sleep(0.1) + assert get_deployment_successes_for_current_minute(router, "lit-3058-fail") == 1 + assert cache.increment_calls == 1 + + cache.release_first_increment.set() + response = await request + + assert response.usage.total_tokens == 30 + assert ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY not in metadata + assert _ratelimit_headers(response)["x-ratelimit-remaining-requests"] == 100 + assert await router.get_model_group_usage("gpt-5-mini") == (None, None) + + +@pytest.mark.asyncio +async def test_increment_deployment_usage_for_response_skips_session_wrappers(): + router = _rpm_tpm_router("lit-3058-ws") + request_kwargs = { + "model": "gpt-5-mini", + "litellm_metadata": {"model_group": "gpt-5-mini", "model_info": {"id": "lit-3058-ws"}}, + } + + await router.increment_deployment_usage_for_response(response=None, request_kwargs=request_kwargs) + + assert await router.get_model_group_usage("gpt-5-mini") == (None, None) + assert ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY not in request_kwargs["litellm_metadata"] + + +@pytest.mark.asyncio +async def test_increment_deployment_usage_writes_only_positive_deltas_for_limited_deployments(): + router = _rpm_tpm_router("lit-3058-delta") + unlimited = Router( + model_list=[ + { + "model_name": "gpt-5-mini", + "litellm_params": {"model": "gpt-5-mini", "api_key": "sk-fake"}, + "model_info": {"id": "lit-3058-unlimited"}, + } + ] + ) + + tpm_key = await router._increment_deployment_usage( + deployment_id="lit-3058-delta", + deployment_name="gpt-5-mini", + model_group="gpt-5-mini", + total_tokens=25, + rpm_increment=1, + parent_otel_span=None, + ) + assert tpm_key is not None + assert await router.get_model_group_usage("gpt-5-mini") == (25, 1) + + assert ( + await router._increment_deployment_usage( + deployment_id="lit-3058-delta", + deployment_name="gpt-5-mini", + model_group="gpt-5-mini", + total_tokens=0, + rpm_increment=0, + parent_otel_span=None, + ) + is None + ) + assert await router.get_model_group_usage("gpt-5-mini") == (25, 1) + + assert ( + await unlimited._increment_deployment_usage( + deployment_id="lit-3058-unlimited", + deployment_name="gpt-5-mini", + model_group="gpt-5-mini", + total_tokens=25, + rpm_increment=1, + parent_otel_span=None, + ) + is None + ) + assert await unlimited.get_model_group_usage("gpt-5-mini") == (None, None) + + +def _shared_redis_stub(store: dict) -> MagicMock: + from litellm.caching.redis_cache import RedisCache + + async def increment_pipeline(increment_list, **kwargs): + for op in increment_list: + store[op["key"]] = store.get(op["key"], 0.0) + op["increment_value"] + return [store[op["key"]] for op in increment_list] + + async def batch_get(keys, **kwargs): + return {key: store.get(key) for key in keys} + + redis_stub = MagicMock(spec=RedisCache) + redis_stub.async_increment_pipeline = increment_pipeline + redis_stub.async_batch_get_cache = batch_get + return redis_stub + + +@pytest.mark.asyncio +async def test_headers_on_fresh_worker_reflect_shared_redis_usage(): + store: dict = {} + worker_a = _rpm_tpm_router("lit-3058-workers") + worker_b = _rpm_tpm_router("lit-3058-workers") + worker_a.cache = DualCache(redis_cache=_shared_redis_stub(store), in_memory_cache=InMemoryCache()) + worker_b.cache = DualCache(redis_cache=_shared_redis_stub(store), in_memory_cache=InMemoryCache()) + + messages = [{"role": "user", "content": "hi"}] + tokens_on_a = 0 + for _ in range(3): + response = await worker_a.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong") + tokens_on_a += response.usage.total_tokens + + response = await worker_b.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong") + headers = _ratelimit_headers(response) + assert headers["x-ratelimit-remaining-requests"] == 96 + assert headers["x-ratelimit-remaining-tokens"] == 1000 - tokens_on_a - response.usage.total_tokens + + counted_tokens = tokens_on_a + response.usage.total_tokens + for _ in range(2): + response = await worker_a.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong") + counted_tokens += response.usage.total_tokens + + stream = await worker_b.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong", stream=True) + stream_headers = _ratelimit_headers(stream) + assert stream_headers["x-ratelimit-remaining-requests"] == 93 + assert stream_headers["x-ratelimit-remaining-tokens"] == 1000 - counted_tokens + assert [chunk async for chunk in stream] + + @pytest.mark.asyncio async def test_get_model_group_io_token_usage_sums_across_deployments(): """ @@ -1154,8 +1432,8 @@ async def test_set_response_headers_native_input_token_header_does_not_suppress_ await router.set_response_headers(response=resp, model_group="gpt-3.5-turbo") headers = resp._hidden_params["additional_headers"] - assert headers["x-ratelimit-remaining-tokens"] == 958 - assert headers["x-ratelimit-remaining-requests"] == 99 + assert headers["x-ratelimit-remaining-tokens"] == 1000 + assert headers["x-ratelimit-remaining-requests"] == 100 # the provider's native header is left untouched assert headers["x-ratelimit-remaining-input-tokens"] == 5 @@ -1187,7 +1465,7 @@ async def test_set_response_headers_native_token_header_does_not_suppress_io_hea headers = resp._hidden_params["additional_headers"] assert headers["x-ratelimit-remaining-tokens"] == 5 - assert headers["x-ratelimit-remaining-requests"] == 99 + assert headers["x-ratelimit-remaining-requests"] == 100 assert headers["x-ratelimit-remaining-input-tokens"] == 900 assert headers["x-ratelimit-remaining-output-tokens"] == 450 @@ -1196,8 +1474,7 @@ async def test_set_response_headers_native_token_header_does_not_suppress_io_hea async def test_set_response_headers_handles_missing_usage(model_list): """ Streaming chunks and some response shapes may lack a `usage` attribute or - populated `total_tokens`. The in-flight subtraction must default to 0 - tokens (still subtract 1 from requests) and never raise. + populated `total_tokens`. Header composition must not depend on usage and never raise. """ from pydantic import BaseModel @@ -1218,7 +1495,7 @@ async def test_set_response_headers_handles_missing_usage(model_list): headers = resp._hidden_params["additional_headers"] assert headers["x-ratelimit-remaining-tokens"] == 1000 - assert headers["x-ratelimit-remaining-requests"] == 99 + assert headers["x-ratelimit-remaining-requests"] == 100 @pytest.mark.asyncio diff --git a/tests/test_litellm/integrations/test_opentelemetry.py b/tests/test_litellm/integrations/test_opentelemetry.py index 7812590b3e7..7d1faf93116 100644 --- a/tests/test_litellm/integrations/test_opentelemetry.py +++ b/tests/test_litellm/integrations/test_opentelemetry.py @@ -15,10 +15,11 @@ from unittest.mock import MagicMock, patch # Adds the grandparent directory to sys.path to allow importing project modules from opentelemetry import trace +from opentelemetry.sdk._logs import LogData from opentelemetry.sdk._logs import LoggerProvider as OTLoggerProvider from opentelemetry.sdk._logs.export import InMemoryLogExporter, SimpleLogRecordProcessor from opentelemetry.sdk.metrics import MeterProvider -from opentelemetry.sdk.metrics.export import InMemoryMetricReader +from opentelemetry.sdk.metrics.export import InMemoryMetricReader, MetricsData from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter @@ -5921,13 +5922,11 @@ class TestOpenTelemetryMetricAttributeFiltering(unittest.TestCase): } ) - def test_no_filter_returns_attrs_object_unchanged(self): - """The no-config path is a hot-path no-op: it returns the same dict - object, so default emission pays zero copy cost. Locking identity makes - a future refactor that always copies/filters trip here.""" + def test_no_filter_keeps_every_attribute(self): + """The no-config path drops nothing: every attribute the caller set reaches the meter.""" otel = OpenTelemetry(config=OpenTelemetryConfig(exporter="console")) attrs = {"gen_ai.request.model": "m", "hidden_params": "{}"} - self.assertIs(otel._filter_metric_attributes(attrs), attrs) + self.assertEqual(otel._filter_metric_attributes(attrs), attrs) def test_token_type_discriminator_rejected_from_either_list(self): """gen_ai.token.type is a structural discriminator stamped onto the @@ -6068,6 +6067,118 @@ class TestOTELServiceTierAttributes(unittest.TestCase): self.assertEqual(attributes[self.RESPONSE_KEY], "tier-added-by-provider-later") +class TestOpenTelemetryProviderlessCallAttributes(unittest.TestCase): + """Regression for the OTLP exporter rejecting a None gen_ai.system or gen_ai.request.model + attribute on every export cycle.""" + + HERE = os.path.dirname(__file__) + POLL_INTERVAL = 0.05 + POLL_TIMEOUT = 2.0 + + def _providerless_kwargs(self) -> tuple[dict[str, object], dict[str, object]]: + with open(os.path.join(self.HERE, "open_telemetry", "data", "captured_kwargs.json")) as f: + kwargs = json.load(f) + with open(os.path.join(self.HERE, "open_telemetry", "data", "captured_response.json")) as f: + response_obj = json.load(f) + kwargs["litellm_params"]["custom_llm_provider"] = None + return kwargs, response_obj + + def _modelless_kwargs(self) -> tuple[dict[str, object], dict[str, object]]: + kwargs, response_obj = self._providerless_kwargs() + kwargs["model"] = None + return kwargs, response_obj + + def _recorded_metrics(self, kwargs: dict[str, object], response_obj: dict[str, object]) -> MetricsData | None: + metric_reader = InMemoryMetricReader() + meter_provider = MeterProvider(metric_readers=[metric_reader]) + tracer_provider = TracerProvider() + tracer_provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter())) + otel = OpenTelemetry( + config=OpenTelemetryConfig(exporter="console", enable_metrics=True), + tracer_provider=tracer_provider, + meter_provider=meter_provider, + ) + otel.tracer = tracer_provider.get_tracer(__name__) + + start = datetime.utcnow() + otel._handle_success(kwargs, response_obj, start, start + timedelta(seconds=1)) + + deadline = time.time() + self.POLL_TIMEOUT + while time.time() < deadline: + data = metric_reader.get_metrics_data() + if data and getattr(data, "resource_metrics", None): + return data + time.sleep(self.POLL_INTERVAL) + return None + + def _emitted_log_records(self, semconv_opt_in: str) -> tuple[LogData, ...]: + log_exporter = InMemoryLogExporter() + logger_provider = OTLoggerProvider() + logger_provider.add_log_record_processor(SimpleLogRecordProcessor(log_exporter)) + with patch.dict(os.environ, {"OTEL_SEMCONV_STABILITY_OPT_IN": semconv_opt_in}): + handler = OpenTelemetry( + config=OpenTelemetryConfig(exporter="console", enable_events=True), + logger_provider=logger_provider, + ) + handler.message_logging = True + + kwargs, response_obj = self._providerless_kwargs() + span = handler.tracer.start_span("test") + with self.assertNoLogs("opentelemetry.attributes", level="WARNING"): + handler._emit_semantic_logs(kwargs, response_obj, span) + span.end() + handler._logger_provider.force_flush(2000) + return log_exporter.get_finished_logs() + + def _assert_every_attribute_encodes(self, attrs: dict[str, object]) -> None: + from opentelemetry.exporter.otlp.proto.common._internal import _encode_attributes + + self.assertEqual(len(_encode_attributes(attrs) or []), len(attrs)) + + def _recorded_data_points(self, kwargs: dict[str, object], response_obj: dict[str, object]) -> list[object]: + data = self._recorded_metrics(kwargs, response_obj) + self.assertIsNotNone(data, "no metrics were recorded") + data_points = [ + dp + for rm in data.resource_metrics + for sm in rm.scope_metrics + for m in sm.metrics + for dp in m.data.data_points + ] + self.assertTrue(data_points, "no metric data points were recorded") + return data_points + + def test_metrics_are_encodable_and_carry_no_provider_label(self): + kwargs, response_obj = self._providerless_kwargs() + for dp in self._recorded_data_points(kwargs, response_obj): + self.assertNotIn("gen_ai.system", dp.attributes) + self.assertEqual(dp.attributes["gen_ai.request.model"], kwargs["model"]) + self._assert_every_attribute_encodes(dict(dp.attributes)) + + def test_metrics_are_encodable_and_carry_no_model_label_when_the_call_has_none(self): + for dp in self._recorded_data_points(*self._modelless_kwargs()): + self.assertNotIn("gen_ai.request.model", dp.attributes) + self._assert_every_attribute_encodes(dict(dp.attributes)) + + def test_legacy_content_events_are_encodable_and_carry_no_provider_label(self): + logs = self._emitted_log_records("") + self.assertTrue(logs, "no content events were emitted") + for log in logs: + attrs = dict(log.log_record.attributes or {}) + self.assertNotIn("gen_ai.system", attrs) + self.assertNotIn(None, attrs.values()) + self._assert_every_attribute_encodes(attrs) + + def test_inference_details_event_is_encodable_and_carries_no_provider_label(self): + logs = self._emitted_log_records("gen_ai_latest_experimental") + self.assertEqual(len(logs), 1) + attrs = dict(logs[0].log_record.attributes or {}) + self.assertEqual(attrs["event_name"], "gen_ai.client.inference.operation.details") + self.assertNotIn("gen_ai.provider.name", attrs) + self.assertNotIn(None, attrs.values()) + self._assert_every_attribute_encodes(attrs) + + class TestDynamicTracerProviderCache(unittest.TestCase): """Every credential-scoped TracerProvider that owns its exporter also owns a BatchSpanProcessor worker thread that only stops on shutdown, so the cache holding them diff --git a/tests/test_litellm/integrations/test_s3.py b/tests/test_litellm/integrations/test_s3.py index 58b15b79e76..fd677b9dfdf 100644 --- a/tests/test_litellm/integrations/test_s3.py +++ b/tests/test_litellm/integrations/test_s3.py @@ -1,16 +1,24 @@ +import copy +import json from datetime import datetime from unittest.mock import MagicMock, patch +import pytest + import litellm from litellm.constants import MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES, MAX_S3_OBJECT_KEY_BYTES -from litellm.integrations.s3 import S3Logger +from litellm.integrations.s3 import S3Logger, prompts_only_payload, resolve_s3_log_prompts_only TEST_KMS_KEY_ARN = "arn:aws:kms:us-east-1:111122223333:key/test-key-id" +TEST_MESSAGES = [{"role": "user", "content": "Reply with exactly the word PINEAPPLE."}] +TEST_RESPONSE = {"choices": [{"message": {"role": "assistant", "content": "PINEAPPLE"}}]} def _standard_logging_payload(response_id: str = "chatcmpl-test-id") -> dict: return { "id": response_id, + "messages": copy.deepcopy(TEST_MESSAGES), + "response": copy.deepcopy(TEST_RESPONSE), "metadata": {"user_api_key_team_alias": None}, } @@ -22,7 +30,9 @@ def _log_event_kwargs(response_id: str = "chatcmpl-test-id") -> dict: } -def _run_log_event(callback_params: dict, response_id: str = "chatcmpl-test-id") -> MagicMock: +def _run_log_event( + callback_params: dict, response_id: str = "chatcmpl-test-id", log_kwargs: dict[str, object] | None = None +) -> MagicMock: original = litellm.s3_callback_params litellm.s3_callback_params = callback_params try: @@ -31,7 +41,7 @@ def _run_log_event(callback_params: dict, response_id: str = "chatcmpl-test-id") mock_boto3_client.return_value = mock_s3_client logger = S3Logger() logger.log_event( - kwargs=_log_event_kwargs(response_id), + kwargs=_log_event_kwargs(response_id) if log_kwargs is None else log_kwargs, response_obj={"id": response_id}, start_time=datetime(2026, 7, 30, 12, 0, 0), end_time=datetime(2026, 7, 30, 12, 0, 1), @@ -182,3 +192,123 @@ def test_put_object_keeps_the_configured_path_intact_when_only_the_id_has_to_shr key = mock_s3_client.put_object.call_args.kwargs["Key"] assert key.startswith(long_path + "/2026-07-30/") assert len(key.encode("utf-8")) == MAX_S3_OBJECT_KEY_BYTES + + +def _uploaded_body(mock_s3_client: MagicMock) -> dict[str, object]: + return json.loads(mock_s3_client.put_object.call_args.kwargs["Body"]) + + +def test_log_event_prompts_only_drops_response_and_keeps_messages(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("S3_LOG_PROMPTS_ONLY", raising=False) + log_kwargs = _log_event_kwargs() + original_payload = copy.deepcopy(log_kwargs["standard_logging_object"]) + + mock_s3_client = _run_log_event( + {"s3_bucket_name": "test-bucket", "s3_region_name": "us-east-1", "s3_log_prompts_only": True}, + log_kwargs=log_kwargs, + ) + + body = _uploaded_body(mock_s3_client) + assert body["messages"] == TEST_MESSAGES + assert body["response"] is None + assert body["id"] == "chatcmpl-test-id" + assert log_kwargs["standard_logging_object"] == original_payload + + +def test_log_event_default_keeps_response(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("S3_LOG_PROMPTS_ONLY", raising=False) + + mock_s3_client = _run_log_event({"s3_bucket_name": "test-bucket", "s3_region_name": "us-east-1"}) + + body = _uploaded_body(mock_s3_client) + assert body["response"] == TEST_RESPONSE + assert body["messages"] == TEST_MESSAGES + + +def test_log_event_reads_prompts_only_env_var_at_log_time(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("S3_LOG_PROMPTS_ONLY", raising=False) + original = litellm.s3_callback_params + litellm.s3_callback_params = {"s3_bucket_name": "test-bucket", "s3_region_name": "us-east-1"} + try: + with patch("boto3.client") as mock_boto3_client: + mock_s3_client = MagicMock() + mock_boto3_client.return_value = mock_s3_client + logger = S3Logger() + monkeypatch.setenv("S3_LOG_PROMPTS_ONLY", "true") + logger.log_event( + kwargs=_log_event_kwargs(), + response_obj={"id": "chatcmpl-test-id"}, + start_time=datetime(2026, 7, 30, 12, 0, 0), + end_time=datetime(2026, 7, 30, 12, 0, 1), + print_verbose=lambda *args, **kwargs: None, + ) + finally: + litellm.s3_callback_params = original + + body = _uploaded_body(mock_s3_client) + assert body["response"] is None + assert body["messages"] == TEST_MESSAGES + + +def test_log_event_explicit_false_param_beats_env_var(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("S3_LOG_PROMPTS_ONLY", "true") + + mock_s3_client = _run_log_event( + {"s3_bucket_name": "test-bucket", "s3_region_name": "us-east-1", "s3_log_prompts_only": False} + ) + + assert _uploaded_body(mock_s3_client)["response"] == TEST_RESPONSE + + +def test_s3_logger_init_does_not_mutate_global_callback_params(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("MY_S3_BUCKET", "resolved-bucket") + callback_params = {"s3_bucket_name": "os.environ/MY_S3_BUCKET", "s3_region_name": "us-east-1"} + snapshot = copy.deepcopy(callback_params) + original = litellm.s3_callback_params + litellm.s3_callback_params = callback_params + try: + with patch("boto3.client"): + logger = S3Logger() + finally: + litellm.s3_callback_params = original + + assert logger.bucket_name == "resolved-bucket" + assert callback_params == snapshot + + +@pytest.mark.parametrize( + "configured,env_value,expected", + [ + (True, None, True), + (False, "true", False), + ("true", None, True), + ("False", "true", False), + ("1", None, True), + ("0", None, False), + (" yes ", None, True), + (None, None, False), + (None, "true", True), + (None, "false", False), + (None, "", False), + ("", "true", False), + ], +) +def test_resolve_s3_log_prompts_only(configured: object, env_value: str | None, expected: bool): + environ = {} if env_value is None else {"S3_LOG_PROMPTS_ONLY": env_value} + assert resolve_s3_log_prompts_only(configured, environ) is expected + + +def test_resolve_s3_log_prompts_only_unparseable_value_fails_toward_prompts_only(): + assert resolve_s3_log_prompts_only("enabled", {}) is True + + +def test_prompts_only_payload_returns_copy_with_response_cleared(): + payload = _standard_logging_payload() + snapshot = copy.deepcopy(payload) + + stripped = prompts_only_payload(payload) + + assert stripped["response"] is None + assert stripped["messages"] == TEST_MESSAGES + assert stripped is not payload + assert payload == snapshot diff --git a/tests/test_litellm/integrations/test_s3_v2.py b/tests/test_litellm/integrations/test_s3_v2.py index 08d37297ab1..52fbbe40b0e 100644 --- a/tests/test_litellm/integrations/test_s3_v2.py +++ b/tests/test_litellm/integrations/test_s3_v2.py @@ -1,8 +1,11 @@ import asyncio +import copy +import json import re import sys import textwrap import uuid +from collections.abc import Awaitable, Callable from contextlib import asynccontextmanager from datetime import datetime from pathlib import Path @@ -10,6 +13,7 @@ from unittest.mock import AsyncMock, MagicMock, call, patch import httpx import pytest +import respx from litellm.integrations.s3_v2 import S3Logger from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler @@ -2310,3 +2314,137 @@ def _s3_logger_for_region(region_name: str) -> S3Logger: ) def test_build_object_url_uses_partition_dns_suffix(region_name: str, expected_url: str) -> None: assert _s3_logger_for_region(region_name)._build_object_url("2025-01-01/key.json") == expected_url + + +def _prompts_only_logger(s3_log_prompts_only: bool | None = None) -> S3Logger: + return S3Logger( + s3_bucket_name="test-bucket", + s3_aws_access_key_id="test-key", + s3_aws_secret_access_key="test-secret", + s3_region_name="us-east-1", + s3_log_prompts_only=s3_log_prompts_only, + ) + + +def _chat_payload() -> StandardLoggingPayload: + return StandardLoggingPayload( + id="chatcmpl-prompts-only", + messages=[{"role": "user", "content": "Reply with exactly the word PINEAPPLE."}], + response={"choices": [{"message": {"role": "assistant", "content": "PINEAPPLE"}}]}, + metadata={"user_api_key_team_alias": None}, + ) + + +async def _queued_body_via_async_upload( + logger: S3Logger, log_event: Callable[..., Awaitable[None]] +) -> dict[str, object]: + payload = _chat_payload() + original = copy.deepcopy(payload) + await log_event( + kwargs={"standard_logging_object": payload}, + response_obj=None, + start_time=datetime(2026, 7, 30, 12, 0, 0), + end_time=datetime(2026, 7, 30, 12, 0, 1), + ) + assert payload == original, "the caller's standard_logging_object must not be mutated" + (element,) = logger.log_queue + + response = MagicMock() + response.status_code = 200 + response.raise_for_status = MagicMock() + logger.async_httpx_client = AsyncMock() + logger.async_httpx_client.put.return_value = response + await logger.async_upload_data_to_s3(element) + return json.loads(logger.async_httpx_client.put.call_args.kwargs["data"]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("event_name", ["async_log_success_event", "async_log_failure_event"]) +async def test_prompts_only_drops_response_but_keeps_messages_in_uploaded_object( + monkeypatch: pytest.MonkeyPatch, event_name: str +): + import litellm + + monkeypatch.setattr(litellm, "s3_callback_params", {"s3_log_prompts_only": True}) + logger = _prompts_only_logger() + + log_event: Callable[..., Awaitable[None]] = ( + logger.async_log_success_event if event_name == "async_log_success_event" else logger.async_log_failure_event + ) + body = await _queued_body_via_async_upload(logger, log_event) + + assert body["messages"] == _chat_payload()["messages"] + assert body["response"] is None + assert body["id"] == "chatcmpl-prompts-only" + + +@pytest.mark.asyncio +async def test_prompts_only_default_off_keeps_response_in_uploaded_object(monkeypatch: pytest.MonkeyPatch): + import litellm + + monkeypatch.setattr(litellm, "s3_callback_params", {}) + monkeypatch.delenv("S3_LOG_PROMPTS_ONLY", raising=False) + logger = _prompts_only_logger() + + body = await _queued_body_via_async_upload(logger, logger.async_log_success_event) + + assert body["response"] == _chat_payload()["response"] + assert body["messages"] == _chat_payload()["messages"] + + +@pytest.mark.asyncio +async def test_prompts_only_explicit_false_in_params_beats_env_var(monkeypatch: pytest.MonkeyPatch): + import litellm + + monkeypatch.setattr(litellm, "s3_callback_params", {"s3_log_prompts_only": False}) + monkeypatch.setenv("S3_LOG_PROMPTS_ONLY", "true") + logger = _prompts_only_logger() + + body = await _queued_body_via_async_upload(logger, logger.async_log_success_event) + + assert body["response"] == _chat_payload()["response"] + + +@pytest.mark.asyncio +async def test_prompts_only_env_var_applies_when_param_unset(monkeypatch: pytest.MonkeyPatch): + import litellm + + monkeypatch.setattr(litellm, "s3_callback_params", {}) + logger = _prompts_only_logger() + monkeypatch.setenv("S3_LOG_PROMPTS_ONLY", "true") + + body = await _queued_body_via_async_upload(logger, logger.async_log_success_event) + + assert body["response"] is None + assert body["messages"] == _chat_payload()["messages"] + + +@respx.mock +def test_prompts_only_constructor_kwarg_applies_to_sync_upload(monkeypatch: pytest.MonkeyPatch): + import litellm + + monkeypatch.setattr(litellm, "s3_callback_params", {}) + monkeypatch.delenv("S3_LOG_PROMPTS_ONLY", raising=False) + logger = _prompts_only_logger(s3_log_prompts_only=True) + payload = _chat_payload() + + element = logger.create_s3_batch_logging_element( + start_time=datetime(2026, 7, 30, 12, 0, 0), + standard_logging_payload=payload, + ) + assert element is not None + assert payload["response"] == _chat_payload()["response"] + + put_route = respx.put(url__regex=r"https://test-bucket\.s3\..*").mock(return_value=httpx.Response(200)) + logger.upload_data_to_s3(element) + + body = json.loads(put_route.calls.last.request.content) + assert body["response"] is None + assert body["messages"] == _chat_payload()["messages"] + + +@pytest.mark.parametrize("callback_name", ["s3", "s3_v2"]) +def test_prompts_only_toggle_is_exposed_to_admin_ui_for_both_s3_callbacks(callback_name: str): + from litellm.integrations.custom_logger import CustomLogger + + assert "S3_LOG_PROMPTS_ONLY" in CustomLogger.get_callback_env_vars(callback_name) diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 7522e9a62e5..cb18a5f192c 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -2482,7 +2482,10 @@ class TestAnthropicMessagesHandlerStreamingScanKey: open_key = handler.get_streaming_scan_key([self._text_delta("hi"), tool_use]) ended_key = handler.get_streaming_scan_key([self._text_delta("hi"), tool_use, self._stop("tool_use")]) assert open_key == StreamingScanKey(texts=("hi",)) + assert open_key.tool_calls_in_flight is True + assert handler.get_streaming_scan_key([self._text_delta("hi")]).tool_calls_in_flight is False assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0] + assert ended_key.tool_calls_in_flight is False assert ended_key != open_key diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py index c3400dc40c3..a241fc03d57 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py @@ -2719,3 +2719,74 @@ class TestRustChatCompletionsHook: "model": "m", "messages": [], } + + +def _served_model_stream_chunks(model: str | None) -> list[dict[str, object]]: + return [ + { + "type": "message_start", + "message": { + "id": "msg_served", + "type": "message", + "role": "assistant", + "content": [], + "usage": {"input_tokens": 10, "output_tokens": 1}, + **({"model": model} if model is not None else {}), + }, + }, + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": ""}, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "Hello"}, + }, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": {"output_tokens": 2}, + }, + {"type": "message_stop"}, + ] + + +def test_message_start_model_is_carried_on_stream_chunks(): + iterator: Final = ModelResponseIterator(None, sync_stream=True) + + parsed: Final = [iterator.chunk_parser(chunk) for chunk in _served_model_stream_chunks("claude-served-1")] + + assert all(chunk.model == "claude-served-1" for chunk in parsed) + + +def test_message_start_without_model_leaves_chunk_model_unset(): + iterator: Final = ModelResponseIterator(None, sync_stream=True) + + parsed: Final = [iterator.chunk_parser(chunk) for chunk in _served_model_stream_chunks(None)] + + assert all(chunk.model is None for chunk in parsed) + + +def test_served_model_reaches_assembled_stream_through_custom_stream_wrapper(): + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + + served_model: Final = "claude-served-1" + sse_lines: Final = [f"data: {json.dumps(chunk)}\n".encode() for chunk in _served_model_stream_chunks(served_model)] + iterator: Final = ModelResponseIterator(iter(sse_lines), sync_stream=True) + wrapper: Final = CustomStreamWrapper( + completion_stream=iter(iterator), + model="anthropic/claude-requested", + custom_llm_provider="anthropic", + logging_obj=MagicMock(), + ) + + chunks: Final = list(wrapper) + + assert len(chunks) > 1 + for chunk in chunks[1:]: + assert chunk._hidden_params["provider_response_model"] == served_model + assembled: Final = litellm.stream_chunk_builder(chunks=list(chunks), messages=[{"role": "user", "content": "hi"}]) + assert assembled._hidden_params["provider_response_model"] == served_model diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py index 7660a8649b5..fc5d807bc23 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py @@ -1200,6 +1200,7 @@ def _fake_user_api_key_auth( team_models=None, team_id=None, model_max_budget=None, + team_model_max_budget=None, end_user_model_max_budget=None, end_user_id=None, user_model_max_budget=None, @@ -1220,6 +1221,7 @@ def _fake_user_api_key_auth( auth.team_id = team_id auth.team_model_aliases = None auth.model_max_budget = model_max_budget + auth.team_model_max_budget = team_model_max_budget auth.end_user_model_max_budget = end_user_model_max_budget auth.end_user_id = end_user_id auth.user_model_max_budget = user_model_max_budget @@ -1860,6 +1862,78 @@ async def test_summary_model_rate_limit_skipped_for_legacy_limiter(): assert not result.applied_edits[0].get("error") +async def test_summary_model_denied_when_team_over_model_budget(): + """The team per-model budget gates the summary subrequest, whose spend is + charged to the team counter via the propagated `user_api_key_team_model_max_budget`. + The key's own `model_max_budget` is handed to the limiter so a key-level + override keeps taking precedence over the team cap here as it does in auth.""" + import litellm + + messages = _simple_messages() + mock_call = AsyncMock(return_value=_make_mock_response("x")) + key_budget = {"claude-opus-4-8": {"budget_limit": 1}} + team_budget = {"claude-haiku-4-5": {"budget_limit": 5, "time_period": "1d"}} + + auth = _fake_user_api_key_auth( + key_models=["all-proxy-models"], + model_max_budget=key_budget, + team_model_max_budget=team_budget, + team_id="team-over-budget", + token="hashed-token", + ) + + limiter = MagicMock() + limiter.is_key_within_model_budget = AsyncMock(return_value=True) + limiter.is_team_within_model_budget = AsyncMock( + side_effect=litellm.BudgetExceededError( + message="over budget", current_cost=10, max_budget=5 + ) + ) + + with ( + patch( # test-quality-ok: apply_compact_20260112 reads the summary model setting as a module global, no seam + "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting", + return_value="claude-haiku-4-5", + ), + patch("litellm.token_counter", return_value=200_000), # test-quality-ok: forces the over-threshold branch + patch( # test-quality-ok: the summary call is the observable that must NOT happen when the team is over budget + "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model", + mock_call, + ), + patch( # test-quality-ok: the limiter is a proxy_server module global the editor imports, no injection seam + "litellm.proxy.proxy_server.model_max_budget_limiter", limiter + ), + ): + result = await apply_compact_20260112( + model=MODEL, + messages=messages, + tools=None, + system=None, + edit_spec=_EDIT_SPEC_DEFAULT, + user_api_key_auth=auth, + ) + + mock_call.assert_not_awaited() + assert result.applied_edits[0].get("error") == "summary_model_budget_exceeded" + limiter.is_team_within_model_budget.assert_awaited_once_with( + team_id="team-over-budget", + team_model_max_budget=team_budget, + key_model_max_budget=key_budget, + model="claude-haiku-4-5", + ) + import inspect + + from litellm.proxy.hooks.model_max_budget_limiter import ( + _PROXY_VirtualKeyModelMaxBudgetLimiter, + ) + + real_params = inspect.signature( + _PROXY_VirtualKeyModelMaxBudgetLimiter.is_team_within_model_budget + ).parameters + for kwarg in ("team_id", "team_model_max_budget", "key_model_max_budget", "model"): + assert kwarg in real_params, f"compact.py passes {kwarg}=, which the limiter does not accept" + + async def test_scoped_budget_metadata_propagated_to_summary_call(): """The end-user/project scope identifiers and the end-user budget the post-call spend and rate-limit hooks key on are forwarded to the summary subrequest, and diff --git a/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py b/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py index 3f54b695fef..5795e29a8bc 100644 --- a/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py +++ b/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py @@ -5,12 +5,15 @@ from typing import NamedTuple import pytest import litellm +from litellm.cost_calculator import completion_cost from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig from litellm.llms.bedrock.common_utils import BedrockModelInfo from litellm.types.utils import ( Choices, Message, ModelResponse, + PromptTokensDetailsWrapper, + Usage, ) @@ -152,3 +155,55 @@ def test_bedrock_gpt_5_6_offers_tools_and_reasoning_effort_but_not_thinking(prof assert "reasoning_effort" in supported assert "thinking" not in supported assert "output_config" not in supported + + +# Cache-read prices are the `*-cache-read-input-tokens` usagetype rows of the AWS Price List API, us-east-1, +# https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrock/current/us-east-1/index.json on 2026-09-15 +@pytest.mark.parametrize( + "model,expected_cache_read", + [ + ("amazon.nova-lite-v1:0", 1.5e-8), + ("us.amazon.nova-lite-v1:0", 1.5e-8), + ("amazon.nova-micro-v1:0", 8.75e-9), + ("us.amazon.nova-micro-v1:0", 8.75e-9), + ("amazon.nova-pro-v1:0", 2e-7), + ("us.amazon.nova-pro-v1:0", 2e-7), + ("us.amazon.nova-premier-v1:0", 6.25e-7), + ], +) +def test_bedrock_nova_cache_read_prices( + model, expected_cache_read, local_model_cost_map +): + model_info = litellm.model_cost[model] + assert model_info["cache_read_input_token_cost"] == expected_cache_read + usage = Usage( + prompt_tokens=1_000, + completion_tokens=100, + total_tokens=1_100, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=400), + ) + response = _bedrock_response(model, usage) + + cost = completion_cost( + completion_response=response, + model=model, + custom_llm_provider="bedrock", + ) + expected_cost = ( + 600 * model_info["input_cost_per_token"] + + 400 * expected_cache_read + + 100 * model_info["output_cost_per_token"] + ) + assert cost == pytest.approx(expected_cost) + + uncached_usage = Usage( + prompt_tokens=1_000, + completion_tokens=100, + total_tokens=1_100, + ) + uncached_cost = completion_cost( + completion_response=_bedrock_response(model, uncached_usage), + model=model, + custom_llm_provider="bedrock", + ) + assert cost < uncached_cost diff --git a/tests/test_litellm/llms/gemini/test_cost_calculator.py b/tests/test_litellm/llms/gemini/test_cost_calculator.py index 2d56757c601..5eed11dff03 100644 --- a/tests/test_litellm/llms/gemini/test_cost_calculator.py +++ b/tests/test_litellm/llms/gemini/test_cost_calculator.py @@ -1,8 +1,7 @@ -import os - import pytest import litellm +from litellm.cost_calculator import completion_cost from litellm.llms.gemini.cost_calculator import ( cost_per_google_maps_grounding_request, cost_per_web_search_request, @@ -18,6 +17,7 @@ from litellm.types.utils import ( ImageResponse, ImageUsage, ImageUsageInputTokensDetails, + ModelResponse, PromptTokensDetailsWrapper, Usage, ) @@ -452,6 +452,42 @@ def test_map_traffic_type_to_service_tier( ) +# Alias targets are the `modelVersion` returned by +# POST https://generativelanguage.googleapis.com/v1beta/models/:generateContent on 2026-09-15 +@pytest.mark.parametrize( + "alias,target", + [ + ("gemini/gemini-flash-latest", "gemini/gemini-3.8-flash"), + ("gemini/gemini-flash-lite-latest", "gemini/gemini-3.5-flash-lite"), + ("gemini/gemini-pro-latest", "gemini/gemini-3.1-pro-preview"), + ], +) +def test_latest_aliases_cost_the_same_as_their_current_target( + monkeypatch, alias, target +): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + + usage = Usage( + prompt_tokens=1_000, + completion_tokens=500, + total_tokens=1_500, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=400), + ) + + def cost_of(model: str) -> float: + return completion_cost( + completion_response=ModelResponse(model=model, usage=usage), + model=model, + custom_llm_provider="gemini", + ) + + alias_cost = cost_of(alias) + target_cost = cost_of(target) + assert alias_cost == pytest.approx(target_cost) + assert alias_cost > 0 + + @pytest.mark.parametrize( "prefixed,bare", [ diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index cb884fb7cc1..9b88161aca9 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -2206,10 +2206,35 @@ class TestStreamingScanKey: [self._chunk("hi"), tool_chunk, self._chunk(None, finish_reason="stop")] ) assert open_key == StreamingScanKey(texts=("hi",)) + assert open_key.tool_calls_in_flight is True + assert handler.get_streaming_scan_key([self._chunk("hi")]).tool_calls_in_flight is False assert ended_key.texts == ("hi",) assert len(ended_key.tool_calls) == 1 and "get_weather" in ended_key.tool_calls[0] + assert ended_key.tool_calls_in_flight is False assert ended_key != open_key + def test_legacy_function_call_delta_is_held_like_a_tool_call(self): + from litellm.types.utils import Delta, FunctionCall, ModelResponseStream, StreamingChoices + + handler = OpenAIChatCompletionsHandler() + function_chunk = ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta(content=None, function_call=FunctionCall(name="run_shell", arguments='{"cmd": "rm"}')), + finish_reason=None, + ) + ] + ) + open_key = handler.get_streaming_scan_key([self._chunk("hi"), function_chunk]) + ended_key = handler.get_streaming_scan_key( + [self._chunk("hi"), function_chunk, self._chunk(None, finish_reason="function_call")] + ) + assert open_key.tool_calls_in_flight is True + assert open_key.tool_calls == () + assert len(ended_key.tool_calls) == 1 and "run_shell" in ended_key.tool_calls[0] + assert ended_key.tool_calls_in_flight is False + def test_text_after_the_first_choice_finishes_still_changes_the_key(self): handler = OpenAIChatCompletionsHandler() first_done = [self._chunk("a", index=0), self._chunk("b", finish_reason="stop", index=0)] diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index 48d86384633..c714b5d378a 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -3211,3 +3211,42 @@ class TestOpenAIResponsesHandlerStreamingScanKey: def test_output_item_done_round_is_never_deduped(self): done = {"type": "response.output_item.done", "sequence_number": 1, "item": {"type": "function_call"}} assert OpenAIResponsesHandler().get_streaming_scan_key([self._delta(0, "hi"), done]) is None + + @pytest.mark.parametrize("terminal_type", ["response.incomplete", "response.failed"]) + def test_non_completed_terminal_envelopes_key_their_output_items(self, terminal_type): + handler = OpenAIResponsesHandler() + arguments_delta = { + "type": "response.function_call_arguments.delta", + "sequence_number": 1, + "item_id": "fc_1", + "delta": '{"city":', + } + function_call = {"type": "function_call", "call_id": "call_1", "name": "get_weather", "arguments": '{"city":'} + terminal = {"type": terminal_type, "sequence_number": 2, "response": {"id": "resp_1", "output": [function_call]}} + mid_stream_key = handler.get_streaming_scan_key([arguments_delta]) + ended_key = handler.get_streaming_scan_key([arguments_delta, terminal]) + assert ended_key.stream_ended is True + assert ended_key.tool_calls_in_flight is False + assert len(ended_key.tool_calls) == 1 + assert ended_key != mid_stream_key + + def test_streamed_tool_call_events_flag_tool_calls_in_flight_until_the_stream_ends(self): + handler = OpenAIResponsesHandler() + added = { + "type": "response.output_item.added", + "sequence_number": 1, + "item": {"type": "function_call", "id": "fc_1", "call_id": "call_1", "name": "get_weather"}, + } + arguments_delta = { + "type": "response.function_call_arguments.delta", + "sequence_number": 2, + "item_id": "fc_1", + "delta": '{"city":', + } + function_call = {"type": "function_call", "call_id": "call_1", "name": "get_weather", "arguments": "{}"} + assert handler.get_streaming_scan_key([self._delta(0, "hi")]).tool_calls_in_flight is False + assert handler.get_streaming_scan_key([self._delta(0, "hi"), added]).tool_calls_in_flight is True + assert handler.get_streaming_scan_key([self._delta(0, "hi"), arguments_delta]).tool_calls_in_flight is True + ended_key = handler.get_streaming_scan_key([self._delta(0, "hi"), added, self._completed(3, [function_call])]) + assert ended_key.tool_calls_in_flight is False + assert len(ended_key.tool_calls) == 1 diff --git a/tests/test_litellm/models/test_models.py b/tests/test_litellm/models/test_models.py index aa6449c98dd..9b803c14062 100644 --- a/tests/test_litellm/models/test_models.py +++ b/tests/test_litellm/models/test_models.py @@ -362,6 +362,14 @@ class TestVerificationToken: assert deleted.deleted_at is not None assert deleted.token == "t1" + def test_total_spend_is_carried_separately_from_resettable_spend(self): + token = LiteLLM_VerificationToken(token="t1", spend=0.0, total_spend=12.5) + assert token.model_dump()["total_spend"] == 12.5 + assert token.model_dump()["spend"] == 0.0 + + deleted = LiteLLM_DeletedVerificationToken.model_validate({**token.model_dump(), "deleted_by": "admin"}) + assert deleted.total_spend == 12.5 + class TestConfigTable: def test_config_creation(self): diff --git a/tests/test_litellm/ocr/test_ocr_file_input.py b/tests/test_litellm/ocr/test_ocr_file_input.py index 3526d8c00d6..8f82a64bd85 100644 --- a/tests/test_litellm/ocr/test_ocr_file_input.py +++ b/tests/test_litellm/ocr/test_ocr_file_input.py @@ -12,32 +12,16 @@ Tests that: import base64 import os import tempfile -from collections.abc import Generator from io import BytesIO from pathlib import Path from typing import Final -from unittest.mock import AsyncMock, MagicMock, Mock +from unittest.mock import AsyncMock, MagicMock import orjson import pytest from starlette.datastructures import FormData -from litellm.ocr.input import convert_file_document_to_url_document, get_mime_type - - -@pytest.fixture(autouse=True, params=["native", "disabled", "unavailable"]) -def document_runtime(request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch) -> Generator[None]: - from litellm.rust_bridge import bindings, configuration - - configuration.reset_rust_configuration() - monkeypatch.delenv("LITELLM_RUST", raising=False) - if request.param == "disabled": - monkeypatch.setenv("LITELLM_RUST", "0") - monkeypatch.setattr(bindings, "get_native_bridge", Mock(side_effect=AssertionError("Rust is disabled"))) - elif request.param == "unavailable": - monkeypatch.setattr(bindings, "get_native_bridge", lambda: None) - yield - configuration.reset_rust_configuration() +from litellm.ocr.legacy import convert_file_document_to_url_document, get_mime_type class TestGetMimeType: @@ -503,10 +487,9 @@ class TestProxySecurityGuard: async def test_proxy_upload_stops_reading_at_size_limit() -> None: from starlette.datastructures import UploadFile - from litellm.ocr.input import get_max_file_bytes - from litellm.proxy.ocr_endpoints.endpoints import _parse_multipart_form + from litellm.proxy.ocr_endpoints.endpoints import _MAX_FILE_BYTES, _parse_multipart_form - limit: Final = get_max_file_bytes() + limit: Final = _MAX_FILE_BYTES with tempfile.TemporaryFile() as stream: stream.truncate(limit * 2) upload: Final = UploadFile(file=stream, filename="large.pdf") diff --git a/tests/test_litellm/proxy/auth/test_fallback_budget.py b/tests/test_litellm/proxy/auth/test_fallback_budget.py new file mode 100644 index 00000000000..00c1a7cdefc --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_fallback_budget.py @@ -0,0 +1,202 @@ +import pytest + +from litellm import Router +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.fallback_budget import ( + RouterFallbackBudgetCheck, + is_token_within_budget_for_model, + router_fallback_budget_check, +) + +FREE_MODEL = { + "model_name": "free-model", + "litellm_params": { + "model": "ollama/llama2", + "api_base": "http://localhost:11434", + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + }, + "model_info": { + "id": "free-model-id", + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + }, +} + +PAID_MODEL = { + "model_name": "paid-model", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "k"}, + "model_info": {"id": "paid-model-id"}, +} + + +def _router() -> Router: + return Router(model_list=[FREE_MODEL, PAID_MODEL], fallbacks=[{"free-model": ["paid-model"]}]) + + +def _token(**overrides) -> UserAPIKeyAuth: + fields = { + "api_key": "hashed", + "token": "hashed", + "spend": 0.0, + "max_budget": None, + "user_id": "u1", + "user_spend": 0.0, + "user_max_budget": None, + } + fields.update(overrides) + return UserAPIKeyAuth(**fields) + + +ENFORCED = RouterFallbackBudgetCheck(is_enforced=lambda: True) +NOT_ENFORCED = RouterFallbackBudgetCheck(is_enforced=lambda: False) + + +@pytest.mark.asyncio +async def test_paid_target_allowed_when_under_budget(): + token = _token(spend=1.0, max_budget=50.0, user_spend=1.0, user_max_budget=50.0) + assert await is_token_within_budget_for_model(model="paid-model", valid_token=token, llm_router=_router()) is True + + +@pytest.mark.asyncio +async def test_paid_target_refused_when_over_key_budget(): + token = _token(spend=100.0, max_budget=50.0) + assert await is_token_within_budget_for_model(model="paid-model", valid_token=token, llm_router=_router()) is False + + +@pytest.mark.asyncio +async def test_paid_target_refused_when_over_user_budget(): + token = _token(user_spend=1900.0, user_max_budget=50.0) + assert await is_token_within_budget_for_model(model="paid-model", valid_token=token, llm_router=_router()) is False + + +@pytest.mark.asyncio +async def test_zero_cost_target_allowed_even_when_over_budget(): + """Refusing a free target would deny a request on spend some other model accrued.""" + token = _token(user_spend=1900.0, user_max_budget=50.0) + assert await is_token_within_budget_for_model(model="free-model", valid_token=token, llm_router=_router()) is True + + +@pytest.mark.asyncio +async def test_no_budget_configured_is_always_within_budget(): + token = _token(spend=9999.0, user_spend=9999.0) + assert await is_token_within_budget_for_model(model="paid-model", valid_token=token, llm_router=_router()) is True + + +@pytest.mark.asyncio +async def test_team_key_does_not_inherit_personal_budget_by_default(monkeypatch): + """Mirrors _PROXY_MaxBudgetLimiter: a team key ignores the owner's personal cap.""" + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "general_settings", {}, raising=False) + token = _token(team_id="t1", user_spend=1900.0, user_max_budget=50.0) + assert await is_token_within_budget_for_model(model="paid-model", valid_token=token, llm_router=_router()) is True + + +@pytest.mark.asyncio +async def test_team_key_inherits_personal_budget_when_opted_in(monkeypatch): + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "general_settings", {"apply_user_budget_to_team_keys": True}, raising=False) + token = _token(team_id="t1", user_spend=1900.0, user_max_budget=50.0) + assert await is_token_within_budget_for_model(model="paid-model", valid_token=token, llm_router=_router()) is False + + +@pytest.mark.asyncio +async def test_check_is_a_no_op_while_not_enforced(): + request = {"metadata": {"user_api_key_auth": _token(user_spend=1900.0, user_max_budget=50.0)}} + assert await NOT_ENFORCED(model="paid-model", request_kwargs=request, llm_router=_router()) is True + + +@pytest.mark.asyncio +async def test_request_without_a_key_is_unrestricted(): + assert await ENFORCED(model="paid-model", request_kwargs={}, llm_router=_router()) is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("metadata_field", ["metadata", "litellm_metadata"]) +async def test_enforced_check_reads_the_key_from_request_metadata(metadata_field: str): + over = {metadata_field: {"user_api_key_auth": _token(user_spend=1900.0, user_max_budget=50.0)}} + under = {metadata_field: {"user_api_key_auth": _token(user_spend=1.0, user_max_budget=50.0)}} + + assert await ENFORCED(model="paid-model", request_kwargs=over, llm_router=_router()) is False + assert await ENFORCED(model="paid-model", request_kwargs=under, llm_router=_router()) is True + + +@pytest.mark.asyncio +async def test_a_stale_low_counter_still_refuses_a_paid_target(monkeypatch): + """ + The counter can read low (e.g. restored from an older Redis snapshot). Passing the budget makes + `get_current_spend` verify against authoritative spend instead of trusting that read, so the + paid target is still refused. + """ + from litellm.proxy import proxy_server + + seen: list[dict] = [] + + async def _stale_counter(**kwargs): + seen.append(kwargs) + # a stale-low counter read; the authoritative spend is what the budget must be judged on + return 0.0 if kwargs.get("max_budget") is None else kwargs["fallback_spend"] + + monkeypatch.setattr(proxy_server, "get_current_spend", _stale_counter, raising=False) + token = _token(user_spend=1900.0, user_max_budget=50.0) + + assert await is_token_within_budget_for_model(model="paid-model", valid_token=token, llm_router=_router()) is False + assert [call["max_budget"] for call in seen] == [50.0] + + +@pytest.mark.asyncio +async def test_check_fails_closed_when_the_spend_lookup_breaks(monkeypatch): + from litellm.proxy import proxy_server + + async def _boom(**kwargs): + raise RuntimeError("spend counter unavailable") + + monkeypatch.setattr(proxy_server, "get_current_spend", _boom, raising=False) + request = {"metadata": {"user_api_key_auth": _token(user_spend=1.0, user_max_budget=50.0)}} + + assert await ENFORCED(model="paid-model", request_kwargs=request, llm_router=_router()) is False + + +@pytest.mark.asyncio +async def test_router_skips_the_paid_fallback_target_when_over_budget(): + from litellm.router_utils.fallback_event_handlers import _is_fallback_target_within_budget + + router = Router( + model_list=[FREE_MODEL, PAID_MODEL], + fallbacks=[{"free-model": ["paid-model"]}], + fallback_budget_check=ENFORCED, + ) + over = {"metadata": {"user_api_key_auth": _token(user_spend=1900.0, user_max_budget=50.0)}} + under = {"metadata": {"user_api_key_auth": _token(user_spend=1.0, user_max_budget=50.0)}} + + assert await _is_fallback_target_within_budget(router, "paid-model", "free-model", over) is False + assert await _is_fallback_target_within_budget(router, "paid-model", "free-model", under) is True + + +@pytest.mark.asyncio +async def test_router_without_a_budget_check_attempts_every_fallback(): + from litellm.router_utils.fallback_event_handlers import _is_fallback_target_within_budget + + router = _router() # fallback_budget_check defaults to None + over = {"metadata": {"user_api_key_auth": _token(user_spend=1900.0, user_max_budget=50.0)}} + + assert await _is_fallback_target_within_budget(router, "paid-model", "free-model", over) is True + + +@pytest.mark.asyncio +async def test_enforcement_is_on_by_default_and_opt_out_restores_the_leak(monkeypatch): + """ + Leaving the paid fallback unguarded is the budget bypass this module exists to close, so an + unconfigured proxy has to enforce. `enforce_fallback_budget: false` is the deliberate opt-out. + """ + from litellm.proxy import proxy_server + + over = {"metadata": {"user_api_key_auth": _token(user_spend=1900.0, user_max_budget=50.0)}} + + monkeypatch.setattr(proxy_server, "general_settings", {}, raising=False) + assert await router_fallback_budget_check(model="paid-model", request_kwargs=over, llm_router=_router()) is False + + monkeypatch.setattr(proxy_server, "general_settings", {"enforce_fallback_budget": False}, raising=False) + assert await router_fallback_budget_check(model="paid-model", request_kwargs=over, llm_router=_router()) is True diff --git a/tests/test_litellm/proxy/auth/test_team_grants.py b/tests/test_litellm/proxy/auth/test_team_grants.py index 7b6717f804f..f74531beaf0 100644 --- a/tests/test_litellm/proxy/auth/test_team_grants.py +++ b/tests/test_litellm/proxy/auth/test_team_grants.py @@ -31,6 +31,7 @@ def _full_team(model_aliases=ALIASES) -> LiteLLM_TeamTable: max_budget=50.0, soft_budget=25.0, spend=12.5, + model_max_budget={"gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}}, models=["gpt-4o", "gpt-4o-mini"], blocked=True, metadata={"tier": "gold"}, @@ -72,6 +73,7 @@ def test_team_grants_cover_every_team_field_the_key_path_gets(): assert token.team_max_budget == 50.0 assert token.team_soft_budget == 25.0 assert token.team_spend == 12.5 + assert token.team_model_max_budget == {"gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}} assert token.team_models == ["gpt-4o", "gpt-4o-mini"] assert token.team_blocked is True assert token.team_metadata == {"tier": "gold"} diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 896acc5fcef..ba3e98ee718 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -4374,6 +4374,74 @@ async def test_centralized_common_checks_carries_team_and_user_budget_state_on_t } +class _RecordingTeamModelBudgetLimiter: + def __init__(self): + self.calls = [] + + async def is_team_within_model_budget(self, team_id, team_model_max_budget, key_model_max_budget, model): + self.calls.append((team_id, dict(team_model_max_budget), key_model_max_budget, model)) + return True + + +@pytest.mark.asyncio +async def test_centralized_common_checks_enforces_team_model_max_budget_from_the_resolved_team(): + """The team's model_max_budget is enforced at the single authz gate, off the + team object auth resolved (not the possibly stale token copy), and the key's + own model_max_budget is handed to the limiter so a matching key entry can + override the team cap.""" + from fastapi import Request + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as _proxy_server_mod + + team_caps = {"gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}} + key_caps = {"claude-sonnet-4-6": {"max_budget": 1.0, "budget_duration": "1d"}} + token = UserAPIKeyAuth( + api_key="sk-test", + token="hashed", + team_id="t1", + team_model_max_budget={"gpt-4o": {"max_budget": 999.0, "budget_duration": "30d"}}, + model_max_budget=key_caps, + ) + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + user_api_key_cache = DualCache() + await user_api_key_cache.async_set_cache( + key="team_id:t1", + value=LiteLLM_TeamTableCachedObj(team_id="t1", model_max_budget=team_caps), + ) + limiter = _RecordingTeamModelBudgetLimiter() + attrs = { + **_proxy_attrs_for_centralized_checks(user_custom_auth=None), + "prisma_client": MagicMock(), + "user_api_key_cache": user_api_key_cache, + "model_max_budget_limiter": limiter, + } + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with ( + patch("litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock), # test-quality-ok: stubs the sibling check so only the team model-budget gate is under test + patch( # test-quality-ok: stubs the budget reservation so only the team model-budget gate is under test + "litellm.proxy.auth.user_api_key_auth._reserve_budget_after_common_checks", + new_callable=AsyncMock, + ), + ): + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-4o"}, + route="/chat/completions", + ) + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + assert limiter.calls == [("t1", team_caps, key_caps, "gpt-4o")] + + @pytest.mark.asyncio async def test_centralized_common_checks_skipped_for_custom_auth_without_flag(): """Existing RPS guarantee: custom-auth deployments without diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index 943a6c905c0..1ccf9be37b9 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -291,6 +291,23 @@ def test_reset_budget_for_key(reset_budget_job, mock_prisma_client): assert set(write["data"].keys()) == {"spend", "budget_reset_at"} +def test_reset_budget_for_key_leaves_lifetime_total_spend_alone(reset_budget_job, mock_prisma_client): + """A period reset zeroes spend but must neither write nor touch the lifetime total_spend.""" + now = datetime.now(timezone.utc) + key = LiteLLM_VerificationToken( + token="tok-key-1", spend=100.0, total_spend=340.0, budget_duration="30d", budget_reset_at=now + ) + mock_prisma_client.data["key"] = [key] + + asyncio.run(reset_budget_job.reset_budget_for_litellm_keys()) + + (write,) = _batch_writes(mock_prisma_client, "key") + assert write["data"]["spend"] == {"decrement": 100.0} + assert "total_spend" not in write["data"] + assert key.spend == 0.0 + assert key.total_spend == 340.0 + + def test_reset_budget_for_key_honors_injected_reset_time(mock_prisma_client, mock_proxy_logging): """Injected BudgetResetSettings drives the written reset time end to end (DI, no globals). diff --git a/tests/test_litellm/proxy/db/test_create_views.py b/tests/test_litellm/proxy/db/test_create_views.py index ecc6d70123e..54418e10bdf 100644 --- a/tests/test_litellm/proxy/db/test_create_views.py +++ b/tests/test_litellm/proxy/db/test_create_views.py @@ -71,6 +71,7 @@ async def test_create_views_creates_view_on_does_not_exist(): mock_db.execute_raw.assert_called_once() created_sql = mock_db.execute_raw.call_args[0][0] assert 'CREATE VIEW "LiteLLM_VerificationTokenView"' in created_sql + assert "t.model_max_budget AS team_model_max_budget" in created_sql @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index c547d06904b..8f72e1d6248 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -1658,6 +1658,57 @@ async def test_commit_key_spend_updates_includes_last_active(): assert before_call <= last_active <= after_call +@pytest.mark.asyncio +async def test_commit_spend_updates_to_db_increments_key_total_spend_alongside_spend(): + """ + The key table write must increment the lifetime total_spend by the same amount as the + resettable spend, in the same update so the two cannot drift. + """ + db_writer = DBSpendUpdateWriter() + + mock_batcher = MagicMock() + mock_batcher.litellm_verificationtoken = MagicMock() + mock_batcher.litellm_verificationtoken.update_many = MagicMock() + + mock_transaction = AsyncMock() + mock_transaction.__aenter__ = AsyncMock(return_value=mock_transaction) + mock_transaction.__aexit__ = AsyncMock(return_value=False) + mock_transaction.batch_ = MagicMock( + return_value=AsyncMock( + __aenter__=AsyncMock(return_value=mock_batcher), + __aexit__=AsyncMock(return_value=False), + ) + ) + + mock_prisma_client = MagicMock() + mock_prisma_client.db = MagicMock() + mock_prisma_client.db.tx = MagicMock(return_value=mock_transaction) + + db_spend_update_transactions = { + "user_list_transactions": {}, + "end_user_list_transactions": {}, + "key_list_transactions": {"hashed_token_abc": 0.05, "hashed_token_def": 1.25}, + "team_list_transactions": {}, + "team_member_list_transactions": {}, + "org_list_transactions": {}, + "tag_list_transactions": {}, + "agent_list_transactions": {}, + } + + await db_writer._commit_spend_updates_to_db( + prisma_client=mock_prisma_client, + n_retry_times=0, + proxy_logging_obj=MagicMock(), + db_spend_update_transactions=db_spend_update_transactions, + ) + + calls = mock_batcher.litellm_verificationtoken.update_many.call_args_list + assert [c.kwargs["where"] for c in calls] == [{"token": "hashed_token_abc"}, {"token": "hashed_token_def"}] + for call, expected_cost in zip(calls, (0.05, 1.25)): + assert call.kwargs["data"]["spend"] == {"increment": expected_cost} + assert call.kwargs["data"]["total_spend"] == call.kwargs["data"]["spend"] + + @pytest.mark.asyncio async def test_update_database_creates_single_task(): """ @@ -2813,7 +2864,7 @@ async def test_commit_spend_updates_to_db_does_not_stamp_key_settings_updated_at mock_batcher.litellm_verificationtoken.update_many.assert_called_once() call_kwargs = mock_batcher.litellm_verificationtoken.update_many.call_args[1] assert call_kwargs["where"] == {"token": token} - assert set(call_kwargs["data"]) == {"spend", "last_active"} + assert set(call_kwargs["data"]) == {"spend", "total_spend", "last_active"} assert call_kwargs["data"]["spend"] == {"increment": response_cost} diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index d173c5f5c70..1d3d7a452b6 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -5603,6 +5603,7 @@ def test_initialize_bedrock_wires_streaming_flags(): streaming_buffer_until_moderated=False, streaming_sampling_rate=3, streaming_end_of_stream_only=True, + streaming_buffer_release_on_scan=True, ), {"guardrail_name": "bedrock-streaming"}, ) @@ -5616,9 +5617,11 @@ def test_initialize_bedrock_wires_streaming_flags(): assert configured.streaming_buffer_until_moderated is False assert configured.streaming_sampling_rate == 3 assert configured.streaming_end_of_stream_only is True + assert configured.streaming_buffer_release_on_scan is True assert defaulted.streaming_buffer_until_moderated is True assert defaulted.streaming_sampling_rate == 5 assert defaulted.streaming_end_of_stream_only is False + assert defaulted.streaming_buffer_release_on_scan is False def test_initialize_bedrock_rejects_non_positive_sampling_rate(): @@ -5721,6 +5724,44 @@ async def test_buffered_default_hook_scans_before_any_chunk(): assert len([e for e in events if e != "scan"]) >= 1 +@pytest.mark.asyncio +async def test_buffered_release_on_scan_hook_releases_each_window_after_its_scan(): + guardrail = BedrockGuardrail( + guardrail_name="bedrock-release-on-scan", + guardrailIdentifier="test-id", + guardrailVersion="DRAFT", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + streaming_buffer_release_on_scan=True, + streaming_sampling_rate=1, + ) + + assert guardrail._streams_incrementally() is True + events = await _run_streaming_hook_recording_order(guardrail) + + assert events == ["scan", ("chunk", "Hello"), "scan", ("chunk", " world"), ("chunk", "")] + + +@pytest.mark.asyncio +async def test_buffered_release_on_scan_defers_to_end_of_stream_only(): + guardrail = BedrockGuardrail( + guardrail_name="bedrock-release-on-scan-end-only", + guardrailIdentifier="test-id", + guardrailVersion="DRAFT", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + streaming_buffer_release_on_scan=True, + streaming_end_of_stream_only=True, + streaming_sampling_rate=1, + ) + + assert guardrail._streams_incrementally() is False + events = await _run_streaming_hook_recording_order(guardrail) + + assert events.count("scan") == 1 + assert events[0] == "scan" + + @pytest.mark.asyncio async def test_masking_keeps_buffered_path_even_when_unbuffered_configured(): guardrail = BedrockGuardrail( diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py index 9849ad7ec88..a1aae119d56 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_crowdstrike_aidr.py @@ -1622,10 +1622,23 @@ def test_initialize_guardrail_rejects_unsupported_mode_instead_of_running_other_ def test_initialize_guardrail_defaults_streaming_params() -> None: handler = _initialize_from_config(mode="post_call") + assert handler.streaming_buffer_until_moderated is False + assert handler.streaming_buffer_release_on_scan is False assert handler.streaming_end_of_stream_only is False assert handler.streaming_sampling_rate == 5 +def test_initialize_guardrail_forwards_buffer_streaming_params() -> None: + handler = _initialize_from_config( + mode="post_call", + streaming_buffer_until_moderated=True, + streaming_buffer_release_on_scan=True, + ) + + assert handler.streaming_buffer_until_moderated is True + assert handler.streaming_buffer_release_on_scan is True + + @pytest.mark.parametrize( "configured", [ diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py index 87a1b84acc5..427fa43ffd5 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py @@ -1291,8 +1291,9 @@ class TestToolPermissionGuardrailAnthropicMessages: async def test_rewrite_mode_keeps_the_stream_identity_it_had_before_the_shared_helper(self): """Well-formed SSE must round-trip exactly as it did before the helpers were shared. - The shared module can stamp the upstream message id and model onto the assembled response - for callers that ask for it; this path never did, and a client reads those bytes. + The shared module can stamp the upstream message id onto the assembled response for + callers that ask for it; this path never did, and a client reads those bytes. The model, + though, is now the upstream's, matching what the untouched passthrough shows clients. """ with patch.object(self.rewriting, "should_run_guardrail", return_value=True): out = await self._drain(self.rewriting, self._sse_chunks("Read")) @@ -1304,7 +1305,7 @@ class TestToolPermissionGuardrailAnthropicMessages: if line.startswith("data: ") and json.loads(line[6:]).get("type") == "message_start" )["message"] assert message_start["id"].startswith("chatcmpl-"), "the rewritten stream must not adopt the upstream message id" - assert message_start["model"] == "unknown-model", "the rewritten stream must not adopt the upstream model" + assert message_start["model"] == "claude-sonnet-4-5", "the rewritten stream reports the model the upstream served" @pytest.mark.asyncio async def test_message_start_without_a_dict_message_fails_closed(self): diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py index 2b163ee5233..db937f18e96 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_streaming_buffer_until_moderated.py @@ -11,7 +11,7 @@ released unchanged after moderation passes. """ import json -from typing import Any, List, Literal, Optional +from typing import Any, AsyncGenerator, List, Literal, Optional import pytest @@ -19,14 +19,25 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, ModifyResponseException, ) +from litellm.llms.base_llm.guardrail_translation.base_translation import StreamingScanKey from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( UnifiedLLMGuardrails, + _is_redundant_scan, +) +from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + Delta, + Function, + FunctionCall, + GenericGuardrailAPIInputs, + ModelResponseStream, + StreamingChoices, ) -from litellm.types.utils import GenericGuardrailAPIInputs BLOCK_MESSAGE = "Blocked by policy: this response was withheld." ORIGINAL_MARKER = "ORIGINAL-SECRET-ANSWER" +TOOL_ARGUMENTS_MARKER = "TOOL-ARGS-SECRET" class _BlockingGuardrail(CustomGuardrail): @@ -60,6 +71,85 @@ class _PassingGuardrail(CustomGuardrail): return inputs +class _CountingPassingGuardrail(_PassingGuardrail): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.scan_count = 0 + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + self.scan_count += 1 + return inputs + + +class _ToolCallRecordingGuardrail(_CountingPassingGuardrail): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.tool_call_scan_indexes: List[int] = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + self.scan_count += 1 + if inputs.get("tool_calls"): + self.tool_call_scan_indexes.append(self.scan_count) + return inputs + + +class _SecondScanBlockingGuardrail(_CountingPassingGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + self.scan_count += 1 + if self.scan_count == 2: + raise ModifyResponseException( + message=BLOCK_MESSAGE, + model="gpt-4", + request_data=request_data, + guardrail_name=self.guardrail_name, + ) + return inputs + + +class _MarkerBlockingGuardrail(_CountingPassingGuardrail): + """Blocks as soon as the inspected input field (texts or tool_calls) carries the marker.""" + + def __init__(self, *args, marker: str, field: Literal["texts", "tool_calls"] = "texts", **kwargs): + super().__init__(*args, **kwargs) + self.marker = marker + self.field = field + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + self.scan_count += 1 + if self.marker in json.dumps(inputs.get(self.field, [])): + raise ModifyResponseException( + message=BLOCK_MESSAGE, + model="gpt-4o", + request_data=request_data, + guardrail_name=self.guardrail_name, + ) + return inputs + + def _sse_event(event_type: str, data: dict) -> bytes: return f"event: {event_type}\ndata: {json.dumps(data)}\n\n".encode() @@ -115,6 +205,212 @@ def _decode(chunks: List[Any]) -> str: return "".join(c.decode() if isinstance(c, bytes) else str(c) for c in chunks) +def _chat_chunk(content: str = "", finish_reason: str | None = None) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-windowed", + created=1724900000, + model="gpt-4", + choices=[ + StreamingChoices( + index=0, + delta=Delta(role="assistant", content=content), + finish_reason=finish_reason, + ) + ], + ) + + +def _tool_call_chunk( + arguments: str, finish_reason: str | None = None, legacy_function_call: bool = False +) -> ModelResponseStream: + delta = ( + Delta(role="assistant", content=None, function_call=FunctionCall(name="run_shell", arguments=arguments)) + if legacy_function_call + else Delta( + role="assistant", + content=None, + tool_calls=[ + ChatCompletionDeltaToolCall( + id="call_1", + type="function", + index=0, + function=Function(name="run_shell", arguments=arguments), + ) + ], + ) + ) + return ModelResponseStream( + id="chatcmpl-windowed", + created=1724900000, + model="gpt-4", + choices=[StreamingChoices(index=0, delta=delta, finish_reason=finish_reason)], + ) + + +async def _windowed_chat_stream( + yielded_count: List[int], + collected: List[Any], + content_chunks: List[str], + tool_argument_chunks: List[str] | None = None, + legacy_function_call: bool = False, +) -> AsyncGenerator[ModelResponseStream, None]: + for content in content_chunks: + yielded_count.append(len(collected)) + yield _chat_chunk(content) + for arguments in tool_argument_chunks or []: + yielded_count.append(len(collected)) + yield _tool_call_chunk(arguments, legacy_function_call=legacy_function_call) + yielded_count.append(len(collected)) + yield _chat_chunk(finish_reason="tool_calls" if tool_argument_chunks else "stop") + + +def _tool_argument_text(chunks: List[Any]) -> str: + return "".join( + tool_call.function.arguments or "" + for chunk in chunks + if isinstance(chunk, ModelResponseStream) + for choice in chunk.choices + for tool_call in choice.delta.tool_calls or [] + ) + + +def _function_call_argument_text(chunks: list[Any]) -> str: + return "".join( + choice.delta.function_call.arguments or "" + for chunk in chunks + if isinstance(chunk, ModelResponseStream) + for choice in chunk.choices + if choice.delta.function_call is not None + ) + + +async def _run_windowed( + guardrail: CustomGuardrail, + content_chunks: List[str], + end_of_stream_only: bool = False, + tool_argument_chunks: List[str] | None = None, + legacy_function_call: bool = False, +) -> tuple[List[Any], List[int]]: + guardrail.streaming_buffer_until_moderated = True + guardrail.streaming_buffer_release_on_scan = True + guardrail.streaming_end_of_stream_only = end_of_stream_only + guardrail.streaming_sampling_rate = 2 + unified = UnifiedLLMGuardrails() + user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/v1/chat/completions") + request_data = { + "messages": [{"role": "user", "content": "hi"}], + "guardrail_to_apply": guardrail, + "metadata": {"guardrails": [guardrail.guardrail_name]}, + } + collected: List[Any] = [] + yielded_count: List[int] = [] + async for chunk in unified.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=_windowed_chat_stream( + yielded_count, collected, content_chunks, tool_argument_chunks, legacy_function_call + ), + request_data=request_data, + ): + collected.append(chunk) + return collected, yielded_count + + +def _responses_message_stream_events(text_chunks: List[str]) -> List[dict]: + message = {"type": "message", "id": "msg_1", "status": "completed", "role": "assistant"} + content = [{"type": "output_text", "text": "".join(text_chunks), "annotations": []}] + return [ + {"type": "response.output_item.added", "output_index": 0, "item": {**message, "content": []}}, + *( + { + "type": "response.output_text.delta", + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": text, + } + for text in text_chunks + ), + {"type": "response.output_item.done", "output_index": 0, "item": {**message, "content": content}}, + { + "type": "response.completed", + "response": { + "id": "resp_1", + "model": "gpt-4o", + "status": "completed", + "output": [{**message, "content": content}], + }, + }, + ] + + +def _responses_truncated_function_call_events(text: str, argument_chunks: List[str]) -> List[dict]: + message = {"type": "message", "id": "msg_1", "status": "completed", "role": "assistant"} + content = [{"type": "output_text", "text": text, "annotations": []}] + function_call = {"type": "function_call", "id": "fc_1", "call_id": "call_1", "name": "run_shell"} + return [ + {"type": "response.output_item.added", "output_index": 0, "item": {**message, "content": []}}, + { + "type": "response.output_text.delta", + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": text, + }, + {"type": "response.output_item.added", "output_index": 1, "item": {**function_call, "arguments": ""}}, + *( + {"type": "response.function_call_arguments.delta", "item_id": "fc_1", "output_index": 1, "delta": arguments} + for arguments in argument_chunks + ), + { + "type": "response.incomplete", + "response": { + "id": "resp_1", + "model": "gpt-4o", + "status": "incomplete", + "output": [ + {**message, "content": content}, + {**function_call, "arguments": "".join(argument_chunks), "status": "incomplete"}, + ], + }, + }, + ] + + +async def _replay(events: List[dict]) -> AsyncGenerator[dict, None]: + for event in events: + yield event + + +async def _run_windowed_responses(guardrail: CustomGuardrail, events: List[dict]) -> str: + guardrail.streaming_buffer_until_moderated = True + guardrail.streaming_buffer_release_on_scan = True + guardrail.streaming_sampling_rate = 2 + unified = UnifiedLLMGuardrails() + user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/v1/responses") + request_data = { + "input": "hi", + "guardrail_to_apply": guardrail, + "metadata": {"guardrails": [guardrail.guardrail_name]}, + } + collected: List[Any] = [] + async for chunk in unified.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=_replay(events), + request_data=request_data, + ): + collected.append(chunk) + return json.dumps([chunk if isinstance(chunk, dict) else str(chunk) for chunk in collected]) + + +def _chat_text(chunks: List[Any]) -> str: + return "".join( + choice.delta.content or "" + for chunk in chunks + if isinstance(chunk, ModelResponseStream) + for choice in chunk.choices + ) + + async def _run(guardrail: CustomGuardrail) -> str: # Rubrik's real config: end-of-stream-only moderation. Without buffering # this releases every chunk before moderation runs (content leaks on @@ -159,6 +455,110 @@ async def test_buffered_clean_releases_all_content(): assert BLOCK_MESSAGE not in raw +@pytest.mark.asyncio +async def test_windowed_buffer_releases_after_each_passing_scan(): + guardrail = _CountingPassingGuardrail(guardrail_name="windowed-pass", event_hook="post_call") + content_chunks = ["one ", "two ", "three ", "four ", "five ", "six "] + + collected, yielded_count = await _run_windowed(guardrail, content_chunks) + + assert yielded_count[2] >= 2 + assert yielded_count == [0, 0, 2, 2, 4, 4, 6] + assert _chat_text(collected) == "".join(content_chunks) + assert guardrail.scan_count > 1 + + +@pytest.mark.asyncio +async def test_windowed_buffer_drops_blocked_window(): + guardrail = _SecondScanBlockingGuardrail(guardrail_name="windowed-block", event_hook="post_call") + content_chunks = ["one ", "two ", "MARKER ", "four ", "five ", "six "] + + collected, _ = await _run_windowed(guardrail, content_chunks) + raw = _decode(collected) + + assert _chat_text(collected) == "one two " + assert "MARKER" not in raw + assert BLOCK_MESSAGE in raw + assert '"error"' not in raw + + +@pytest.mark.asyncio +async def test_windowed_buffer_holds_tool_call_windows_until_end_of_stream_scan(): + guardrail = _ToolCallRecordingGuardrail(guardrail_name="windowed-tools", event_hook="post_call") + content_chunks = ["one ", "two ", "three "] + tool_argument_chunks = ['{"cmd": "', TOOL_ARGUMENTS_MARKER, '"}'] + + collected, yielded_count = await _run_windowed(guardrail, content_chunks, tool_argument_chunks=tool_argument_chunks) + + assert yielded_count == [0, 0, 2, 2, 2, 2, 2] + assert _chat_text(collected) == "".join(content_chunks) + assert _tool_argument_text(collected) == "".join(tool_argument_chunks) + assert guardrail.tool_call_scan_indexes == [guardrail.scan_count] + + +@pytest.mark.asyncio +async def test_windowed_buffer_holds_legacy_function_call_windows_until_end_of_stream(): + guardrail = _PassingGuardrail(guardrail_name="windowed-functions", event_hook="post_call") + content_chunks = ["one ", "two ", "three "] + function_argument_chunks = ['{"cmd": "', TOOL_ARGUMENTS_MARKER, '"}'] + + collected, yielded_count = await _run_windowed( + guardrail, content_chunks, tool_argument_chunks=function_argument_chunks, legacy_function_call=True + ) + + assert yielded_count == [0, 0, 2, 2, 2, 2, 2] + assert _chat_text(collected) == "".join(content_chunks) + assert _function_call_argument_text(collected) == "".join(function_argument_chunks) + + +def test_tool_call_only_scan_key_is_not_skipped_as_empty(): + assert _is_redundant_scan(StreamingScanKey(texts=("",)), None) is True + assert _is_redundant_scan(StreamingScanKey(texts=("",), tool_calls=("run_shell:{}",)), None) is False + + +@pytest.mark.asyncio +async def test_windowed_responses_output_item_done_round_keeps_text_window_withheld(): + guardrail = _MarkerBlockingGuardrail( + guardrail_name="windowed-responses", event_hook="post_call", marker=ORIGINAL_MARKER + ) + events = _responses_message_stream_events(["one ", f"{ORIGINAL_MARKER} "]) + + raw = await _run_windowed_responses(guardrail, events) + + assert ORIGINAL_MARKER not in raw, f"unscanned window leaked: {raw!r}" + assert BLOCK_MESSAGE in raw + assert guardrail.scan_count >= 1 + + +@pytest.mark.asyncio +async def test_windowed_responses_incomplete_stream_scans_tool_call_before_release(): + guardrail = _MarkerBlockingGuardrail( + guardrail_name="windowed-responses-tools", + event_hook="post_call", + marker=TOOL_ARGUMENTS_MARKER, + field="tool_calls", + ) + events = _responses_truncated_function_call_events("hi ", ['{"cmd": "', TOOL_ARGUMENTS_MARKER, '"}']) + + raw = await _run_windowed_responses(guardrail, events) + + assert '"hi "' in raw + assert TOOL_ARGUMENTS_MARKER not in raw, f"unscanned tool call leaked: {raw!r}" + assert BLOCK_MESSAGE in raw + + +@pytest.mark.asyncio +async def test_windowed_buffer_with_explicit_end_of_stream_only_stays_fully_buffered(): + guardrail = _CountingPassingGuardrail(guardrail_name="windowed-eos", event_hook="post_call") + content_chunks = ["one ", "two ", "three ", "four ", "five ", "six "] + + collected, yielded_count = await _run_windowed(guardrail, content_chunks, end_of_stream_only=True) + + assert yielded_count == [0, 0, 0, 0, 0, 0, 0] + assert _chat_text(collected) == "".join(content_chunks) + assert guardrail.scan_count == 1 + + @pytest.mark.asyncio async def test_buffered_mode_disabled_for_content_rewriting_guardrail(): """Buffered replay yields the withheld *original* chunks verbatim, which diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index 530f8ffd854..bf641fd6cd0 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -682,6 +682,22 @@ async def test_provider_specific_params_includes_embedding_toggle(): assert field["default_value"] is False +@pytest.mark.asyncio +async def test_provider_specific_params_exposes_bedrock_streaming_flags(): + from litellm.proxy.guardrails.guardrail_endpoints import get_provider_specific_params + + provider_params = await get_provider_specific_params() + + bedrock = provider_params["bedrock"] + assert "guardrailIdentifier" in bedrock + assert "guardrailVersion" in bedrock + assert bedrock["streaming_buffer_release_on_scan"]["type"] == "boolean" + assert bedrock["streaming_buffer_release_on_scan"]["default_value"] is False + assert bedrock["streaming_buffer_until_moderated"]["default_value"] is True + assert bedrock["streaming_end_of_stream_only"]["type"] == "boolean" + assert bedrock["streaming_sampling_rate"]["type"] == "number" + + @pytest.mark.asyncio async def test_provider_specific_params_includes_hide_secrets(): """hide-secrets lives in the enterprise package so it is not in diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 8d7ab89f354..85023b94207 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -7,6 +7,7 @@ import logging import os import sys import time +from collections.abc import Sequence from contextlib import contextmanager from datetime import datetime, timedelta from typing import Any, Dict, List, Optional @@ -32,6 +33,7 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import ( ) from litellm.proxy.utils import InternalUsageCache, ProxyLogging, hash_token from litellm.types.caching import RedisPipelineIncrementOperation +from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.utils import ( EmbeddingResponse, ModelResponse, @@ -4054,6 +4056,125 @@ async def _seed_max_parallel_requests_slots( ) +@pytest.mark.asyncio +async def test_completed_responses_post_call_releases_parallel_slot() -> None: + api_key = hash_token("sk-responses-post-call") + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + user_api_key_dict = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1) + data = { + "model": "gpt-4o-mini", + "input": "hello", + "litellm_call_id": "responses-owner", + } + parallel_key = f"{{api_key:{api_key}}}:max_parallel_requests" + + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data=data, + call_type="aresponses", + ) + assert handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=parallel_key) + ) == 1 + + await handler.async_post_call_success_hook( + data=data, + user_api_key_dict=user_api_key_dict, + response=ResponsesAPIResponse( + id="resp_parallel_slot", + created_at=0, + model="gpt-4o-mini", + object="response", + output=[], + status="completed", + ), + ) + assert handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=parallel_key) + ) == 0 + + await handler.async_log_success_event( + kwargs={"litellm_call_id": data["litellm_call_id"]}, + response_obj=None, + start_time=None, + end_time=None, + ) + assert handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=parallel_key) + ) == 0 + + +@pytest.mark.asyncio +async def test_concurrent_success_callbacks_release_parallel_slot_once_when_redis_fails() -> None: + from unittest.mock import AsyncMock + + api_key = hash_token("sk-concurrent-release") + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + user_api_key_dict = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=2) + call_id = "concurrent-release-owner" + parallel_key = f"{{api_key:{api_key}}}:max_parallel_requests" + release_started = asyncio.Event() + allow_redis_failure = asyncio.Event() + + async def failing_release( + keys: Sequence[str], args: Sequence[object] + ) -> list[int]: + release_started.set() + await allow_redis_failure.wait() + raise ConnectionError("redis unavailable") + + release_script = AsyncMock(side_effect=failing_release) + handler.parallel_release_script = release_script + await local_cache.async_set_cache(key=parallel_key, value=2, local_only=True) + stash = get_or_create_request_stash() + stash.owner_litellm_call_id = call_id + stash.parallel_slot = ParallelSlotAcquisition( + slot_id="slot-concurrent-release", + counter_keys=[parallel_key], + ) + data = {"litellm_call_id": call_id} + + post_call_task = asyncio.create_task( + handler.async_post_call_success_hook( + data=data, + user_api_key_dict=user_api_key_dict, + response=ResponsesAPIResponse( + id="resp_concurrent_release", + created_at=0, + model="gpt-4o-mini", + object="response", + output=[], + status="completed", + ), + ) + ) + await asyncio.wait_for(release_started.wait(), timeout=5) + logging_task = asyncio.create_task( + handler.async_log_success_event( + kwargs=data, + response_obj=None, + start_time=None, + end_time=None, + ) + ) + allow_redis_failure.set() + await asyncio.wait_for( + asyncio.gather(post_call_task, logging_task), + timeout=5, + ) + + assert release_script.await_count == 1 + assert await local_cache.async_get_cache(key=parallel_key) == 1 + assert stash.parallel_slot is None + + async def _build_seeded_limiter(): """Build a v3 limiter whose api-key slot registry already holds the pre-call slot.""" api_key = hash_token("sk-disconnect") diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_utils.py b/tests/test_litellm/proxy/management_endpoints/test_common_utils.py index 7352ca0e9ee..2b614632346 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_utils.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_utils.py @@ -35,6 +35,7 @@ from litellm.proxy.management_endpoints.common_utils import ( admin_can_invite_user, ) from litellm.proxy.management_endpoints.common_utils import _has_non_empty_value +from litellm.types.utils import BudgetConfig class TestUpdateMetadataFieldsEmptyCollections: @@ -1162,3 +1163,54 @@ async def test_router_weights_validate_current_deployment_scope( assert exc.value.detail == error else: await validation + + +@pytest.mark.parametrize( + "model_max_budget, error", + [ + ({"gpt-4o": BudgetConfig(max_budget=-1.0, budget_duration="1d")}, "non-negative finite"), + ({"gpt-4o": BudgetConfig(max_budget=float("inf"), budget_duration="1d")}, "non-negative finite"), + ({"gpt-4o": BudgetConfig(max_budget=float("nan"), budget_duration="1d")}, "non-negative finite"), + ({"gpt-4o": BudgetConfig(budget_duration="1d")}, "non-negative finite"), + ({"gpt-4o": BudgetConfig(max_budget=5.0)}, "requires a budget_duration"), + ({"gpt-4o": BudgetConfig(max_budget=5.0, budget_duration="fortnight")}, "budget_duration"), + ({" ": BudgetConfig(max_budget=5.0, budget_duration="1d")}, "non-empty model names"), + ({"gpt-4o": BudgetConfig(max_budget=5.0, budget_duration="1d", tpm_limit=1000)}, "not enforced on a team"), + ({"gpt-4o": BudgetConfig(max_budget=5.0, budget_duration="1d", rpm_limit=10)}, "not enforced on a team"), + ], + ids=["negative", "inf", "nan", "no_cap", "no_duration", "bad_duration", "blank_model", "tpm_limit", "rpm_limit"], +) +def test_validate_team_model_max_budget_rejects_unenforceable_entries(model_max_budget, error) -> None: + from litellm.proxy.management_endpoints.common_utils import validate_team_model_max_budget + + with pytest.raises(HTTPException) as exc: + validate_team_model_max_budget(model_max_budget=model_max_budget, premium_user=True) + assert exc.value.status_code == 400 + assert error in exc.value.detail["error"] + + +def test_validate_team_model_max_budget_accepts_a_zero_cap_and_prefixed_models() -> None: + from litellm.proxy.management_endpoints.common_utils import validate_team_model_max_budget + + assert ( + validate_team_model_max_budget( + model_max_budget={ + "gpt-4o": BudgetConfig(max_budget=0.0, budget_duration="1d"), + "openai/gpt-4o-mini": BudgetConfig(max_budget=2.5, budget_duration="30d"), + }, + premium_user=True, + ) + is None + ) + + +def test_validate_team_model_max_budget_is_license_gated_only_when_set() -> None: + from litellm.proxy.management_endpoints.common_utils import validate_team_model_max_budget + + validate_team_model_max_budget(model_max_budget=None, premium_user=False) + validate_team_model_max_budget(model_max_budget={}, premium_user=False) + with pytest.raises(HTTPException) as exc: + validate_team_model_max_budget( + model_max_budget={"gpt-4o": BudgetConfig(max_budget=1.0, budget_duration="1d")}, premium_user=False + ) + assert exc.value.status_code == 403 diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 4e70063015d..cc0a7631b59 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -1431,6 +1431,60 @@ async def test_key_info_returns_object_permission(monkeypatch): ) +def _stored_key_with_lifetime_spend(token: str, spend: float, total_spend: float) -> LiteLLM_VerificationToken: + return LiteLLM_VerificationToken.model_validate( + {"token": token, "user_id": "user123", "spend": spend, "total_spend": total_spend} + ) + + +@pytest.mark.asyncio +async def test_key_info_returns_lifetime_total_spend_next_to_resettable_spend(monkeypatch): + """After a budget reset the period spend is 0 while total_spend keeps the lifetime figure.""" + from litellm.proxy.management_endpoints.key_management_endpoints import info_key_fn + + mock_prisma_client = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=_stored_key_with_lifetime_spend(token="hashed_key", spend=0.0, total_spend=3.75) + ) + + result = await info_key_fn( + key="sk-test-key-456", + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-test-key-456"), + ) + + assert result["info"]["spend"] == 0.0 + assert result["info"]["total_spend"] == 3.75 + + +@pytest.mark.asyncio +async def test_list_keys_full_object_returns_lifetime_total_spend(): + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[_stored_key_with_lifetime_spend(token="hashed_key", spend=0.0, total_spend=3.75)] + ) + mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=1) + + result = await _list_key_helper( + prisma_client=mock_prisma_client, + page=1, + size=50, + user_id=None, + team_id=None, + organization_id=None, + key_alias=None, + key_hash=None, + exclude_team_id=None, + return_full_object=True, + admin_team_ids=None, + ) + + listed_key = result["keys"][0] + assert isinstance(listed_key, UserAPIKeyAuth) + assert listed_key.spend == 0.0 + assert listed_key.total_spend == 3.75 + + @pytest.mark.asyncio async def test_get_new_token_with_valid_key(monkeypatch): """Test get_new_token function when provided with a valid key that starts with 'sk-'""" @@ -4923,6 +4977,23 @@ def test_transform_verification_tokens_to_deleted_records(): assert json.loads(record2["budget_fallbacks"]) == {"gpt-4": ["gpt-4o-mini"]} +def test_transform_verification_tokens_to_deleted_records_keeps_organization_id(): + live_row = MagicMock() + live_row.model_dump.return_value = { + "token": "hashed-token-org", + "user_id": "user-123", + "team_id": None, + "organization_id": "org-finops", + } + + records = _transform_verification_tokens_to_deleted_records( + keys=[live_row], + user_api_key_dict=UserAPIKeyAuth(user_id="admin-1", api_key="sk-admin"), + ) + + assert records[0]["organization_id"] == "org-finops" + + def test_transform_verification_tokens_to_deleted_records_empty_list(): user_api_key_dict = UserAPIKeyAuth( user_id="user-123", @@ -6022,6 +6093,244 @@ async def test_list_keys_with_invalid_status(): assert "deleted" in str(exc_info.value.message) +@pytest.mark.asyncio +@pytest.mark.parametrize("status_filter", ["active", "expired", "revoked"]) +async def test_list_keys_accepts_live_status_filters(monkeypatch, status_filter): + from unittest.mock import Mock + + from litellm.proxy.management_endpoints.key_management_endpoints import list_keys + + live_row = MagicMock() + live_row.model_dump.return_value = {"token": "hashed_live_token", "object_permission_id": None} + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[live_row]) + mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=1) + mock_prisma_client.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + response = await list_keys( + request=Mock(), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + page=1, + size=10, + user_id=None, + team_id=None, + organization_id=None, + key_hash=None, + key_alias=None, + search=None, + return_full_object=False, + include_team_keys=False, + include_created_by_keys=False, + sort_by=None, + sort_order="desc", + expand=None, + status=status_filter, + project_id=None, + access_group_id=None, + agent_id=None, + substring_matching=False, + expires=None, + ) + + assert response["keys"] == ["hashed_live_token"] + assert response["total_count"] == 1 + mock_prisma_client.db.litellm_deletedverificationtoken.find_many.assert_not_called() + + +def _status_filter_where(status_filter: str | None) -> Mapping[str, object]: + from litellm.proxy.management_endpoints.key_management_endpoints import _build_key_filter_conditions + + return _build_key_filter_conditions( + user_id=None, + team_id=None, + organization_id=None, + key_alias=None, + key_hash=None, + exclude_team_id=None, + admin_team_ids=None, + status_filter=status_filter, + ) + + +def test_build_key_filter_conditions_status_filter_partitions_live_keys(): + not_blocked = {"OR": [{"blocked": None}, {"blocked": False}]} + + revoked_where = _status_filter_where("revoked") + assert {"blocked": True} in revoked_where["AND"] + + expired_clause = next(clause for clause in _status_filter_where("expired")["AND"] if "AND" in clause) + assert expired_clause["AND"][0] == not_blocked + assert expired_clause["AND"][1]["AND"][0] == {"expires": {"not": None}} + assert "lt" in expired_clause["AND"][1]["AND"][1]["expires"] + + active_clause = next(clause for clause in _status_filter_where("active")["AND"] if "AND" in clause) + assert active_clause["AND"][0] == not_blocked + assert active_clause["AND"][1]["OR"][0] == {"expires": None} + assert "gte" in active_clause["AND"][1]["OR"][1]["expires"] + + +def test_build_key_filter_conditions_deleted_status_adds_no_live_clause(): + assert _status_filter_where("deleted") == _status_filter_where(None) + + +@pytest.mark.asyncio +async def test_list_key_helper_revoked_status_filters_live_table_on_blocked(): + mock_prisma_client = AsyncMock() + mock_find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many + mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) + + await _list_key_helper( + prisma_client=mock_prisma_client, + page=1, + size=50, + user_id=None, + team_id=None, + organization_id=None, + key_alias=None, + key_hash=None, + exclude_team_id=None, + return_full_object=True, + admin_team_ids=None, + include_created_by_keys=False, + status="revoked", + ) + + mock_prisma_client.db.litellm_deletedverificationtoken.find_many.assert_not_called() + where = mock_find_many.call_args.kwargs["where"] + assert {"blocked": True} in where["AND"] + + +def _archived_key_row(token: str, user_id: str) -> MagicMock: + row = MagicMock() + row.model_dump.return_value = { + "id": "archive-row-1", + "token": token, + "key_alias": "finops-2024", + "user_id": user_id, + "team_id": None, + "organization_id": "org-finops", + "blocked": None, + "deleted_at": datetime(2024, 11, 15, 10, 0, tzinfo=timezone.utc), + "deleted_by": "admin-1", + } + return row + + +@pytest.mark.asyncio +async def test_info_key_fn_serves_deleted_key_from_archive(monkeypatch): + from litellm.proxy.management_endpoints.key_management_endpoints import info_key_fn + + hashed = "hashed_deleted_token" + mock_prisma_client = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_deletedverificationtoken.find_first = AsyncMock( + return_value=_archived_key_row(hashed, "user-x") + ) + + result = await info_key_fn( + key=hashed, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin"), + ) + + mock_prisma_client.db.litellm_deletedverificationtoken.find_first.assert_awaited_once() + assert mock_prisma_client.db.litellm_deletedverificationtoken.find_first.await_args.kwargs["where"] == { + "token": hashed + } + info = result["info"] + assert info["status"] == "deleted" + assert info["key_alias"] == "finops-2024" + assert info["organization_id"] == "org-finops" + assert info["deleted_by"] == "admin-1" + assert info["deleted_at"] is not None + assert "token" not in info + + +@pytest.mark.asyncio +async def test_info_key_fn_archived_key_keeps_owner_authorization(monkeypatch): + from litellm.proxy.management_endpoints.key_management_endpoints import info_key_fn + + hashed = "hashed_deleted_token" + mock_prisma_client = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_deletedverificationtoken.find_first = AsyncMock( + return_value=_archived_key_row(hashed, "owner-1") + ) + + with pytest.raises(ProxyException) as exc_info: + await info_key_fn( + key=hashed, + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="someone-else", api_key="sk-other" + ), + ) + assert exc_info.value.code == "403" + + owner_result = await info_key_fn( + key=hashed, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="owner-1", api_key="sk-own"), + ) + assert owner_result["info"]["status"] == "deleted" + + +@pytest.mark.asyncio +async def test_info_key_fn_unknown_key_still_404s(monkeypatch): + from litellm.proxy.management_endpoints.key_management_endpoints import info_key_fn + + mock_prisma_client = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_deletedverificationtoken.find_first = AsyncMock(return_value=None) + + with pytest.raises(ProxyException) as exc_info: + await info_key_fn( + key="hashed_missing", + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin"), + ) + assert exc_info.value.code == "404" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("blocked", "expires", "expected_status"), + [ + (True, None, "revoked"), + (True, "2020-01-01T00:00:00Z", "revoked"), + (False, "2020-01-01T00:00:00Z", "expired"), + (None, datetime(2020, 1, 1, tzinfo=timezone.utc), "expired"), + (False, None, "active"), + (None, "2999-01-01T00:00:00Z", "active"), + ], +) +async def test_info_key_fn_reports_live_key_status(monkeypatch, blocked, expires, expected_status): + from litellm.proxy.management_endpoints.key_management_endpoints import info_key_fn + + mock_prisma_client = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + live_row = MagicMock(spec=LiteLLM_VerificationToken) + live_row.model_dump.return_value = { + "token": "hashed_live", + "user_id": "user-x", + "team_id": None, + "object_permission_id": None, + "blocked": blocked, + "expires": expires, + } + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=live_row) + + result = await info_key_fn( + key="hashed_live", + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin"), + ) + + assert result["info"]["status"] == expected_status + mock_prisma_client.db.litellm_deletedverificationtoken.find_first.assert_not_called() + + @pytest.mark.asyncio async def test_list_keys_non_admin_user_id_auto_set(): """ diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index ebbedc6541e..6478b18e553 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -14651,3 +14651,246 @@ async def test_team_info_reports_parent_organization_models_only_to_team_manager ) assert response["team_info"].organization_models == expected_models + + +_EXISTING_TEAM_MODEL_CAPS: Final = { + "gpt-4o": {"max_budget": 10.0, "budget_duration": "1d"}, + "claude-sonnet-4-6": {"max_budget": 5.0, "budget_duration": "7d"}, +} + + +@pytest.mark.parametrize( + "requested", + [ + {**_EXISTING_TEAM_MODEL_CAPS, "gpt-4o": {"max_budget": 20.0, "budget_duration": "1d"}}, + {**_EXISTING_TEAM_MODEL_CAPS, "gpt-4o": {"max_budget": 10.0, "budget_duration": "30d"}}, + {**_EXISTING_TEAM_MODEL_CAPS, "gpt-4o": {"budget_duration": "1d"}}, + {"claude-sonnet-4-6": _EXISTING_TEAM_MODEL_CAPS["claude-sonnet-4-6"]}, + {}, + None, + {**_EXISTING_TEAM_MODEL_CAPS, "openai/gpt-4o": {"max_budget": 1000.0, "budget_duration": "1d"}}, + {**_EXISTING_TEAM_MODEL_CAPS, "openai/gpt-4o": {"max_budget": 10.0, "budget_duration": "30d"}}, + {**_EXISTING_TEAM_MODEL_CAPS, "anthropic/claude-sonnet-4-6": {"budget_duration": "7d"}}, + ], + ids=[ + "raise", + "change_duration", + "drop_cap_value", + "remove_model", + "clear_all", + "clear_with_null", + "raise_via_provider_alias", + "rewindow_via_provider_alias", + "uncap_via_provider_alias", + ], +) +def test_team_admin_cannot_loosen_team_model_caps(requested) -> None: + from litellm.proxy.management_endpoints.team_endpoints import _check_team_model_budget_update_authority + + with pytest.raises(HTTPException) as exc: + _check_team_model_budget_update_authority( + data=UpdateTeamRequest(team_id="t1", model_max_budget=requested), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="team-admin"), + existing_model_max_budget=_EXISTING_TEAM_MODEL_CAPS, + ) + assert exc.value.status_code == 403 + assert "proxy admin" in exc.value.detail["error"] + + +@pytest.mark.parametrize( + "requested", + [ + {**_EXISTING_TEAM_MODEL_CAPS, "gpt-4o": {"max_budget": 2.0, "budget_duration": "1d"}}, + {**_EXISTING_TEAM_MODEL_CAPS, "gpt-4o-mini": {"max_budget": 1.0, "budget_duration": "1d"}}, + dict(_EXISTING_TEAM_MODEL_CAPS), + {**_EXISTING_TEAM_MODEL_CAPS, "openai/gpt-4o": {"max_budget": 2.0, "budget_duration": "1d"}}, + ], + ids=["lower", "add_model", "unchanged", "tighten_via_provider_alias"], +) +def test_team_admin_can_tighten_or_keep_team_model_caps(requested) -> None: + from litellm.proxy.management_endpoints.team_endpoints import _check_team_model_budget_update_authority + + assert ( + _check_team_model_budget_update_authority( + data=UpdateTeamRequest(team_id="t1", model_max_budget=requested), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="team-admin"), + existing_model_max_budget=_EXISTING_TEAM_MODEL_CAPS, + ) + is None + ) + + +def test_team_model_cap_authority_skips_omitted_field_malformed_rows_and_proxy_admins() -> None: + from litellm.proxy.management_endpoints.team_endpoints import _check_team_model_budget_update_authority + + team_admin = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="team-admin") + outcomes = ( + _check_team_model_budget_update_authority( + data=UpdateTeamRequest(team_id="t1", max_budget=1.0), + user_api_key_dict=team_admin, + existing_model_max_budget=_EXISTING_TEAM_MODEL_CAPS, + ), + _check_team_model_budget_update_authority( + data=UpdateTeamRequest(team_id="t1", model_max_budget={}), + user_api_key_dict=team_admin, + existing_model_max_budget={"gpt-4o": "not-a-budget", "gpt-4o-mini": {"budget_duration": "1d"}}, + ), + _check_team_model_budget_update_authority( + data=UpdateTeamRequest(team_id="t1", model_max_budget=None), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + existing_model_max_budget=_EXISTING_TEAM_MODEL_CAPS, + ), + ) + assert outcomes == (None, None, None) + + +@pytest.mark.asyncio +async def test_new_team_persists_model_max_budget(mock_db_client, mock_admin_auth): + mock_db_client.jsonify_team_object = lambda db_data: db_data + mock_db_client.get_data = AsyncMock(return_value=None) + mock_db_client.update_data = AsyncMock(return_value=MagicMock()) + mock_db_client.db = MagicMock() + mock_db_client.db.litellm_modeltable = MagicMock() + mock_db_client.db.litellm_modeltable.create = AsyncMock(return_value=MagicMock(id="model123")) + + team_create_result = MagicMock(team_id="team-model-caps") + team_create_result.model_dump.return_value = {"team_id": "team-model-caps"} + mock_team_create = AsyncMock(return_value=team_create_result) + mock_db_client.db.litellm_teamtable = MagicMock() + mock_db_client.db.litellm_teamtable.create = mock_team_create + _wire_team_create_tx(mock_db_client) + mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0) + mock_db_client.db.litellm_teamtable.update = AsyncMock(return_value=team_create_result) + mock_db_client.db.litellm_usertable = MagicMock() + mock_db_client.db.litellm_usertable.update = AsyncMock(return_value=MagicMock()) + + from fastapi import Request + + from litellm.proxy._types import NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import new_team + + with patch("litellm.proxy.proxy_server.premium_user", True): # test-quality-ok: proxy_server module global is the endpoint's only injection point + await new_team( + data=NewTeamRequest( + team_alias="model-caps", + model_max_budget={"gpt-4o": {"max_budget": 10.0, "budget_duration": "1d"}}, + ), + http_request=MagicMock(spec=Request), + user_api_key_dict=mock_admin_auth, + ) + + team_data = mock_team_create.call_args.kwargs["data"] + assert team_data["model_max_budget"] == { + "gpt-4o": {"max_budget": 10.0, "budget_duration": "1d", "tpm_limit": None, "rpm_limit": None} + } + + +@pytest.mark.asyncio +async def test_new_team_rejects_unenforceable_model_max_budget(mock_db_client, mock_admin_auth): + from fastapi import Request + + from litellm.proxy._types import NewTeamRequest, ProxyException + from litellm.proxy.management_endpoints.team_endpoints import new_team + + mock_db_client.db.litellm_teamtable.create = AsyncMock() + + with patch("litellm.proxy.proxy_server.premium_user", True), pytest.raises(ProxyException) as exc: # test-quality-ok: proxy_server module global is the endpoint's only injection point + await new_team( + data=NewTeamRequest(team_alias="model-caps", model_max_budget={"gpt-4o": {"max_budget": 10.0}}), + http_request=MagicMock(spec=Request), + user_api_key_dict=mock_admin_auth, + ) + + assert exc.value.code == "400" + assert "budget_duration" in str(exc.value.message) + mock_db_client.db.litellm_teamtable.create.assert_not_awaited() + + +def _existing_team_with_model_caps(caps): + existing = MagicMock() + existing.team_id = "standalone-team-123" + existing.organization_id = None + existing.max_budget = None + existing.model_id = None + existing.model_max_budget = caps + existing.model_dump.return_value = { + "team_id": "standalone-team-123", + "organization_id": None, + "model_max_budget": caps, + "members_with_roles": [{"user_id": "team-admin-model-caps", "role": "admin"}], + } + return existing + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cleared_with", [{}, None], ids=["empty_mapping", "null"]) +async def test_update_team_clearing_model_max_budget_writes_an_empty_mapping( + disable_audit_logging_for_mocked_team, cleared_with +): + from fastapi import Request + + from litellm.proxy.management_endpoints.team_endpoints import update_team + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: proxy_server module global is the endpoint's only injection point + ): + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( + return_value=_existing_team_with_model_caps(_EXISTING_TEAM_MODEL_CAPS) + ) + mock_prisma.jsonify_team_object = lambda db_data: db_data + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_set_cache = AsyncMock() + updated = _existing_team_with_model_caps({}) + updated.litellm_model_table = None + mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=updated) + + await update_team( + data=UpdateTeamRequest(team_id="standalone-team-123", model_max_budget=cleared_with), + http_request=MagicMock(spec=Request), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"), + ) + + assert mock_prisma.db.litellm_teamtable.update.call_args.kwargs["data"]["model_max_budget"] == {} + + +@pytest.mark.asyncio +async def test_update_team_model_max_budget_raise_blocked_for_team_admin(): + from fastapi import Request + + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.team_endpoints import update_team + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch("litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock()), # test-quality-ok: stubs the audit write so the test observes only the team update result + ): + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( + return_value=_existing_team_with_model_caps(_EXISTING_TEAM_MODEL_CAPS) + ) + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_prisma.db.litellm_teamtable.update = AsyncMock() + + with pytest.raises(ProxyException) as exc: + await update_team( + data=UpdateTeamRequest( + team_id="standalone-team-123", + model_max_budget={ + **_EXISTING_TEAM_MODEL_CAPS, + "gpt-4o": {"max_budget": 100.0, "budget_duration": "1d"}, + }, + ), + http_request=MagicMock(spec=Request), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="team-admin-model-caps", models=[] + ), + ) + + assert exc.value.code == "403" + assert "proxy admin" in str(exc.value.message).lower() + mock_prisma.db.litellm_teamtable.update.assert_not_awaited() diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index 2d7397594aa..ba8b5fa3ac4 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -2541,7 +2541,7 @@ class TestRecordPartialUsageForFailure: function_id="test-partial-usage-failure", ) - def _interrupted_chunks(self): + def _interrupted_chunks(self, *, model: str = "claude-sonnet-5"): return [ self._sse( "message_start", @@ -2551,7 +2551,7 @@ class TestRecordPartialUsageForFailure: "id": "msg_abc", "type": "message", "role": "assistant", - "model": "claude-sonnet-5", + "model": model, "content": [], "stop_reason": None, "stop_sequence": None, @@ -2588,7 +2588,7 @@ class TestRecordPartialUsageForFailure: AnthropicPassthroughLoggingHandler.record_partial_usage_for_failure( litellm_logging_obj=logging_obj, request_body={"model": "claude-unpriced-test-model", "stream": True}, - all_chunks=self._interrupted_chunks(), + all_chunks=self._interrupted_chunks(model="claude-unpriced-test-model"), ) usage = logging_obj.model_call_details["combined_usage_object"] diff --git a/tests/test_litellm/proxy/spend_tracking/test_carried_budget_state.py b/tests/test_litellm/proxy/spend_tracking/test_carried_budget_state.py index fe852be775c..0bdf43b396c 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_carried_budget_state.py +++ b/tests/test_litellm/proxy/spend_tracking/test_carried_budget_state.py @@ -47,6 +47,19 @@ def test_team_and_user_state_round_trips_through_metadata(): ) +def test_team_model_max_budget_rides_on_the_token(): + """The team's per-model caps must reach the token, or the auth check and the spend hook never see them.""" + token = UserAPIKeyAuth(token="hashed", team_id="t1") + team_model_max_budget = {"gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}} + carry_team_and_user_budget_state( + valid_token=token, + team_object=LiteLLM_TeamTable(team_id="t1", model_max_budget=team_model_max_budget), + user_object=None, + ) + + assert token.team_model_max_budget == team_model_max_budget + + def test_missing_objects_leave_no_metadata_and_no_snapshot(): token = UserAPIKeyAuth(token="hashed", team_id="t1", user_id="u1") carry_team_and_user_budget_state(valid_token=token, team_object=None, user_object=None) diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 72668dd3528..f4490519554 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -44,7 +44,11 @@ from litellm.litellm_core_utils.get_provider_specific_headers import ( from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( TRUSTED_CALLBACK_VARS_FIELD, ) -from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY, SESSION_ID_OMITTED_METADATA_KEY +from litellm.constants import ( + ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, + SESSION_ID_GENERATED_METADATA_KEY, + SESSION_ID_OMITTED_METADATA_KEY, +) from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.fireworks_ai.common_utils import get_fireworks_session_id from litellm.types.utils import CredentialItem @@ -7513,6 +7517,7 @@ _PLANTED_STAMPS = { "original_model_group": "spoofed-group", "request_retry_count": -100, "_client_output_ceiling": {"api_base": "https://attacker.example"}, + ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: 10**9, "client_key": "client_value", } @@ -7545,6 +7550,7 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_from_bo assert "original_model_group" not in updated["metadata"] assert "_client_output_ceiling" not in updated["metadata"] assert "request_retry_count" not in updated["metadata"] + assert ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY not in updated["metadata"] assert updated["metadata"]["client_key"] == "client_value" diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index d1928b9cd52..19b0e88bea7 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -13463,6 +13463,45 @@ async def test_load_config_router_authorizes_fallback_targets_against_the_callin assert router.fallback_access_check is router_fallback_access_check +@pytest.mark.asyncio +async def test_load_config_router_budget_checks_fallback_targets_against_the_calling_key(tmp_path, monkeypatch): + """A config-loaded router refuses a paid fallback target for an over-budget caller.""" + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.proxy_server import ProxyConfig + + config_file = tmp_path / "config.yaml" + config_file.write_text( + yaml.dump({"model_list": [{"model_name": "m", "litellm_params": {"model": "openai/m", "api_key": "k"}}]}) + ) + + router, _, _ = await ProxyConfig().load_config(router=None, config_file_path=str(config_file)) + + over_budget = { + "metadata": { + "user_api_key_auth": UserAPIKeyAuth( + api_key="hashed", token="hashed", user_id="u1", user_spend=99.0, user_max_budget=1.0 + ) + } + } + under_budget = { + "metadata": { + "user_api_key_auth": UserAPIKeyAuth( + api_key="hashed", token="hashed", user_id="u1", user_spend=0.0, user_max_budget=100.0 + ) + } + } + + # on by default: an over-budget caller is refused the paid fallback with no config at all + monkeypatch.setattr(proxy_server, "general_settings", {}, raising=False) + assert await router.fallback_budget_check(model="m", request_kwargs=over_budget, llm_router=router) is False + assert await router.fallback_budget_check(model="m", request_kwargs=under_budget, llm_router=router) is True + + # explicit opt-out restores the unguarded behaviour + monkeypatch.setattr(proxy_server, "general_settings", {"enforce_fallback_budget": False}, raising=False) + assert await router.fallback_budget_check(model="m", request_kwargs=over_budget, llm_router=router) is True + + @pytest.mark.asyncio async def test_load_config_user_api_key_cache_max_size_keeps_more_than_200_entries(tmp_path, monkeypatch): """The auth cache used to be pinned at InMemoryCache's 200 entry default, so a diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py index ce6ecc2ea65..672dd1eb674 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py @@ -401,19 +401,20 @@ async def test_check_view_exists_creates_token_view_when_missing( prisma_client.db.execute_raw = AsyncMock() prisma_client.health_check = AsyncMock(return_value=[{"?column?": 1}]) result = await prisma_client.check_view_exists() + created_sql = prisma_client.db.execute_raw.await_args.args[0] actual = { "result": result, "create_called": prisma_client.db.execute_raw.await_count, - "create_sql_starts_with_create_view": prisma_client.db.execute_raw.await_args.args[ - 0 - ] - .strip() - .startswith('CREATE VIEW "LiteLLM_VerificationTokenView"'), + "create_sql_starts_with_create_view": created_sql.strip().startswith( + 'CREATE VIEW "LiteLLM_VerificationTokenView"' + ), + "projects_team_model_max_budget": "t.model_max_budget AS team_model_max_budget" in created_sql, } assert actual == { "result": None, "create_called": 1, "create_sql_starts_with_create_view": True, + "projects_team_model_max_budget": True, } diff --git a/tests/test_litellm/router_utils/test_fallback_event_handlers.py b/tests/test_litellm/router_utils/test_fallback_event_handlers.py index 9318f306c89..dfe06bffd09 100644 --- a/tests/test_litellm/router_utils/test_fallback_event_handlers.py +++ b/tests/test_litellm/router_utils/test_fallback_event_handlers.py @@ -27,6 +27,7 @@ class StreamingWrapper: class FakeRouter: fallback_access_check = None + fallback_budget_check = None def log_retry(self, kwargs, e): return kwargs @@ -37,6 +38,7 @@ class FakeRouter: class AlwaysFailRouter: fallback_access_check = None + fallback_budget_check = None def log_retry(self, kwargs, e): return kwargs @@ -101,6 +103,7 @@ async def test_run_async_fallback_raises_when_all_fallbacks_fail(): class RecordingRouter: fallback_access_check = None + fallback_budget_check = None def __init__(self): self.received_kwargs = None @@ -162,6 +165,7 @@ async def test_run_async_fallback_skips_original_model_group(): class AttemptRecordingRouter: fallback_access_check = None + fallback_budget_check = None def __init__(self): self.attempted_model_groups = [] @@ -471,6 +475,8 @@ class AccessCheckedRouter(AttemptRecordingRouter): self.allowed_models = allowed_models self.access_checks = [] + fallback_budget_check = None + async def fallback_access_check(self, *, model, request_kwargs, llm_router): self.access_checks.append((model, request_kwargs["metadata"]["user_api_key"], llm_router is self)) return model in self.allowed_models @@ -542,6 +548,7 @@ async def test_run_async_fallback_does_not_consult_access_check_for_same_model_g class RecordingFailRouter: fallback_access_check = None + fallback_budget_check = None def __init__(self): self.attempted_models = [] @@ -1053,6 +1060,7 @@ class TestTriggerCooldownForFailedDeployment: class TestRunAsyncFallbackTriggersCooldown: class RouterWithLoggingKwarg: fallback_access_check = None + fallback_budget_check = None def __init__(self): self.cooldown_time = 60.0 diff --git a/tests/test_litellm/test_vertex_ai_xai_grok_prompt_caching_metadata.py b/tests/test_litellm/test_vertex_ai_xai_grok_prompt_caching_metadata.py new file mode 100644 index 00000000000..72e98711f0c --- /dev/null +++ b/tests/test_litellm/test_vertex_ai_xai_grok_prompt_caching_metadata.py @@ -0,0 +1,38 @@ +from typing import Final + +import pytest + +import litellm +from litellm import get_model_info +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.utils import supports_prompt_caching + +MODEL: Final = "vertex_ai/xai/grok-4.6" +GROK_KEY_PREFIXES: Final = ("vertex_ai/xai/grok-", "azure_ai/grok-", "xai/grok-") + + +@pytest.mark.usefixtures("local_model_cost_map") +def test_grok_models_with_cache_read_price_advertise_prompt_caching() -> None: + cached_grok_models = tuple( + key + for key, entry in litellm.model_cost.items() + if key.startswith(GROK_KEY_PREFIXES) and entry.get("cache_read_input_token_cost") + ) + assert cached_grok_models, "expected at least one grok model with a cache read price" + + missing_flag = tuple(key for key in cached_grok_models if supports_prompt_caching(model=key) is not True) + assert missing_flag == (), ( + f"grok models with cache_read_input_token_cost fail supports_prompt_caching: {missing_flag}" + ) + + +@pytest.mark.usefixtures("local_model_cost_map") +def test_vertex_ai_grok_4_6_supports_prompt_caching_via_get_model_info() -> None: + routed_model, provider, _, _ = get_llm_provider(model=MODEL) + assert (routed_model, provider) == ("xai/grok-4.6", "vertex_ai") + + info = get_model_info(model=routed_model, custom_llm_provider=provider) + assert info["litellm_provider"] == "vertex_ai" + assert info.get("supports_prompt_caching") is True + + assert supports_prompt_caching(model=MODEL) is True diff --git a/tests/test_litellm_rust/ocr/test_requests.py b/tests/test_litellm_rust/ocr/test_requests.py index 4f4b39fa6c6..58bb6a77537 100644 --- a/tests/test_litellm_rust/ocr/test_requests.py +++ b/tests/test_litellm_rust/ocr/test_requests.py @@ -518,32 +518,34 @@ async def test_native_ocr_inherits_named_credentials_without_overwriting_argumen assert ocr_server.requests[0].body["pages"] == [0, 2] -@pytest.mark.parametrize("source", ["sdk", "proxy"]) @pytest.mark.parametrize( - "filename,mime", [("scan.PNG", "image/png"), ("document.pdf", "application/pdf"), ("note.txt", "text/plain")] + "filename,field,mime", + [("scan.PNG", "image_url", "image/png"), ("document.pdf", "document_url", "application/pdf")], ) -def test_ocr_file_helpers_use_native_document_preparation(source: str, filename: str, mime: str) -> None: +def test_native_ocr_infers_mime_type_from_reader_name( + ocr_server: RecordingServer, filename: str, field: str, mime: str +) -> None: from io import BytesIO - from litellm.ocr.input import convert_file_document_to_url_document, get_mime_type - from litellm.proxy.ocr_endpoints.endpoints import _build_document_from_upload - file: Final = BytesIO(b"abc") file.name = filename - document: Final = ( - convert_file_document_to_url_document({"type": "file", "file": file}) - if source == "sdk" - else _build_document_from_upload(b"abc", filename, "application/octet-stream; charset=utf-8") - ) - field: Final = "image_url" if mime.startswith("image/") else "document_url" - assert get_mime_type(filename) == mime - assert document == {"type": field, field: f"data:{mime};base64,YWJj"} + call_native_ocr(ocr_server, document={"type": "file", "file": file}) + assert ocr_server.requests[0].body["document"] == {"type": field, field: f"data:{mime};base64,YWJj"} + + +def test_native_ocr_encodes_str_reader_results_as_utf8(ocr_server: RecordingServer) -> None: + from io import StringIO + + call_native_ocr(ocr_server, document={"type": "file", "file": StringIO("abc"), "mime_type": "text/plain"}) + assert ocr_server.requests[0].body["document"] == { + "type": "document_url", + "document_url": "data:text/plain;base64,YWJj", + } @pytest.mark.parametrize("attribute", ["read", "name"]) -def test_native_file_preparation_preserves_property_errors(attribute: str) -> None: - from litellm.ocr.input import convert_file_document_to_url_document - +def test_native_file_preparation_preserves_property_errors(ocr_server: RecordingServer, attribute: str) -> None: + ocr_server.expected_requests = 0 failure: Final = LookupError("file property failed") class File: @@ -555,16 +557,47 @@ def test_native_file_preparation_preserves_property_errors(attribute: str) -> No def read(self): return b"abc" - with pytest.raises(LookupError) as caught: - convert_file_document_to_url_document({"type": "file", "file": File()}) - assert caught.value is failure + with pytest.raises(litellm.APIConnectionError, match="file property failed") as caught: + call_native_ocr(ocr_server, document={"type": "file", "file": File()}) + assert caught.value.__context__ is failure + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_native_file_preparation_preserves_reader_exception( + ocr_server: RecordingServer, asynchronous: bool +) -> None: + ocr_server.expected_requests = 0 + failure: Final = RuntimeError("reader failed") + + class Reader: + def read(self) -> bytes: + raise failure + + document: Final = {"type": "file", "file": Reader()} + with pytest.raises(litellm.APIConnectionError, match="reader failed") as caught: + await call_native_aocr(ocr_server, document=document) if asynchronous else call_native_ocr( + ocr_server, document=document + ) + assert caught.value.__context__ is failure + + +def test_native_file_preparation_rejects_unsupported_reader_results(ocr_server: RecordingServer) -> None: + ocr_server.expected_requests = 0 + + class Reader: + def read(self) -> int: + return 1 + + with pytest.raises(litellm.APIConnectionError, match="bytes or str") as caught: + call_native_ocr(ocr_server, document={"type": "file", "file": Reader()}) + assert isinstance(caught.value.__context__, TypeError) @pytest.mark.parametrize("kind", ["bytes", "path", "reader"]) -def test_native_file_preparation_rejects_oversized_input(kind: str, tmp_path: Path) -> None: - from litellm.ocr.input import FileDocument, convert_file_document_to_url_document, get_max_file_bytes - - limit: Final = get_max_file_bytes() +def test_native_file_preparation_rejects_oversized_input(ocr_server: RecordingServer, kind: str, tmp_path: Path) -> None: + ocr_server.expected_requests = 0 + limit: Final = 50 * 1024 * 1024 path: Final = tmp_path / "large.pdf" with path.open("wb") as stream: stream.truncate(limit + 1) @@ -573,53 +606,25 @@ def test_native_file_preparation_rejects_oversized_input(kind: str, tmp_path: Pa def read(self) -> bytes: return b"a" * (limit + 1) - document: Final[FileDocument] = { + document: Final = { "type": "file", "file": path if kind == "path" else Reader() if kind == "reader" else b"a" * (limit + 1), } - with pytest.raises(ValueError, match="exceeds the size limit"): - convert_file_document_to_url_document(document) + with pytest.raises(litellm.BadRequestError, match="exceeds the size limit"): + call_native_ocr(ocr_server, document=document) -@pytest.mark.parametrize("kind", ["str", "path", "reader"]) -def test_native_upload_binding_rejects_filesystem_inputs(kind: str, tmp_path: Path) -> None: +def test_native_file_preparation_reports_missing_paths(ocr_server: RecordingServer, tmp_path: Path) -> None: + ocr_server.expected_requests = 0 + missing: Final = tmp_path / "missing.pdf" + with pytest.raises(litellm.APIConnectionError, match=f"File not found: {missing}") as caught: + call_native_ocr(ocr_server, document={"type": "file", "file": missing}) + assert isinstance(caught.value.__context__, FileNotFoundError) + + +def test_native_file_preparation_rejects_empty_readers(ocr_server: RecordingServer) -> None: from io import BytesIO - from typing import cast # noqa: TID251 # deliberately invalid inputs exercise the native runtime boundary - from litellm.ocr.input import convert_upload_to_url_document - - path: Final = tmp_path / "secret.pdf" - path.write_bytes(b"server secret") - source: Final = str(path) if kind == "str" else path if kind == "path" else BytesIO(b"abc") - with pytest.raises(TypeError): - convert_upload_to_url_document(cast(bytes, source), "document.pdf", None) - - -@pytest.mark.parametrize("extra_bytes", [0, 1]) -def test_native_upload_enforces_file_size_limit(extra_bytes: int) -> None: - import base64 - - from litellm.ocr.input import convert_upload_to_url_document, get_max_file_bytes - - content: Final = b"a" * (get_max_file_bytes() + extra_bytes) - if extra_bytes: - with pytest.raises(ValueError, match="exceeds the size limit"): - convert_upload_to_url_document(content, "scan.pdf", None) - return - document: Final = convert_upload_to_url_document(content, "scan.pdf", None) - assert document["type"] == "document_url" - assert base64.b64decode(document["document_url"].split(",", 1)[1]) == content - - -def test_native_file_preparation_preserves_reader_exception() -> None: - from litellm.ocr.input import convert_file_document_to_url_document - - failure: Final = RuntimeError("reader failed") - - class Reader: - def read(self) -> bytes: - raise failure - - with pytest.raises(RuntimeError) as caught: - convert_file_document_to_url_document({"type": "file", "file": Reader()}) - assert caught.value is failure + ocr_server.expected_requests = 0 + with pytest.raises(litellm.BadRequestError, match="File is empty"): + call_native_ocr(ocr_server, document={"type": "file", "file": BytesIO(b"")}) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/add_plugin_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/add_plugin_form.tsx index 5627f402a65..0c114dce700 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/add_plugin_form.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/skills/_components/add_plugin_form.tsx @@ -232,7 +232,7 @@ const AddPluginForm: React.FC = ({ visible, onClose, accessT name="skillUrl" label={labelWithHint( "Source URL", - "Paste an HTTPS git repository URL from GitHub, GitLab, Bitbucket, or a self-hosted host (e.g. github.com/org/repo or github.com/org/repo/tree/main/my-skill), or an HTTPS link to a .zip archive of the skill hosted on S3 or any static file server.", + "Paste an HTTPS git repository URL from GitHub, GitLab, Bitbucket, or a self-hosted host (e.g. github.com/org/repo or github.com/org/repo/tree/main/my-skill), or an HTTPS link to a .zip archive of the skill hosted on S3 or any static file server. For a private repository use its SSH clone URL (git@ghe.example.com:org/repo.git) so Claude Code clones it with your own SSH key.", )} > {({ ref, onChange, ...field }) => ( diff --git a/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysPage.test.tsx b/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysPage.test.tsx index 31cd407a5e6..cf0c13ee152 100644 --- a/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysPage.test.tsx +++ b/ui/litellm-dashboard/src/components/DeletedKeysPage/DeletedKeysPage.test.tsx @@ -22,6 +22,7 @@ const mockDeletedKey: DeletedKeyResponse = { key_name: "test-key", key_alias: "Test Key Alias", spend: 5.5, + total_spend: 5.5, max_budget: 100, expires: "2024-12-31T23:59:59Z", models: ["gpt-3.5-turbo"], diff --git a/ui/litellm-dashboard/src/components/Teams.test.tsx b/ui/litellm-dashboard/src/components/Teams.test.tsx index 851c9e6d487..f2d7cff2ec4 100644 --- a/ui/litellm-dashboard/src/components/Teams.test.tsx +++ b/ui/litellm-dashboard/src/components/Teams.test.tsx @@ -7,6 +7,7 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import { useTeamMetadataSchema } from "@/app/(dashboard)/hooks/teams/useTeamMetadataSchema"; import { toast } from "@/lib/toast"; import { fetchAvailableModelsForTeamOrKey } from "./key_team_helpers/fetch_available_models_team_key"; +import { MODEL_MAX_BUDGET_PREMIUM_HINT } from "./key_team_helpers/ModelMaxBudgetEditor"; import { fetchMCPAccessGroups, getDefaultTeamSettings, @@ -1547,6 +1548,36 @@ describe("Teams - the exact bytes the create call sends", () => { expect(await screen.findByText("Please input a team name")).toBeInTheDocument(); expect(teamCreateCall).not.toHaveBeenCalled(); }); + + it("locks the per-model budget editor and says why when the proxy has no enterprise license", async () => { + await openCreateModal({ premiumUser: false }); + + expect(screen.getByRole("button", { name: /Add Model Budget/i })).toBeDisabled(); + expect(screen.getByText(MODEL_MAX_BUDGET_PREMIUM_HINT)).toBeInTheDocument(); + }); + + it("sends the per-model budget a licensed operator fills in, keyed by model", async () => { + const user = userEvent.setup({ delay: null }); + await openCreateModal({ premiumUser: true }); + + await user.click(screen.getByRole("button", { name: /Add Model Budget/i })); + await chooseSelectOption(user, screen.getByPlaceholderText("Select model"), "gpt-4"); + fireEvent.change(screen.getByPlaceholderText("Max spend ($)"), { target: { value: "3" } }); + + const payload = await submit(); + + expect(payload.model_max_budget).toStrictEqual({ "gpt-4": { budget_limit: 3, time_period: "30d" } }); + }); + + it("leaves model_max_budget out when a started row is removed again", async () => { + const user = userEvent.setup({ delay: null }); + await openCreateModal({ premiumUser: true }); + + await user.click(screen.getByRole("button", { name: /Add Model Budget/i })); + await user.click(screen.getByRole("button", { name: "Remove model budget" })); + + expect(wireBody(await submit())).not.toHaveProperty("model_max_budget"); + }); }); describe("Teams - the create form keeps the organization and models picks while it is open", () => { diff --git a/ui/litellm-dashboard/src/components/Teams.tsx b/ui/litellm-dashboard/src/components/Teams.tsx index 4f3367d8b98..7214d16f665 100644 --- a/ui/litellm-dashboard/src/components/Teams.tsx +++ b/ui/litellm-dashboard/src/components/Teams.tsx @@ -48,6 +48,7 @@ import BudgetDurationDropdown, { } from "./common_components/budget_duration_dropdown"; import { Organization, getDefaultTeamSettings, getGuardrailsList, getPoliciesList, teamDeleteCall } from "./networking"; import NumericalInput from "./shared/numerical_input"; +import { ModelMaxBudget, ModelMaxBudgetField } from "./key_team_helpers/ModelMaxBudgetEditor"; import VectorStoreSelector from "./vector_store_management/VectorStoreSelector"; import SearchToolSelector from "./search_tools/SearchToolSelector"; import SkillSelector from "./skills/SkillSelector"; @@ -271,6 +272,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser const [policiesList, setPoliciesList] = useState([]); const [loggingSettings, setLoggingSettings] = useState([]); const [modelAliases, setModelAliases] = useState<{ [key: string]: string }>({}); + const [modelMaxBudget, setModelMaxBudget] = useState({}); const [routerSettings, setRouterSettings] = useState(null); const [routerSettingsKey, setRouterSettingsKey] = useState(0); @@ -348,6 +350,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser setSearchToolSettingsOpen(false); setLoggingSettings([]); setModelAliases({}); + setModelMaxBudget({}); setRouterSettings(null); setRouterSettingsKey((prev) => prev + 1); }; @@ -525,6 +528,10 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser formValues.model_aliases = modelAliases; } + if (Object.keys(modelMaxBudget).length > 0) { + formValues.model_max_budget = modelMaxBudget; + } + // Add router_settings if any are defined if (routerSettings?.router_settings) { // Only include router_settings if it has at least one non-null value @@ -813,6 +820,14 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser /> )} + {({ ref, value, ...field }) => ( diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx index 617b9209a41..6c742309eb7 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -79,6 +79,7 @@ const mockKey: KeyResponse = { key_name: "test-key", key_alias: "Test Key Alias", spend: 5.5, + total_spend: 42.25, max_budget: 100, expires: "2999-12-31T23:59:59Z", models: ["gpt-3.5-turbo", "gpt-4"], @@ -236,6 +237,14 @@ it("should display key information correctly", async () => { }); }); +it("shows lifetime spend in its own column next to the period spend meter", async () => { + renderWithProviders(); + + expect(await screen.findByText("Lifetime Spend")).toBeInTheDocument(); + expect(screen.getByText("$42.2500")).toBeInTheDocument(); + expect(screen.getByText("$5.5000")).toBeInTheDocument(); +}); + it("should display user email correctly", async () => { renderWithProviders(); @@ -638,6 +647,23 @@ describe("server-side filtering – the LIT-4080 regression guard", () => { }); }); + it("threads the Status drawer filter into the useKeys query and the URL", async () => { + const onUrlUpdate = vi.fn(); + renderWithProviders(, { onUrlUpdate }); + + openFilters(); + const user = userEvent.setup(); + await chooseSelectOption(user, await screen.findByRole("combobox", { name: "Status" }), "Revoked (blocked)"); + fireEvent.click(screen.getByTestId("filter-drawer-apply")); + + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ status: "revoked" })); + }); + await waitFor(() => { + expect(lastSearchParam(onUrlUpdate, "filter_status")).toBe("revoked"); + }); + }); + it("sends the search box as the combined alias-or-ID search rather than the key-alias filter", async () => { renderWithProviders(); @@ -745,6 +771,25 @@ describe("Status column reflects blocked / expiry / scim metadata", () => { expect(screen.queryByText(/Blocked by SCIM/i)).not.toBeInTheDocument(); }); + it("renders Deleted for an archived key, even when the archived row was also blocked", async () => { + mockUseKeys.mockReturnValue( + keysResult([ + { ...mockKey, blocked: true, metadata: {}, deleted_at: "2024-11-15T10:00:00Z", deleted_by: "admin-1" }, + ]), + ); + + renderWithProviders(); + + const tag = await screen.findByTestId(`key-status-${mockKey.token_id}`); + expect(tag).toHaveTextContent("Deleted"); + + const user = userEvent.setup(); + await user.hover(tag); + await waitFor(() => { + expect(screen.getByText(/by admin-1/)).toBeInTheDocument(); + }); + }); + it("marks a SCIM-blocked key with the SCIM tooltip reason", async () => { mockUseKeys.mockReturnValue(keysResult([{ ...mockKey, blocked: true, metadata: { scim_blocked: true } }])); @@ -790,6 +835,24 @@ describe("table state lives in the URL so it survives leaving and returning to t expect(screen.getByTestId("filter-chip-team_id")).toHaveTextContent("Test Team"); }); + it("restores the status filter from the URL and sends it to /key/list", async () => { + renderWithProviders(, { searchParams: { filter_status: "deleted" } }); + + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ status: "deleted" })); + }); + expect(screen.getByTestId("filter-chip-status")).toHaveTextContent("Deleted"); + }); + + it("ignores a hand-edited status the backend would reject instead of 400ing the page", async () => { + renderWithProviders(, { searchParams: { filter_status: "bogus" } }); + + await waitFor(() => { + expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ status: undefined })); + }); + expect(screen.queryByTestId("filter-chip-status")).not.toBeInTheDocument(); + }); + it("writes the search term to the URL", async () => { const onUrlUpdate = vi.fn(); renderWithProviders(, { onUrlUpdate }); diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx index 907ee28bd05..1f52bdd7335 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx @@ -14,6 +14,7 @@ import { import { SearchSelect } from "@/components/shared/SearchSelect"; import { PageHeader } from "@/components/shared/PageHeader"; import { Input } from "@/components/ui/input"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { useDebouncedValue } from "@tanstack/react-pacer/debouncer"; import { ColumnFiltersState, functionalUpdate, OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table"; import { KeyRound } from "lucide-react"; @@ -28,7 +29,7 @@ interface VirtualKeysTableProps { headerActions?: React.ReactNode; } -const FILTER_COLUMNS = ["team_id", "org_id", "user_id", "key_hash"] as const; +const FILTER_COLUMNS = ["team_id", "org_id", "user_id", "key_hash", "status"] as const; type FilterColumn = (typeof FILTER_COLUMNS)[number]; const FILTER_LABELS: Record = { @@ -36,8 +37,28 @@ const FILTER_LABELS: Record = { org_id: "Organization", user_id: "User ID", key_hash: "Key ID", + status: "Status", }; +const KEY_STATUS_VALUES = ["active", "expired", "revoked", "deleted"] as const; +type KeyStatusFilter = (typeof KEY_STATUS_VALUES)[number]; +const ALL_STATUSES = "all"; + +const KEY_STATUS_LABELS: Record = { + active: "Active", + expired: "Expired", + revoked: "Revoked (blocked)", + deleted: "Deleted", +}; + +const STATUS_FILTER_ITEMS = [ + { value: ALL_STATUSES, label: "All statuses" }, + ...KEY_STATUS_VALUES.map((value) => ({ value, label: KEY_STATUS_LABELS[value] })), +]; + +const isKeyStatusFilter = (value: string): value is KeyStatusFilter => + (KEY_STATUS_VALUES as readonly string[]).includes(value); + const DEFAULT_SORT_BY = "created_at"; const DEFAULT_SORT_ORDER = "desc"; const DEFAULT_PAGE_SIZE = 50; @@ -65,6 +86,7 @@ const TABLE_STATE = { filter_org: parseAsString.withDefault(""), filter_user: parseAsString.withDefault(""), filter_key_id: parseAsString.withDefault(""), + filter_status: parseAsString.withDefault(""), }; const toSortOrder = (active: SortingState[number]): "asc" | "desc" => (active.desc ? "desc" : "asc"); @@ -96,15 +118,16 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { () => ({ pageIndex: tableState.page - 1, pageSize: tableState.page_size }), [tableState.page, tableState.page_size], ); - const { filter_team, filter_org, filter_user, filter_key_id } = tableState; + const { filter_team, filter_org, filter_user, filter_key_id, filter_status } = tableState; const appliedFilters = useMemo( () => ({ team_id: filter_team.trim(), org_id: filter_org.trim(), user_id: filter_user.trim(), key_hash: filter_key_id.trim(), + status: isKeyStatusFilter(filter_status) ? filter_status : "", }), - [filter_team, filter_org, filter_user, filter_key_id], + [filter_team, filter_org, filter_user, filter_key_id, filter_status], ); const columnFilters = useMemo( () => @@ -121,6 +144,7 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { search: searchQuery.trim() || undefined, userID: appliedFilters.user_id || undefined, keyHash: appliedFilters.key_hash || undefined, + status: appliedFilters.status || undefined, sortBy, sortOrder: tableState.sort_order, expand: "user", @@ -164,6 +188,7 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { filter_org: filterValue(next, "org_id"), filter_user: filterValue(next, "user_id"), filter_key_id: filterValue(next, "key_hash"), + filter_status: filterValue(next, "status"), page: null, }; void setTableState(nextFilters); @@ -233,6 +258,9 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { if (columnId === "org_id") { return organizations.find((org) => org.organization_id === raw)?.organization_alias || raw; } + if (columnId === "status" && isKeyStatusFilter(raw)) { + return KEY_STATUS_LABELS[raw]; + } return raw; }, [allTeams, organizations], @@ -340,6 +368,24 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { placeholder="Enter Key ID…" /> + + + )} diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx index 6eea77ae827..c9ee0fa1126 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx @@ -13,6 +13,7 @@ import { IdCell, IdentityCell, ModelsCell, + MoneyCell, SpendBudgetCell, StatusBadge, UserPopoverCell, @@ -43,6 +44,13 @@ export const KEY_TABLE_SORT_FIELDS: readonly string[] = [ ]; const getKeyStatus = (key: KeyResponse): KeyStatus => { + if (key.deleted_at) { + return { + tone: "neutral", + label: "Deleted", + tooltip: `Deleted ${new Date(key.deleted_at).toLocaleString()}${key.deleted_by ? ` by ${key.deleted_by}` : ""}. Kept for audit and spend history; requests using this key are rejected.`, + }; + } if (key.blocked === true) { const isScimBlocked = (key.metadata as Record | null | undefined)?.scim_blocked === true; return { @@ -274,6 +282,20 @@ export const getKeyTableColumns = ({ ); }, }, + { + id: "total_spend", + accessorKey: "total_spend", + meta: { title: "Lifetime Spend" }, + header: () => ( + + ), + size: 130, + enableSorting: false, + cell: (info) => , + }, { id: "budget_reset_at", accessorKey: "budget_reset_at", diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.test.ts b/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.test.ts index 4713fcbc55b..e40e0e0c783 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.test.ts +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.test.ts @@ -156,6 +156,20 @@ describe("getSourceLink", () => { it("returns null when no repo or url", () => { expect(getSourceLink({ source: "github" })).toBeNull(); }); + + it("keeps http and upper-case https urls registered through the api clickable", () => { + expect(getSourceLink({ source: "url", url: "http://git.internal.example/org/repo" })).toBe( + "http://git.internal.example/org/repo", + ); + expect(getSourceLink({ source: "git-subdir", url: "HTTPS://gitlab.com/org/repo", path: "sub/dir" })).toBe( + "HTTPS://gitlab.com/org/repo", + ); + }); + + it("returns null for an ssh clone url, which is not browsable", () => { + expect(getSourceLink({ source: "url", url: "git@ghe.example.com:org/repo.git" })).toBeNull(); + expect(getSourceLink({ source: "url", url: "ssh://git@ghe.example.com/org/repo.git" })).toBeNull(); + }); }); describe("getCategoryBadgeColor", () => { @@ -466,6 +480,70 @@ describe("parseSkillSource", () => { expect(parseSkillSource("gitlab.com/org/repo", "a//b")).toBeNull(); }); + it("keeps an scp-style ssh clone url so private hosts authenticate with the user's key", () => { + expect(parseSkillSource("git@ghe.example.com:org/repo.git")?.parsed).toEqual({ + source: "url", + url: "git@ghe.example.com:org/repo.git", + }); + expect(parseSkillSource("git@ghe.example.com:org/repo.git")?.suggestedName).toBe("repo"); + }); + + it("stores an ssh clone url exactly as typed, so a forced .git suffix cannot break azure devops or codecommit", () => { + for (const url of [ + "git@ghe.example.com:org/repo", + "git@ssh.dev.azure.com:v3/org/project/repo", + "ssh://git@ghe.example.com/org/repo", + "ssh://apka1234@git-codecommit.us-east-1.amazonaws.com/v1/repos/my-repo", + "ssh://git@ghe.example.com:2222/org/nested/repo.git", + ]) { + expect(parseSkillSource(url)?.parsed).toEqual({ source: "url", url }); + } + expect(parseSkillSource("git@ssh.dev.azure.com:v3/org/project/repo")?.suggestedName).toBe("repo"); + }); + + it("accepts an internal host whose last label is not alphabetic, matching the https rule", () => { + expect(parseSkillSource("git@gitlab.internal.k8s2:org/repo.git")?.parsed).toEqual({ + source: "url", + url: "git@gitlab.internal.k8s2:org/repo.git", + }); + expect(parseSkillSource("https://gitlab.internal.k8s2/org/repo")?.parsed).toEqual({ + source: "url", + url: "https://gitlab.internal.k8s2/org/repo", + }); + }); + + it("combines an ssh clone url with an explicit subfolder", () => { + expect(parseSkillSource("git@ghe.example.com:org/repo.git", "plugins/my-skill")?.parsed).toEqual({ + source: "git-subdir", + url: "git@ghe.example.com:org/repo.git", + path: "plugins/my-skill", + }); + expect(parseSkillSource("git@ghe.example.com:org/repo.git", "../etc")).toBeNull(); + }); + + it("rejects ssh-looking input without a host or repo path", () => { + expect(parseSkillSource("git@ghe.example.com:repo.git")).toBeNull(); + expect(parseSkillSource("git@localhost:org/repo.git")).toBeNull(); + expect(parseSkillSource("git@:org/repo.git")).toBeNull(); + expect(parseSkillSource("ssh://ghe.example.com/org/repo.git")).toBeNull(); + }); + + it("rejects ssh remotes with ip hosts or traversal segments", () => { + expect(parseSkillSource("git@10.0.0.5:org/repo.git")).toBeNull(); + expect(parseSkillSource("ssh://git@169.254.169.254/org/repo")).toBeNull(); + expect(parseSkillSource("git@ghe.example.com:../etc")).toBeNull(); + expect(parseSkillSource("ssh://git@ghe.example.com/org/../repo")).toBeNull(); + expect(parseSkillSource("git@ghe.example.com:org/../../etc/passwd")).toBeNull(); + expect(parseSkillSource("git@ghe.example.com:org/.github")?.parsed).toEqual({ + source: "url", + url: "git@ghe.example.com:org/.github", + }); + }); + + it("rejects an ssh remote carrying a password, which would publish a secret on the feed", () => { + expect(parseSkillSource("ssh://git:s3cret@ghe.example.com/org/repo.git")).toBeNull(); + }); + it("returns null for empty and garbage input", () => { expect(parseSkillSource("")).toBeNull(); expect(parseSkillSource(" ")).toBeNull(); @@ -568,7 +646,7 @@ describe("parseSkillSource", () => { // Skill sources are served on the unauthenticated public feeds and cloned by clients, so the // parser must never publish an insecure, credentialed, internal, or malformed clone URL. describe("parseSkillSource — security boundary", () => { - it("rejects non-https schemes", () => { + it("rejects schemes other than https and user-qualified ssh", () => { for (const url of [ "http://gitlab.com/org/repo", "HTTP://gitlab.com/org/repo", diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts b/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts index 6f673d67f87..8cf620d9077 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/helpers.ts @@ -29,17 +29,34 @@ export const SHA256_REGEX = /^[0-9a-fA-F]{64}$/; export const isValidSha256 = (digest: string): boolean => digest.trim() === "" || SHA256_REGEX.test(digest.trim()); -// WHATWG normalizes obfuscated IPv4 (e.g. 2130706433, 0x7f.0.0.1) to dotted-decimal, so this -// catches every IPv4 form; bracketed IPv6 is rejected separately. +// WHATWG normalizes obfuscated IPv4 (e.g. 2130706433, 0x7f.0.0.1) to dotted-decimal on https, so +// this catches every IPv4 form there; on a non-special scheme like ssh it catches the dotted form +// only. Bracketed IPv6 is rejected separately. const IPV4_HOST_REGEX = /^\d{1,3}(\.\d{1,3}){3}$/; const GITHUB_ORG_REGEX = /^[A-Za-z0-9-]+$/; const GITHUB_REPO_REGEX = /^[A-Za-z0-9._-]+$/; +const BROWSABLE_URL_REGEX = /^https?:\/\//i; +const SSH_SCHEME = "ssh://"; +const SSH_SCP_REGEX = /^([a-z0-9._-]+)@([^:/@]+):(?!\/)(.+)$/i; + const buildRepoUrl = (url: URL): string => `${url.protocol}//${url.host}${url.pathname.replace(/\/+$/, "")}`; const pathSegments = (url: URL): string[] => url.pathname.split("/").filter((seg) => seg !== ""); +const toUrl = (candidate: string): URL | null => { + try { + return new URL(candidate); + } catch { + return null; + } +}; + +/** One host rule for every scheme, so an ssh remote is neither more nor less trusted than its https twin. */ +const isSafeHost = (url: URL): boolean => + url.hostname.includes(".") && !url.hostname.startsWith("[") && !IPV4_HOST_REGEX.test(url.hostname); + /** * Validate and normalize a repository URL into a parsed URL, or null. Enforces https (rejects * http/ssh/git/etc.), rejects embedded credentials, and requires a dotted host, so the public @@ -52,20 +69,8 @@ const parseRepoUrl = (raw: string): URL | null => { return null; } const withScheme = /^[a-z][a-z0-9+.-]*:\/\//i.test(trimmed) ? trimmed : `https://${trimmed}`; - let url: URL; - try { - url = new URL(withScheme); - } catch { - return null; - } - if ( - url.protocol !== "https:" || - url.username !== "" || - url.password !== "" || - !url.hostname.includes(".") || - url.hostname.startsWith("[") || - IPV4_HOST_REGEX.test(url.hostname) - ) { + const url = toUrl(withScheme); + if (!url || url.protocol !== "https:" || url.username !== "" || url.password !== "" || !isSafeHost(url)) { return null; } return url; @@ -140,13 +145,12 @@ const parseGitHubSource = (url: URL, subPath?: string): SkillSourcePreview | nul return repoPreview; }; -const parseRawGitSource = (url: URL, subPath?: string): SkillSourcePreview | null => { - if (pathSegments(url).length < 2) { - return null; - } - - const repoUrl = buildRepoUrl(url); - +const buildGitSourcePreview = ( + kind: "Git" | "SSH", + repoUrl: string, + repoName: string, + subPath?: string, +): SkillSourcePreview | null => { const normalized = normalizeSubPath(subPath ?? ""); if (normalized !== "") { if (!SUBDIR_PATH_REGEX.test(normalized)) { @@ -154,18 +158,51 @@ const parseRawGitSource = (url: URL, subPath?: string): SkillSourcePreview | nul } return { parsed: { source: "git-subdir", url: repoUrl, path: normalized }, - label: `Git subdir — ${repoUrl} @ ${normalized}`, + label: `${kind} subdir — ${repoUrl} @ ${normalized}`, suggestedName: toKebabCase(lastSegment(normalized)), }; } return { parsed: { source: "url", url: repoUrl }, - label: `Git repo — ${repoUrl}`, - suggestedName: toKebabCase(lastSegment(url.pathname).replace(/\.git$/, "")), + label: `${kind} repo — ${repoUrl}`, + suggestedName: toKebabCase(repoName), }; }; +const parseRawGitSource = (url: URL, subPath?: string): SkillSourcePreview | null => { + if (pathSegments(url).length < 2) { + return null; + } + const repoName = lastSegment(url.pathname).replace(/\.git$/, ""); + return buildGitSourcePreview("Git", buildRepoUrl(url), repoName, subPath); +}; + +/** + * Parse an scp-style `git@host:org/repo` or `ssh://git@host/org/repo` clone URL, registering it + * exactly as typed: git treats the `.git` suffix as optional, and forcing one on breaks hosts whose + * paths are not `org/repo`, like Azure DevOps `v3/...` and CodeCommit `v1/repos/...`. The scp form is + * rewritten to `ssh://` only to reuse the https host and credential rules, and only a URL that + * survives that round trip unchanged is accepted, which keeps traversal segments off the feed. + */ +const parseSshSource = (raw: string, subPath?: string): SkillSourcePreview | null => { + const trimmed = raw.trim(); + const scp = SSH_SCP_REGEX.exec(trimmed); + const candidate = scp ? `${SSH_SCHEME}${scp[1]}@${scp[2]}/${scp[3]}` : trimmed; + if (!candidate.toLowerCase().startsWith(SSH_SCHEME)) { + return null; + } + const url = toUrl(candidate); + if (!url || url.username === "" || url.password !== "" || !isSafeHost(url)) { + return null; + } + const pathStart = candidate.indexOf("/", SSH_SCHEME.length); + if (pathStart === -1 || url.pathname !== candidate.slice(pathStart) || pathSegments(url).length < 2) { + return null; + } + return buildGitSourcePreview("SSH", trimmed, lastSegment(url.pathname).replace(/\.git$/i, ""), subPath); +}; + const parseArchiveSource = (url: URL): SkillSourcePreview => ({ parsed: { source: "archive", url: url.href }, label: `Zip archive — ${url.host}${url.pathname}`, @@ -175,10 +212,15 @@ const parseArchiveSource = (url: URL): SkillSourcePreview => ({ /** * Parse any git-accessible repository URL or https zip archive URL into a registerable skill * source. A `.zip` path is an `archive` source (S3, Artifactory, any static host). GitHub URLs - * keep their `github`/`git-subdir` shorthand; every other host is treated as a raw repo URL, + * keep their `github`/`git-subdir` shorthand; ssh clone URLs stay ssh so a private host + * authenticates with the user's own key; every other host is treated as a raw repo URL, * with an optional subfolder turning it into git-subdir. */ export const parseSkillSource = (rawUrl: string, subPath?: string): SkillSourcePreview | null => { + const ssh = parseSshSource(rawUrl, subPath); + if (ssh) { + return ssh; + } const url = parseRepoUrl(rawUrl); if (!url) { return null; @@ -268,14 +310,14 @@ export const getSourceDisplayText = (source: PluginSource): string => { }; /** - * Get clickable link for plugin source + * Get clickable link for plugin source. Ssh clone urls are not browsable, so they yield null. */ export const getSourceLink = (source: PluginSource): string | null => { if (source.source === "github" && source.repo) { return `https://github.com/${source.repo}`; } const linksToUrl = source.source === "url" || source.source === "git-subdir" || source.source === "archive"; - return linksToUrl && source.url ? source.url : null; + return linksToUrl && source.url && BROWSABLE_URL_REGEX.test(source.url) ? source.url : null; }; /** diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.test.tsx b/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.test.tsx new file mode 100644 index 00000000000..1ec1f72f4a1 --- /dev/null +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.test.tsx @@ -0,0 +1,42 @@ +import { render, screen } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; + +import { Plugin } from "./types"; + +import SkillDetail from "./skill_detail"; + +const buildSkill = (source: Plugin["source"]): Plugin => ({ + id: "plugin-id", + name: "my-skill", + source, + enabled: true, +}); + +describe("SkillDetail source", () => { + it("links a github source to the repository", () => { + render(); + expect(screen.getByRole("link", { name: "github.com/org/repo" })).toHaveAttribute( + "href", + "https://github.com/org/repo", + ); + }); + + it("renders an ssh clone url as plain text instead of an unusable link", () => { + render( + , + ); + expect(screen.getByText("git@ghe.example.com:org/repo.git")).toBeInTheDocument(); + expect(screen.queryByRole("link")).not.toBeInTheDocument(); + }); + + it("renders an ssh git-subdir source as plain text without a tree path", () => { + render( + , + ); + expect(screen.getByText("git@ghe.example.com:org/repo.git @ plugins/x")).toBeInTheDocument(); + expect(screen.queryByRole("link")).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx b/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx index b37204c6073..c873e660bb3 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx @@ -1,8 +1,38 @@ import React, { useState } from "react"; import { ArrowLeft, Check, Copy, Link2 } from "lucide-react"; import { cn } from "@/lib/cva.config"; -import { buildMarketplaceSettingsSnippet, formatInstallCommand } from "./helpers"; -import { Plugin } from "./types"; +import { buildMarketplaceSettingsSnippet, formatInstallCommand, getSourceDisplayText, getSourceLink } from "./helpers"; +import { Plugin, PluginSource } from "./types"; + +const SkillSource: React.FC<{ source: PluginSource }> = ({ source }) => { + const link = getSourceLink(source); + const href = link && source.source === "git-subdir" && source.path ? `${link}/tree/main/${source.path}` : link; + if (href) { + return ( + + ); + } + if (!source.url) { + return null; + } + return ( +
+
Source
+
{getSourceDisplayText(source)}
+
+ ); +}; interface SkillDetailProps { skill: Plugin; @@ -22,14 +52,6 @@ const SkillDetail: React.FC = ({ skill, onBack }) => { setTimeout(() => setCopiedKey(null), 2000); }; - const sourceUrl = (() => { - const src = skill.source; - if (src.source === "github" && src.repo) return `https://github.com/${src.repo}`; - if (src.source === "git-subdir" && src.url) return src.path ? `${src.url}/tree/main/${src.path}` : src.url; - if ((src.source === "url" || src.source === "archive") && src.url) return src.url; - return null; - })(); - const installCommand = formatInstallCommand(skill); const settingsSnippet = buildMarketplaceSettingsSnippet( @@ -128,20 +150,7 @@ const SkillDetail: React.FC = ({ skill, onBack }) => { - {sourceUrl && ( - - )} + {skill.keywords && skill.keywords.length > 0 && (
diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.tsx index 0fab5555343..4d2a0e88f84 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.tsx @@ -144,6 +144,7 @@ export function ModelMaxBudgetEditor({ onClick={() => removeEntry(entry.id)} disabled={!premiumUser} title={hintWhenLocked} + aria-label="Remove model budget" className="absolute top-2 right-2 text-muted-foreground hover:text-destructive transition-colors p-1" > diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx index eadbca87140..d8c331110c0 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx @@ -39,6 +39,7 @@ export interface KeyResponse { key_name: string; key_alias: string; spend: number; + total_spend: number; max_budget: number; expires: string; models: string[]; @@ -64,6 +65,8 @@ export interface KeyResponse { model_max_budget_usage?: Record | null; soft_budget_cooldown: boolean; blocked: boolean; + deleted_at?: string | null; + deleted_by?: string | null; litellm_budget_table: Record; organization_id: string | null; org_id?: string | null; diff --git a/ui/litellm-dashboard/src/components/settings.test.tsx b/ui/litellm-dashboard/src/components/settings.test.tsx index 08cb9550646..4ba5dd23fd1 100644 --- a/ui/litellm-dashboard/src/components/settings.test.tsx +++ b/ui/litellm-dashboard/src/components/settings.test.tsx @@ -302,6 +302,113 @@ describe("Settings", () => { }); }); + const mockS3Callback = (variables: Record, callbackName = "s3") => { + mockGetCallbacksCall.mockResolvedValue({ + callbacks: [{ name: callbackName, variables }], + available_callbacks: { + s3: { + litellm_callback_name: "s3", + litellm_callback_params: [ + "AWS_ACCESS_KEY_ID", + "AWS_SECRET_ACCESS_KEY", + "AWS_REGION_NAME", + "S3_LOG_PROMPTS_ONLY", + ], + ui_callback_name: "s3 Bucket (AWS)", + }, + }, + alerts: [], + }); + mockGetCallbackConfigsCall.mockResolvedValue([ + { + id: "s3", + displayName: "S3", + dynamic_params: { + s3_bucket_name: { type: "text", ui_name: "S3 Bucket Name", required: false }, + s3_log_prompts_only: { type: "boolean", ui_name: "Log Prompts Only", required: false }, + }, + }, + ]); + }; + + const openS3EditModal = async (callbackName = "s3") => { + const user = userEvent.setup(); + render(); + await user.click(await screen.findByTestId(`callback-actions-${callbackName}-success`)); + await user.click(await screen.findByTestId("callback-action-edit")); + return user; + }; + + it("should render a saved boolean dynamic param as a checked switch and post false when toggled off", async () => { + mockS3Callback({ S3_LOG_PROMPTS_ONLY: "true" }); + const user = await openS3EditModal(); + + const promptsOnlySwitch = await screen.findByRole("switch", { name: "Log Prompts Only" }); + expect(promptsOnlySwitch).toBeChecked(); + + await user.click(promptsOnlySwitch); + expect(promptsOnlySwitch).not.toBeChecked(); + await user.click(within(screen.getByRole("dialog")).getByRole("button", { name: "Save Changes" })); + + await waitFor(() => { + expect(vi.mocked(setCallbacksCall)).toHaveBeenCalledWith( + "token", + expect.objectContaining({ + environment_variables: expect.objectContaining({ callback: "s3", s3_log_prompts_only: "false" }), + }), + ); + }); + }); + + it("should render an unset boolean dynamic param as an unchecked switch and post true when toggled on", async () => { + mockS3Callback({ S3_LOG_PROMPTS_ONLY: null }); + const user = await openS3EditModal(); + + const promptsOnlySwitch = await screen.findByRole("switch", { name: "Log Prompts Only" }); + expect(promptsOnlySwitch).not.toBeChecked(); + + await user.click(promptsOnlySwitch); + await user.click(within(screen.getByRole("dialog")).getByRole("button", { name: "Save Changes" })); + + await waitFor(() => { + expect(vi.mocked(setCallbacksCall)).toHaveBeenCalledWith( + "token", + expect.objectContaining({ + environment_variables: expect.objectContaining({ callback: "s3", s3_log_prompts_only: "true" }), + }), + ); + }); + }); + + it.each(["True", "1"])("should render a boolean dynamic param stored as %s as a checked switch", async (stored) => { + mockS3Callback({ S3_LOG_PROMPTS_ONLY: stored }); + await openS3EditModal(); + + expect(await screen.findByRole("switch", { name: "Log Prompts Only" })).toBeChecked(); + }); + + it("should resolve the s3_v2 callback to the s3 dynamic params and post under the s3_v2 name", async () => { + mockS3Callback({ S3_LOG_PROMPTS_ONLY: null }, "s3_v2"); + const user = await openS3EditModal("s3_v2"); + + const promptsOnlySwitch = await screen.findByRole("switch", { name: "Log Prompts Only" }); + expect(promptsOnlySwitch).not.toBeChecked(); + expect(within(screen.getByRole("dialog")).getByRole("combobox", { name: "Callback" })).toHaveValue("S3"); + + await user.click(promptsOnlySwitch); + await user.click(within(screen.getByRole("dialog")).getByRole("button", { name: "Save Changes" })); + + await waitFor(() => { + expect(vi.mocked(setCallbacksCall)).toHaveBeenCalledWith( + "token", + expect.objectContaining({ + environment_variables: expect.objectContaining({ callback: "s3_v2", s3_log_prompts_only: "true" }), + litellm_settings: { success_callback: ["s3_v2"] }, + }), + ); + }); + }); + it("should send the typed webhook url for an alert type when the alerting tab is saved", async () => { const user = userEvent.setup(); render(); diff --git a/ui/litellm-dashboard/src/components/settings.tsx b/ui/litellm-dashboard/src/components/settings.tsx index e549770af6e..9247f22ec28 100644 --- a/ui/litellm-dashboard/src/components/settings.tsx +++ b/ui/litellm-dashboard/src/components/settings.tsx @@ -67,19 +67,20 @@ const DynamicParamsFields: React.FC = ({ params, callb return null; } + const callbackConfig = findCallbackConfig(callbackConfigs, selectedCallback); return (
{params.map((param) => { - const callbackConfig = callbackConfigs.find((config) => config.id === selectedCallback); const paramConfig = callbackConfig?.dynamic_params?.[param] || {}; const paramType = paramConfig.type || "text"; const fieldLabel = paramConfig.ui_name || param.replace(/_/g, " ").replace(/\b\w/g, (l) => l.toUpperCase()); const isRequired = paramConfig.required || false; const selectOptions: string[] = Array.isArray(paramConfig.options) ? paramConfig.options : []; const isSelect = paramType === "select" && selectOptions.length > 0; + const isBoolean = paramType === "boolean"; const fieldId = `${fieldIdPrefix}-${param}`; const validationRules = isRequired ? { required: `Please enter the ${fieldLabel.toLowerCase()}` } : undefined; - const registration = isSelect ? undefined : register(param, validationRules); + const registration = isSelect || isBoolean ? undefined : register(param, validationRules); return ( @@ -111,7 +112,22 @@ const DynamicParamsFields: React.FC = ({ params, callb )} /> )} + {isBoolean && ( + ( + field.onChange(checked ? "true" : "false")} + onBlur={field.onBlur} + /> + )} + /> + )} {!isSelect && + !isBoolean && (paramType === "password" ? ( = ({ }) => { const { control } = useFormContext(); const inputId = React.useId(); - const selectedConfig = callbackConfigs.find((config) => config.id === selectedCallback) ?? null; + const selectedConfig = findCallbackConfig(callbackConfigs, selectedCallback) ?? null; return ( = ({ ); }; +const CALLBACK_CONFIG_ALIASES: Record = { s3_v2: "s3" }; + +interface DynamicParamConfig { + type?: string; + ui_name?: string; + required?: boolean; + options?: string[]; +} + +interface CallbackConfigWithParams { + id: string; + dynamic_params?: Record; +} + +const findCallbackConfig = ( + callbackConfigs: readonly T[], + callbackName: string | null, +): T | undefined => { + if (!callbackName) { + return undefined; + } + const configId = CALLBACK_CONFIG_ALIASES[callbackName] ?? callbackName; + return callbackConfigs.find((config) => config.id === configId); +}; + // Shared helper function to get dynamic params for a callback const getDynamicParamsForCallback = ( callbackName: string | null, @@ -231,7 +272,7 @@ const getDynamicParamsForCallback = ( return fallbackVariables ? Object.keys(fallbackVariables) : []; } - const callbackConfig = callbackConfigs.find((config) => config.id === callbackName); + const callbackConfig = findCallbackConfig(callbackConfigs, callbackName); if (callbackConfig?.dynamic_params) { return Object.keys(callbackConfig.dynamic_params); } diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index eb912ffa3cc..16eb70fdf5e 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -1609,6 +1609,99 @@ describe("TeamInfoView", () => { }); }); + describe("per-model budgets", () => { + const teamWithModelBudget = () => + createMockTeamData({ + models: ["gpt-4"], + model_max_budget: { "gpt-4": { max_budget: 5, budget_duration: "1d" } }, + model_max_budget_usage: { "gpt-4": { current_spend: 1.25, budget_limit: 5, time_period: "1d" } }, + }); + + const openSettingsEditor = async (user: ReturnType) => { + await waitFor(() => { + expect(screen.queryAllByText("Test Team").length).toBeGreaterThan(0); + }); + await user.click(screen.getByRole("tab", { name: "Settings" })); + await user.click(await screen.findByRole("button", { name: /edit settings/i })); + await screen.findByLabelText("Team Name"); + }; + + const savedPayload = async () => { + await waitFor(() => { + expect(networking.teamUpdateCall).toHaveBeenCalled(); + }); + return vi.mocked(networking.teamUpdateCall).mock.calls[0][1] as Record; + }; + + it("shows the stored per-model budget and its current spend in the read-only settings view", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue(teamWithModelBudget()); + + renderWithProviders(); + + await waitFor(() => { + expect(screen.queryAllByText("Test Team").length).toBeGreaterThan(0); + }); + await user.click(screen.getByRole("tab", { name: "Settings" })); + + expect(await screen.findByText("Per-Model Budget (gpt-4): $5 per 1d, spent $1.25")).toBeInTheDocument(); + }); + + it("seeds the editor from the stored budget and keeps it read-only without an enterprise license", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue(teamWithModelBudget()); + + renderWithProviders(); + + await openSettingsEditor(user); + + expect(screen.getByPlaceholderText("Max spend ($)")).toHaveValue(5); + expect(screen.getByPlaceholderText("Max spend ($)")).toBeDisabled(); + expect(screen.getByRole("button", { name: /Add Model Budget/i })).toBeDisabled(); + }); + + it("leaves model_max_budget out of a save that did not touch it", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue(teamWithModelBudget()); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); + + renderWithProviders(); + + await openSettingsEditor(user); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + expect(await savedPayload()).not.toHaveProperty("model_max_budget"); + }); + + it("sends the edited cap for the model", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue(teamWithModelBudget()); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); + + renderWithProviders(); + + await openSettingsEditor(user); + fireEvent.change(screen.getByPlaceholderText("Max spend ($)"), { target: { value: "2.5" } }); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + expect((await savedPayload()).model_max_budget).toEqual({ "gpt-4": { budget_limit: 2.5, time_period: "1d" } }); + }); + + it("sends an empty model_max_budget when the last row is removed, so the stored cap is cleared", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue(teamWithModelBudget()); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); + + renderWithProviders(); + + await openSettingsEditor(user); + await user.click(screen.getByRole("button", { name: "Remove model budget" })); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + expect((await savedPayload()).model_max_budget).toEqual({}); + }); + }); + describe("team member settings", () => { it("should populate Default Key Duration from the team's stored metadata", async () => { const user = userEvent.setup({ delay: null }); diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index ce705008678..c476f3492a3 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -51,6 +51,13 @@ import GuardrailsSelect from "./GuardrailsSelect"; import { copyToClipboard as utilCopyToClipboard } from "../../utils/dataUtils"; import AccessGroupSelector from "../common_components/AccessGroupSelector"; import BudgetDurationDropdown, { NEVER_RESETS_BUDGET_DURATION } from "../common_components/budget_duration_dropdown"; +import { + ModelBudgetUsage, + ModelMaxBudget, + ModelMaxBudgetField, + modelMaxBudgetToEntries, +} from "../key_team_helpers/ModelMaxBudgetEditor"; +import { modelMaxBudgetUpdate, StoredModelMaxBudget } from "../key_team_helpers/modelMaxBudgetPayload"; import { computeTeamModelBadges, normalizeTeamModelSelection, @@ -268,6 +275,8 @@ export interface TeamData { max_budget: number | null; soft_budget?: number | null; budget_duration: string | null; + model_max_budget?: StoredModelMaxBudget | null; + model_max_budget_usage?: Record | null; models: string[]; blocked: boolean; spend: number; @@ -563,6 +572,7 @@ const TeamInfoView: React.FC = ({ const [isDeleting, setIsDeleting] = useState(false); const [isTeamSaving, setIsTeamSaving] = useState(false); const [teamModelAliases, setTeamModelAliases] = useState>({}); + const [teamModelMaxBudget, setTeamModelMaxBudget] = useState({}); const routerSettingsRef = React.useRef(null); const [organization, setOrganization] = useState(null); const { userRole, userId } = useAuthorized(); @@ -628,6 +638,7 @@ const TeamInfoView: React.FC = ({ const startEditing = () => { form.reset(teamFormValues()); + setTeamModelMaxBudget((teamData?.team_info?.model_max_budget ?? {}) as ModelMaxBudget); setTeamMemberSettingsOpen(false); setSearchToolSettingsOpen(false); setIsEditing(true); @@ -1078,6 +1089,11 @@ const TeamInfoView: React.FC = ({ updateData.model_aliases = teamModelAliases; } + const modelBudgets = modelMaxBudgetUpdate(teamModelMaxBudget, info.model_max_budget); + if (modelBudgets !== undefined) { + updateData.model_max_budget = modelBudgets; + } + // Handle router_settings - read fresh values from DOM at save time. const currentRouterSettings = routerSettingsRef.current?.getValue(); if (currentRouterSettings?.router_settings) { @@ -1536,6 +1552,15 @@ const TeamInfoView: React.FC = ({ )} + + {({ ref, value, ...field }) => } @@ -2051,6 +2076,17 @@ const TeamInfoView: React.FC = ({ : "No Limit"}
Budget Reset: {info.budget_duration || "Never"}
+ {modelMaxBudgetToEntries(info.model_max_budget as ModelMaxBudget | null | undefined).map( + ({ model, budgetLimit, timePeriod }) => { + const spent = model === null ? undefined : info.model_max_budget_usage?.[model]?.current_spend; + return ( +
+ Per-Model Budget ({model}): ${budgetLimit ?? "?"} per {timePeriod} + {spent !== undefined && `, spent $${spent}`} +
+ ); + }, + )} {info.metadata?.soft_budget_alerting_emails && Array.isArray(info.metadata.soft_budget_alerting_emails) && info.metadata.soft_budget_alerting_emails.length > 0 && ( diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx index 9efcff04832..6a2f778d4a9 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx @@ -174,6 +174,7 @@ describe("KeyEditView", () => { key_name: "sk-...TUuw", key_alias: "asdasdas", spend: 0, + total_spend: 0, max_budget: 0, expires: "null", models: [], diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx index b403255b329..4bf41c1f3a8 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx @@ -119,6 +119,7 @@ describe("KeyInfoView", () => { key_name: "sk-...TUuw", key_alias: "asdasdas", spend: 0, + total_spend: 0, max_budget: 0, expires: "null", models: [], @@ -272,6 +273,23 @@ describe("KeyInfoView", () => { }); }); + it("shows lifetime spend separately from the resettable period spend", async () => { + vi.mocked(useAuthorized).mockReturnValue(baseUseAuthorizedMock); + + renderWithProviders( + {}} + keyId={"test-key-id"} + onKeyDataUpdate={() => {}} + teams={[]} + />, + ); + + expect(await screen.findByText("$0.2500")).toBeInTheDocument(); + expect(screen.getByTestId("key-lifetime-spend")).toHaveTextContent("Lifetime spend: $340.5000"); + }); + it("should render the key's saved router fallbacks", async () => { vi.mocked(useAuthorized).mockReturnValue(baseUseAuthorizedMock); diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx index 0e6dba64110..06e088f956b 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -677,6 +677,9 @@ export default function KeyInfoView({ {currentKeyData.budget_reset_at && (

Resets {formatTimestamp(currentKeyData.budget_reset_at)}

)} +

+ Lifetime spend: ${formatNumberWithCommas(currentKeyData.total_spend ?? 0, 4)} +

@@ -935,6 +938,11 @@ export default function KeyInfoView({

${formatNumberWithCommas(currentKeyData.spend, 4)} USD

+
+

Lifetime Spend

+

${formatNumberWithCommas(currentKeyData.total_spend ?? 0, 4)} USD

+
+

Budget

diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 17ec8367324..872875cc535 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -7859,7 +7859,10 @@ export interface paths { * * Returns: * - key: str - The key that was looked up, echoed back as it was passed in - * - info: dict - The key's row, minus the hashed token + * - info: dict - The key's row, minus the hashed token. Deleted keys are served from the + * LiteLLM_DeletedVerificationToken archive and carry deleted_at / deleted_by + * - status: "active" | "expired" | "revoked" | "deleted" - Derived from blocked, expires and + * whether the row came from the archive * - key_alias: str | None - User-friendly key alias * - spend: float - Amount spent by the key. When budget_duration is set this covers only the * current budget window, not the key's lifetime @@ -7917,7 +7920,9 @@ export interface paths { * * Parameters: * expand: Optional[List[str]] - Expand related objects (e.g. 'user' to include user information) - * status: Optional[str] - Filter by status. Currently supports "deleted" to query deleted keys. + * status: Optional[str] - Filter by status: "active", "expired", "revoked" (blocked) or "deleted". + * "deleted" reads the LiteLLM_DeletedVerificationToken archive; the other values partition the + * live key table, so every live key matches exactly one of them. * * Returns: * { @@ -15682,6 +15687,7 @@ export interface paths { * - prompts: Optional[List[str]] - List of prompts that the team is allowed to use. * - organization_id: Optional[str] - The organization id of the team. Default is None. Create via `/organization/new`. * - model_aliases: Optional[dict] - Model aliases for the team. [Docs](https://docs.litellm.ai/docs/proxy/team_based_routing#create-team-with-model-alias) + * - model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}} * - guardrails: Optional[List[str]] - Guardrails for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails) * - policies: Optional[List[str]] - Policies for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails/guardrail_policies) * - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key. @@ -15908,6 +15914,7 @@ export interface paths { * - tags: Optional[List[str]] - Tags for [tracking spend](https://litellm.vercel.app/docs/proxy/enterprise#tracking-spend-for-custom-tags) and/or doing [tag-based routing](https://litellm.vercel.app/docs/proxy/tag_routing). * - organization_id: Optional[str] - The organization id of the team. Default is None. Create via `/organization/new`. * - model_aliases: Optional[dict] - Model aliases for the team. [Docs](https://docs.litellm.ai/docs/proxy/team_based_routing#create-team-with-model-alias) + * - model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}} * - guardrails: Optional[List[str]] - Guardrails for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails) * - policies: Optional[List[str]] - Policies for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails/guardrail_policies) * - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key. @@ -29648,6 +29655,8 @@ export interface components { object_permission_id?: string | null; /** Org Id */ org_id?: string | null; + /** Organization Id */ + organization_id?: string | null; /** * Permissions * @default {} @@ -29686,6 +29695,11 @@ export interface components { team_id?: string | null; /** Token */ token?: string | null; + /** + * Total Spend + * @default 0 + */ + total_spend: number; /** Tpd Limit */ tpd_limit?: number | null; /** Tpm Limit */ @@ -31259,6 +31273,11 @@ export interface components { team_id?: string | null; /** Token */ token?: string | null; + /** + * Total Spend + * @default 0 + */ + total_spend: number; /** Tpd Limit */ tpd_limit?: number | null; /** Tpm Limit */ @@ -33419,6 +33438,13 @@ export interface components { model_aliases?: { [key: string]: unknown; } | null; + /** + * Model Max Budget + * @description Max budget per model for every key on the team, overridable per key (e.g. {'gpt-4o': {'max_budget': 10, 'budget_duration': '1d'}}) + */ + model_max_budget?: { + [key: string]: components["schemas"]["BudgetConfig"]; + } | null; /** Model Rpm Limit */ model_rpm_limit?: { [key: string]: number; @@ -34177,6 +34203,13 @@ export interface components { model_aliases?: { [key: string]: unknown; } | null; + /** + * Model Max Budget + * @description Max budget per model for every key on the team, overridable per key (e.g. {'gpt-4o': {'max_budget': 10, 'budget_duration': '1d'}}) + */ + model_max_budget?: { + [key: string]: components["schemas"]["BudgetConfig"]; + } | null; /** Model Rpm Limit */ model_rpm_limit?: { [key: string]: number; @@ -39326,6 +39359,13 @@ export interface components { model_aliases?: { [key: string]: unknown; } | null; + /** + * Model Max Budget + * @description Max budget per model for every key on the team, overridable per key (e.g. {'gpt-4o': {'max_budget': 10, 'budget_duration': '1d'}}) + */ + model_max_budget?: { + [key: string]: components["schemas"]["BudgetConfig"]; + } | null; /** Model Rpm Limit */ model_rpm_limit?: { [key: string]: number; @@ -40037,6 +40077,10 @@ export interface components { team_model_aliases?: { [key: string]: unknown; } | null; + /** Team Model Max Budget */ + team_model_max_budget?: { + [key: string]: unknown; + } | null; /** * Team Models * @default [] @@ -40057,6 +40101,11 @@ export interface components { team_tpm_limit?: number | null; /** Token */ token?: string | null; + /** + * Total Spend + * @default 0 + */ + total_spend: number; /** Tpd Limit */ tpd_limit?: number | null; /** Tpm Limit */ @@ -51366,7 +51415,7 @@ export interface operations { sort_order?: string; /** @description Expand related objects (e.g. 'user') */ expand?: string[] | null; - /** @description Filter by status (e.g. 'deleted') */ + /** @description Filter by status: 'active' (not blocked, not expired), 'expired' (not blocked, past expiry), 'revoked' (blocked) or 'deleted' (archived keys). Omit to return live keys regardless of status. */ status?: string | null; /** @description Filter keys by project ID */ project_id?: string | null;