scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-19 00:01:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 88

Merge remote-tracking branch 'origin/main' into litellm_converted_stream_post_call_hook

Browse source
This commit is contained in:
yucheng 2026-09-17 00:31:24 +00:00
commit e1cce943de
296 changed files with 20092 additions and 3114 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
.circleci/config.yml
View file

@ -2983,7 +2983,7 @@ workflows:
name: integration-<< matrix.suite >>
matrix:
parameters:
suite: [management, accounting, providers]
suite: [management, accounting, database, providers]
filters:
branches:
only:

3
ARCHITECTURE.md
View file

@ -45,7 +45,7 @@ sequenceDiagram
ProxyServer->>Auth: user_api_key_auth()
Auth->>Redis: Check API key cache
Redis-->>Auth: Key info + spend limits
ProxyServer->>Hooks: max_budget_limiter, parallel_request_limiter
ProxyServer->>Hooks: parallel_request_limiter, cache_control_check
Hooks->>Redis: Check/increment rate limit counters
ProxyServer->>Router: route_request()
Router->>Main: litellm.acompletion()
@ -145,7 +145,6 @@ graph TD
| Hook | File | Purpose |
|------|------|---------|
| `max_budget_limiter` | `proxy/hooks/max_budget_limiter.py` | Enforce budget limits |
| `parallel_request_limiter` | `proxy/hooks/parallel_request_limiter_v3.py` | Rate limiting per key/user |
| `cache_control_check` | `proxy/hooks/cache_control_check.py` | Cache validation |
| `responses_id_security` | `proxy/hooks/responses_id_security.py` | Response ID validation |

5
litellm-proxy-extras/litellm_proxy_extras/migrations/20260916000000_add_key_total_spend/migration.sql Normal file
View file

@ -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;

2
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -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("{}")

102
litellm-rust/Cargo.lock generated
View file

@ -70,6 +70,12 @@ dependencies = [
"rustversion",
]
[[package]]
name = "arcstr"
version = "1.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "03918c3dbd7701a85c6b9887732e2921175f26c350b4563841d0958c21d57e6d"
[[package]]
name = "async-compression"
version = "0.4.46"
@ -1837,6 +1843,12 @@ version = "0.2.186"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66"
[[package]]
name = "linux-raw-sys"
version = "0.12.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53"
[[package]]
name = "litellm-auth"
version = "0.1.0"
@ -1915,6 +1927,17 @@ dependencies = [
"tokio",
]
[[package]]
name = "litellm-cache-redis"
version = "0.1.0"
dependencies = [
"litellm-cache",
"redis",
"redis-test",
"serde_json",
"tokio",
]
[[package]]
name = "litellm-core"
version = "0.1.0"
@ -1951,6 +1974,7 @@ dependencies = [
name = "litellm-python-bridge"
version = "0.1.0"
dependencies = [
"bytes",
"criterion",
"futures-util",
"litellm-auth",
@ -2139,6 +2163,16 @@ dependencies = [
"minimal-lexical",
]
[[package]]
name = "num-bigint"
version = "0.5.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "93e7820bc0a80a0238e650327316f929ba18d5be054b647490a3a6a339f3e7c0"
dependencies = [
"num-integer",
"num-traits",
]
[[package]]
name = "num-conv"
version = "0.2.2"
@ -2655,6 +2689,36 @@ dependencies = [
"crossbeam-utils",
]
[[package]]
name = "redis"
version = "1.7.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2acbc41a996f7652b2ddd9dfd98cc4ff602cfd742ae35382f07f608405ab50ed"
dependencies = [
"arcstr",
"combine",
"itoa",
"num-bigint",
"percent-encoding",
"ryu",
"sha1_smol",
"socket2 0.6.5",
"url",
"xxhash-rust",
]
[[package]]
name = "redis-test"
version = "1.0.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "804d36862e4323b69f96440cbb13c9894fc90176abdeaf91264e21d5d77f6aca"
dependencies = [
"rand 0.9.5",
"redis",
"socket2 0.6.5",
"tempfile",
]
[[package]]
name = "redox_syscall"
version = "0.5.18"
@ -2845,6 +2909,19 @@ dependencies = [
"semver",
]
[[package]]
name = "rustix"
version = "1.1.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "891efababe418670775f199f0d233d84843c227a0949a883ce15b37c78d6629d"
dependencies = [
"bitflags",
"errno",
"libc",
"linux-raw-sys",
"windows-sys 0.61.2",
]
[[package]]
name = "rustls"
version = "0.21.12"
@ -3095,6 +3172,12 @@ dependencies = [
"digest 0.10.7",
]
[[package]]
name = "sha1_smol"
version = "1.0.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bbfa15b3dddfee50a0fff136974b3e1bde555604ba463834a7eb7deb6417705d"
[[package]]
name = "sha2"
version = "0.10.9"
@ -3298,6 +3381,19 @@ version = "0.13.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca"
[[package]]
name = "tempfile"
version = "3.27.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd"
dependencies = [
"fastrand",
"getrandom 0.4.3",
"once_cell",
"rustix",
"windows-sys 0.61.2",
]
[[package]]
name = "thiserror"
version = "1.0.69"
@ -4181,6 +4277,12 @@ version = "0.13.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "66fee0b777b0f5ac1c69bb06d361268faafa61cd4682ae064a171c16c433e9e4"
[[package]]
name = "xxhash-rust"
version = "0.8.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "aee1b19627c7c60102ab80d3a9cbe18de90bfe03bfa6c3715447681f0e8c8af6"
[[package]]
name = "yoke"
version = "0.8.3"

15
litellm-rust/crates/cache-redis/Cargo.toml Normal file
View file

@ -0,0 +1,15 @@
[package]
name = "litellm-cache-redis"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-cache.workspace = true
redis = "1.7.0"
serde_json.workspace = true
tokio.workspace = true
[dev-dependencies]
redis-test = "1.0.4"

315
litellm-rust/crates/cache-redis/src/cache.rs Normal file
View file

@ -0,0 +1,315 @@
use std::sync::{Arc, Mutex, MutexGuard};
use std::time::Duration;
use litellm_cache::{
BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheEntry, CacheFuture, CacheKwargs,
Error,
};
use redis::Commands;
const DEFAULT_TTL: Duration = Duration::from_secs(600);
const KEY_PREFIX: &str = "litellm-cache:";
pub struct RedisCache<C = redis::Connection> {
connection: Arc<Mutex<C>>,
default_ttl: Duration,
}
impl RedisCache<redis::Connection> {
pub fn new(url: &str, default_ttl: Option<Duration>) -> Result<Self, Error> {
let client = redis::Client::open(url).map_err(|_| Error::Unavailable)?;
let connection = client.get_connection().map_err(|_| Error::Unavailable)?;
Ok(Self::with_connection(connection, default_ttl))
}
}
impl<C> RedisCache<C>
where
C: redis::ConnectionLike + Send + 'static,
{
fn with_connection(connection: C, default_ttl: Option<Duration>) -> Self {
Self {
connection: Arc::new(Mutex::new(connection)),
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
}
}
fn connection(&self) -> Result<MutexGuard<'_, C>, Error> {
self.connection.lock().map_err(|_| Error::Unavailable)
}
fn namespaced_key(key: &str) -> String {
format!("{KEY_PREFIX}{key}")
}
fn namespaced_pattern() -> &'static str {
const PATTERN: &str = "litellm-cache:*";
PATTERN
}
fn encode(value: &CacheEntry) -> Result<Vec<u8>, Error> {
serde_json::to_vec(value).map_err(|_| Error::InvalidEntry)
}
fn decode(value: Vec<u8>) -> Result<CacheEntry, Error> {
serde_json::from_slice(&value).map_err(|_| Error::InvalidEntry)
}
fn ttl_seconds(ttl: Duration) -> u64 {
ttl.as_secs()
.saturating_add(u64::from(ttl.subsec_nanos() > 0))
.max(1)
}
fn run_blocking<T, F>(connection: Arc<Mutex<C>>, operation: F) -> CacheFuture<'static, T>
where
T: Send + 'static,
F: FnOnce(&mut C) -> Result<T, Error> + Send + 'static,
{
Box::pin(async move {
tokio::task::spawn_blocking(move || {
let mut connection = connection.lock().map_err(|_| Error::Unavailable)?;
operation(&mut connection)
})
.await
.map_err(|_| Error::Unavailable)?
})
}
}
impl<C> BaseCache for RedisCache<C>
where
C: redis::ConnectionLike + Send + 'static,
{
type Value = CacheEntry;
fn default_ttl(&self) -> Duration {
self.default_ttl
}
fn set_cache(&self, key: &str, value: Self::Value, kwargs: CacheKwargs) -> Result<(), Error> {
let payload = Self::encode(&value)?;
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
self.connection()?
.set_ex::<_, _, ()>(Self::namespaced_key(key), payload, ttl)
.map_err(|_| Error::Unavailable)
}
fn get_cache(&self, key: &str, _: &CacheKwargs) -> Result<Option<Self::Value>, Error> {
self.connection()?
.get::<_, Option<Vec<u8>>>(Self::namespaced_key(key))
.map_err(|_| Error::Unavailable)?
.map(Self::decode)
.transpose()
}
fn delete_cache(&self, key: &str) -> Result<(), Error> {
self.connection()?
.del::<_, ()>(Self::namespaced_key(key))
.map_err(|_| Error::Unavailable)
}
fn flush_cache(&self) -> Result<(), Error> {
let mut connection = self.connection()?;
let keys = connection
.scan_match(Self::namespaced_pattern())
.map_err(|_| Error::Unavailable)?
.collect::<redis::RedisResult<Vec<String>>>()
.map_err(|_| Error::Unavailable)?;
if keys.is_empty() {
return Ok(());
}
connection
.del::<_, usize>(keys)
.map(|_| ())
.map_err(|_| Error::Unavailable)
}
fn async_set_cache<'a>(
&'a self,
key: &'a str,
value: Self::Value,
kwargs: CacheKwargs,
) -> CacheFuture<'a, ()> {
let payload = Self::encode(&value);
let key = Self::namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
Self::run_blocking(Arc::clone(&self.connection), move |connection| {
connection
.set_ex::<_, _, ()>(key, payload?, ttl)
.map_err(|_| Error::Unavailable)
})
}
fn async_get_cache<'a>(
&'a self,
key: &'a str,
_: &'a CacheKwargs,
) -> CacheFuture<'a, Option<Self::Value>> {
let key = Self::namespaced_key(key);
Box::pin(async move {
Self::run_blocking(Arc::clone(&self.connection), move |connection| {
connection
.get::<_, Option<Vec<u8>>>(key)
.map_err(|_| Error::Unavailable)
})
.await?
.map(Self::decode)
.transpose()
})
}
fn async_set_cache_pipeline<'a>(
&'a self,
cache_list: Vec<(String, Self::Value)>,
kwargs: CacheKwargs,
) -> CacheFuture<'a, ()> {
let entries = cache_list
.into_iter()
.map(|(key, value)| {
Self::encode(&value).map(|payload| (Self::namespaced_key(&key), payload))
})
.collect::<Result<Vec<_>, _>>();
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
Self::run_blocking(Arc::clone(&self.connection), move |connection| {
for (key, payload) in entries? {
connection
.set_ex::<_, _, ()>(key, payload, ttl)
.map_err(|_| Error::Unavailable)?;
}
Ok(())
})
}
fn async_delete_cache<'a>(&'a self, key: &'a str) -> CacheFuture<'a, ()> {
let key = Self::namespaced_key(key);
Self::run_blocking(Arc::clone(&self.connection), move |connection| {
connection.del::<_, ()>(key).map_err(|_| Error::Unavailable)
})
}
fn disconnect(&self) -> CacheFuture<'_, ()> {
Box::pin(async { Ok(()) })
}
fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult> {
Box::pin(async move {
Self::run_blocking(Arc::clone(&self.connection), |connection| {
redis::cmd("PING")
.query::<String>(connection)
.map_err(|_| Error::Unavailable)
})
.await?;
Ok(CacheConnectionResult {
status: CacheConnectionStatus::Success,
message: "Redis cache connection test successful".into(),
error: None,
})
})
}
}
#[cfg(test)]
mod tests {
use super::RedisCache;
use litellm_cache::{BaseCache, CacheEntry, CacheKwargs};
use redis_test::{MockCmd, MockRedisConnection};
use serde_json::json;
use std::time::Duration;
fn entry() -> CacheEntry {
CacheEntry {
timestamp: 123.0,
response: json!({"choices": [{"text": "cached"}]}),
}
}
#[test]
fn cache_entries_round_trip_through_json() {
let entry = entry();
let encoded = RedisCache::<redis::Connection>::encode(&entry).unwrap();
assert_eq!(
RedisCache::<redis::Connection>::decode(encoded).unwrap(),
entry
);
}
#[test]
fn invalid_json_is_rejected() {
assert!(RedisCache::<redis::Connection>::decode(b"not json".to_vec()).is_err());
}
#[test]
fn ttl_seconds_rounds_up_and_keeps_expiration_positive() {
assert_eq!(
RedisCache::<redis::Connection>::ttl_seconds(Duration::ZERO),
1
);
assert_eq!(
RedisCache::<redis::Connection>::ttl_seconds(Duration::from_millis(1500)),
2
);
assert_eq!(
RedisCache::<redis::Connection>::ttl_seconds(Duration::from_secs(15)),
15
);
}
#[test]
fn redis_commands_round_trip_entries_and_delete_only_namespaced_keys() {
let value = entry();
let payload = RedisCache::<redis::Connection>::encode(&value).unwrap();
let connection = MockRedisConnection::new([
MockCmd::new(
redis::cmd("SETEX")
.arg("litellm-cache:key")
.arg(600)
.arg(payload.clone()),
Ok("OK"),
),
MockCmd::new(redis::cmd("GET").arg("litellm-cache:key"), Ok(payload)),
MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)),
])
.assert_all_commands_consumed();
let cache = RedisCache::with_connection(connection, None);
cache
.set_cache("key", value.clone(), CacheKwargs::default())
.unwrap();
assert_eq!(
cache.get_cache("key", &CacheKwargs::default()).unwrap(),
Some(value)
);
cache.delete_cache("key").unwrap();
}
#[test]
fn flush_scans_and_deletes_only_cache_keys() {
let connection = MockRedisConnection::new([
MockCmd::new(
redis::cmd("SCAN")
.cursor_arg(0)
.arg("MATCH")
.arg("litellm-cache:*"),
Ok(redis_test::redis_value!(["0", ["litellm-cache:key"]])),
),
MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)),
])
.assert_all_commands_consumed();
let cache = RedisCache::with_connection(connection, None);
cache.flush_cache().unwrap();
}
#[tokio::test]
async fn test_connection_runs_ping_off_executor() {
let connection = MockRedisConnection::new([MockCmd::new(redis::cmd("PING"), Ok("PONG"))])
.assert_all_commands_consumed();
let cache = RedisCache::with_connection(connection, None);
assert_eq!(
cache.test_connection().await.unwrap().status,
litellm_cache::CacheConnectionStatus::Success
);
}
}

3
litellm-rust/crates/cache-redis/src/lib.rs Normal file
View file

@ -0,0 +1,3 @@
mod cache;
pub use cache::RedisCache;

6
litellm-rust/crates/cache-redis/tests/cache.rs Normal file
View file

@ -0,0 +1,6 @@
use litellm_cache_redis::RedisCache;
#[test]
fn constructor_rejects_invalid_urls() {
assert!(RedisCache::new("not a redis url", None).is_err());
}

13
litellm-rust/crates/core/src/ocr/client.rs
View file

@ -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,
))))
}

126
litellm-rust/crates/core/src/ocr/document.rs
View file

@ -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<OcrDocument, super::Error> {
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<OcrDocument, super::Error> {
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::<u64>()));
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]

6
litellm-rust/crates/core/src/ocr/error.rs
View file

@ -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}")]

50
litellm-rust/crates/core/src/ocr/lifecycle.rs
View file

@ -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<LiteLLMOcrResponse>),
MapFailure(Error),
@ -83,7 +85,8 @@ impl OcrHostOperation {
}
pub enum OcrHostResult {
Request(Result<(Box<LiteLLMOcrRequest>, bool), Error>),
Request(Result<(Box<LiteLLMOcrRequest<OcrDocumentInput>>, bool), Error>),
Document(Result<OcrFileContent, Error>),
Lifecycle(Result<(), HostFailure<Error>>),
AzureAdToken(Result<ResolvedCredential, AuthError>),
PreCall(Result<OcrPreCallRequest, Error>),
@ -313,7 +316,7 @@ struct PendingOperation {
struct OcrExecution {
client: Option<OcrClient>,
request: Option<LiteLLMOcrRequest>,
request: Option<LiteLLMOcrRequest<OcrDocumentInput>>,
operations_tx: mpsc::UnboundedSender<PendingOperation>,
operations_rx: mpsc::UnboundedReceiver<PendingOperation>,
pending_result: Option<oneshot::Sender<OcrHostResult>>,
@ -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<OcrDocumentInput>,
hooks: &ProtocolHooks,
) -> Result<LiteLLMOcrRequest, Error> {
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,

7
litellm-rust/crates/core/src/ocr/mod.rs
View file

@ -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"]

70
litellm-rust/crates/core/src/ocr/types.rs
View file

@ -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<String>,
},
Bytes {
bytes: Bytes,
file_name: Option<String>,
mime_type: Option<String>,
},
HostReader {
mime_type: Option<String>,
},
}
impl From<OcrDocument> 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<String>,
}
#[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<D = OcrDocument> {
pub model: String,
pub document: OcrDocument,
pub document: D,
pub connection: OcrConnection,
pub hooks: Arc<dyn OcrHooks>,
pub litellm_call_id: Option<String>,
@ -101,10 +133,10 @@ pub struct LiteLLMOcrRequest {
pub(crate) adapter: OcrAdapterKind,
}
impl LiteLLMOcrRequest {
impl<D> LiteLLMOcrRequest<D> {
pub fn new(
model: String,
document: OcrDocument,
document: D,
custom_llm_provider: Option<&str>,
optional_params: Map<String, Value>,
) -> Result<Self, Error> {
@ -151,6 +183,36 @@ impl LiteLLMOcrRequest {
..self
}
}
pub fn map_document<T, E>(
self,
map: impl FnOnce(D) -> Result<T, E>,
) -> Result<LiteLLMOcrRequest<T>, 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<T>(self, document: T) -> LiteLLMOcrRequest<T> {
let Ok(request) = self.map_document(|_| Ok::<T, Infallible>(document));
request
}
}
impl From<LiteLLMOcrRequest> for LiteLLMOcrRequest<OcrDocumentInput> {
fn from(request: LiteLLMOcrRequest) -> Self {
let Ok(request) = request
.map_document(|document| Ok::<_, Infallible>(OcrDocumentInput::Document(document)));
request
}
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]

43
litellm-rust/crates/core/src/ocr/wire.rs
View file

@ -68,9 +68,9 @@ pub struct DecodedOcrResponse<T> {
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
pub struct OcrWireRequest {
pub struct OcrWireRequest<D = Value> {
pub model: String,
pub document: Value,
pub document: D,
pub api_key: Option<String>,
pub api_base: Option<String>,
pub custom_llm_provider: Option<String>,
@ -141,10 +141,34 @@ pub fn consumed_optional_params(
}
pub fn decode_request(wire: OcrWireRequest) -> Result<LiteLLMOcrRequest, Error> {
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<D>(wire: OcrWireRequest<D>) -> Result<LiteLLMOcrRequest<D>, 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<LiteLLMOcrRequest, Error>
.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<LiteLLMOcrRequest, Error>
})
}
fn decode_document(value: Value) -> Result<OcrDocument, OcrRequestError> {
pub fn decode_document(value: Value) -> Result<OcrDocument, Error> {
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<String, InputSource>, 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));
}
}
}

154
litellm-rust/crates/core/tests/ocr.rs
View file

@ -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<super::OcrDocumentInput>,
content: Result<super::OcrFileContent, crate::ocr::Error>,
) -> (Result<super::LiteLLMOcrResponse, crate::ocr::Error>, 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::<u64>()));
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"),
});

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

@ -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

370
litellm-rust/crates/python-bridge/src/routes/ocr/document.rs
View file

@ -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<u8>),
#[derive(Debug)]
pub(super) struct PythonFileReader {
reader: Py<PyAny>,
name: Option<String>,
}
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<OcrFileContent> {
let value = self.reader.bind(py).call0()?;
let bytes = if value.is_instance_of::<PyString>() {
Bytes::from(value.extract::<String>()?)
} else if value.is_instance_of::<PyBytes>() {
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<String>)> {
if file.is_instance_of::<PyString>() {
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<Bytes> {
if value.is_exact_instance_of::<PyBytes>() {
return Ok(Bytes::from_owner(value.extract::<PyBackedBytes>()?));
}
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::<PyBytes>() {
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::<String>())
.transpose()?;
let value = reader.call0()?;
let bytes = if value.is_instance_of::<PyString>() {
FileBytes::Native(value.extract::<String>()?.into_bytes())
} else if value.is_instance_of::<PyBytes>() {
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::<PyBackedBytes>()?.as_ref(),
))
}
pub(super) struct FileDocumentInput {
bytes: FileBytes,
name: Option<String>,
mime_type: Option<String>,
pub input: OcrDocumentInput,
pub reader: Option<PythonFileReader>,
}
impl FromPyObject<'_, '_> for FileDocumentInput {
@ -104,79 +63,79 @@ impl FromPyObject<'_, '_> for FileDocumentInput {
Err(error) if error.is_instance_of::<pyo3::exceptions::PyKeyError>(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::<pyo3::exceptions::PyKeyError>(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::<PyString>() {
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::<PathBuf>()?,
mime_type,
},
reader: None,
});
}
if file.is_instance_of::<PyBytes>() {
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::<String>())
.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<OcrDocument> {
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<Py<PyAny>> {
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<Py<PyAny>> {
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::<FileDocumentInput>().err().unwrap();
assert!(error.is_instance_of::<PyTypeError>(py));
}
let document = py.eval(c"{'file': b'abc'}", None, None).unwrap();
let error = py
.eval(c"{'file': 'scan.pdf'}", None, None)
.unwrap()
.extract::<FileDocumentInput>()
.err()
.unwrap();
assert!(error.is_instance_of::<PyValueError>(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::<FileDocumentInput>().err().unwrap();
assert!(error.is_instance_of::<PyTypeError>(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::<usize>()
.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::<FileDocumentInput>().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::<FileDocumentInput>()
.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::<PyTypeError>(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");
}
}

7
litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs
View file

@ -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)

22
litellm-rust/crates/python-bridge/src/routes/ocr/lifecycle.rs
View file

@ -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<litellm_core::ocr::OcrFileContent> {
self.projected()?
.fields
.reader
.as_ref()
.ok_or_else(missing_state)?
.read(py)
}
fn acquire_azure_ad_token(&self, py: Python<'_>) -> PyResult<ResolvedCredential> {
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)?;

1
litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs
View file

@ -9,6 +9,5 @@ use pyo3::prelude::*;
pub(super) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> {
value::register(module)?;
document::register(module)?;
lifecycle::register(module)
}

141
litellm-rust/crates/python-bridge/src/routes/ocr/project.rs
View file

@ -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<PyAny>,
pub document: Py<PyAny>,
pub document: Option<Py<PyAny>>,
pub reader: Option<PythonFileReader>,
pub api_key: Py<PyAny>,
pub azure_ad_token_provider: Option<PythonTokenProvider>,
pub provider: &'static str,
@ -25,7 +27,7 @@ pub(super) struct ProjectedOcrFields {
}
pub(super) struct ProjectedOcrCall {
pub request: LiteLLMOcrRequest,
pub request: LiteLLMOcrRequest<OcrDocumentInput>,
pub fields: ProjectedOcrFields,
}
@ -80,12 +82,12 @@ impl<'py> OcrArguments<'_, 'py> {
}
enum ProjectedDocument {
File { wire: Value, retained: Py<PyAny> },
File(FileDocumentInput),
Other { wire: Value, retained: Py<PyAny> },
}
impl ProjectedDocument {
fn project(py: Python<'_>, document: &Bound<'_, PyAny>) -> PyResult<Self> {
fn project(document: &Bound<'_, PyAny>) -> PyResult<Self> {
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<PyAny>) {
fn into_parts(
self,
) -> PyResult<(
OcrDocumentInput,
Option<Py<PyAny>>,
Option<PythonFileReader>,
)> {
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<ProjectedOcrCall> {
@ -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<PyAny>)> {
ProjectedDocument::project(py, document).map(ProjectedDocument::into_parts)
) -> PyResult<(
OcrDocumentInput,
Option<Py<PyAny>>,
Option<PythonFileReader>,
)> {
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::<PyValueError>(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::<PyKeyError>(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::<PyTypeError>(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<String> = document.getattr("reads").unwrap().extract().unwrap();
assert_eq!(reads, ["type", "mime_type", "file"]);
});

4
litellm/__init__.py
View file

@ -244,6 +244,7 @@ telemetry = True
max_tokens: int = DEFAULT_MAX_TOKENS # OpenAI Defaults
drop_params = drop_params_env_flag(os.environ, verbose_logger)
modify_params = bool(os.getenv("LITELLM_MODIFY_PARAMS", False))
bedrock_neutralize_orphaned_tool_blocks: bool = True
use_chat_completions_url_for_anthropic_messages: bool = bool(
os.getenv("LITELLM_USE_CHAT_COMPLETIONS_URL_FOR_ANTHROPIC_MESSAGES", False)
) # When True, routes OpenAI /v1/messages requests to chat/completions instead of the Responses API
@ -1818,6 +1819,9 @@ if TYPE_CHECKING:
from .llms.azure.responses.o_series_transformation import (
AzureOpenAIOSeriesResponsesAPIConfig as AzureOpenAIOSeriesResponsesAPIConfig,
)
from .llms.azure_ai.responses.transformation import (
AzureAIResponsesAPIConfig as AzureAIResponsesAPIConfig,
)
from .llms.xai.responses.transformation import (
XAIResponsesAPIConfig as XAIResponsesAPIConfig,
)

5
litellm/_lazy_imports_registry.py
View file

@ -234,6 +234,7 @@ LLM_CONFIG_NAMES: Final = (
"OpenAIResponsesAPIConfig",
"AzureOpenAIResponsesAPIConfig",
"AzureOpenAIOSeriesResponsesAPIConfig",
"AzureAIResponsesAPIConfig",
"XAIResponsesAPIConfig",
"LiteLLMProxyResponsesAPIConfig",
"HostedVLLMResponsesAPIConfig",
@ -946,6 +947,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
".llms.azure.responses.o_series_transformation",
"AzureOpenAIOSeriesResponsesAPIConfig",
),
"AzureAIResponsesAPIConfig": (
".llms.azure_ai.responses.transformation",
"AzureAIResponsesAPIConfig",
),
"XAIResponsesAPIConfig": (
".llms.xai.responses.transformation",
"XAIResponsesAPIConfig",

20
litellm/completion_extras/litellm_responses_transformation/transformation.py
View file

@ -502,7 +502,12 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
elif key == "response_format":
text_format = self._transform_response_format_to_text_format(value)
if text_format:
responses_api_request["text"] = text_format
responses_api_request["text"] = self._merge_text(responses_api_request, text_format)
elif key == "verbosity":
responses_api_request["text"] = self._merge_text(
responses_api_request,
MappingProxyType({"verbosity": value}), # pyright: ignore[reportUnknownArgumentType] # untyped value
)
elif key == "tool_choice":
responses_api_request["tool_choice"] = self._normalize_tool_choice_for_responses_api(value)
elif key == "stream_options":
@ -518,6 +523,19 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
elif key == "web_search_options":
self._add_web_search_tool(responses_api_request, value)
@staticmethod
def _merge_text(
responses_api_request: "ResponsesAPIOptionalRequestParams", update: Mapping[str, object]
) -> "ResponseText":
existing: Final = cast( # cast-ok: text field is a ResponseText | dict[str, Any] | None union
"dict[str, object]",
dict(responses_api_request).get("text") or {}, # mutable-ok: one-shot merge seed
)
return cast( # cast-ok: merged mapping is a valid ResponseText shape
"ResponseText",
{**existing, **update}, # mutable-ok: one-shot merged payload
)
def _build_sanitized_litellm_params(self, litellm_params: dict) -> dict[str, object]:
"""Build sanitized litellm_params with merged metadata."""
responses_optional_param_keys: Final = set(ResponsesAPIOptionalRequestParams.__annotations__.keys())

3
litellm/constants.py
View file

@ -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"

8
litellm/experimental_mcp_client/client.py
View file

@ -346,13 +346,17 @@ class MCPClient:
self.update_auth_value(auth_value)
async def discovery_auth_fingerprint(self) -> str:
return self._hash_discovery_auth(await self.prepare_request_auth())
async def prepare_request_auth(self) -> httpx.Request:
"""Preview the authenticated request without sending it, closing the auth flow afterwards."""
request: Final = httpx.Request("POST", self.server_url or "http://localhost/", headers=self._get_auth_headers())
if self._resolved_auth is None:
return self._hash_discovery_auth(request)
return request
flow: Final = self._resolved_auth.async_auth_flow(request)
try:
authenticated: Final = await flow.__anext__()
return self._hash_discovery_auth(authenticated)
return authenticated
finally:
await flow.aclose()

6
litellm/integrations/callback_configs.json
View file

@ -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"

1
litellm/integrations/custom_logger.py
View file

@ -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)

11
litellm/integrations/opentelemetry.py
View file

@ -2965,16 +2965,13 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
)
propagator: Final = TraceContextTextMapPropagator()
carrier: Final = {"traceparent": _traceparent}
carrier: Final = {key: headers[key] for key in ("traceparent", "tracestate") if headers.get(key) is not None}
_parent_context: Final = propagator.extract(carrier=carrier)
return _parent_context
def _get_span_context(self, kwargs, default_span: Span | None = None):
from opentelemetry import context, trace
from opentelemetry.trace.propagation.tracecontext import (
TraceContextTextMapPropagator,
)
litellm_params: Final = kwargs.get("litellm_params", {}) or {}
proxy_server_request: Final = litellm_params.get("proxy_server_request", {}) or {}
@ -2998,11 +2995,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
# Priority 2: HTTP traceparent header
if traceparent is not None:
verbose_logger.debug("OpenTelemetry: Using traceparent header for context propagation")
carrier: Final = {"traceparent": traceparent}
return (
TraceContextTextMapPropagator().extract(carrier=carrier),
None,
)
return self.get_traceparent_from_header(headers=headers), None
# Priority 3: Active span from global context (auto-detection)
try:

32
litellm/integrations/otel/plumbing/context.py
View file

@ -26,6 +26,7 @@ if TYPE_CHECKING:
from litellm.integrations.otel.model.destination import OtelDestination
_PROPAGATOR: Final = TraceContextTextMapPropagator()
_W3C_TRACE_HEADERS: Final = frozenset(("traceparent", "tracestate"))
# The request's root span — the FastAPI-owned SERVER span — captured ONCE when the
# proxy first resolves it, so request-level spans (the LLM call, guardrails) can
@ -310,6 +311,37 @@ def extract_traceparent(headers: Mapping[str, str]) -> Context | None:
return _PROPAGATOR.extract(carrier)
def _outgoing_trace_context(parent_span: object) -> Context | None:
if isinstance(parent_span, Span) and is_recordable_span(parent_span):
return context_from_span(parent_span)
root: Final = request_root_span()
if root is not None:
return context_from_span(root)
current: Final = get_current()
if is_recordable_span(get_current_span(current)):
return current
return None
def inject_trace_context(headers: Mapping[str, str], parent_span: object = None) -> dict[str, str]:
"""``headers`` plus W3C ``traceparent``/``tracestate`` for this request's span.
Parent preference: ``parent_span`` (the request span auth stashed on the key), then
the anchored request root span, then the ambient active span. Only trace context is
injected, never Baggage. Unchanged when no valid span exists anywhere.
"""
context: Final = _outgoing_trace_context(parent_span)
if context is None:
return dict(headers) # mutable-ok: OpenTelemetry propagator requires a mutable carrier
carrier: Final = { # mutable-ok: OpenTelemetry propagator requires a mutable carrier
key: value for key, value in headers.items() if key.lower() not in _W3C_TRACE_HEADERS
}
_PROPAGATOR.inject(carrier, context=context)
return carrier
# The OTLP destinations this request's key or team pointed its traces at, resolved
# once during auth. A ``ContextVar`` for the same reason the root span above is one:
# it rides the request task's context into the ``asyncio.create_task`` children that

69
litellm/integrations/s3.py
View file

@ -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}")

16
litellm/integrations/s3_v2.py
View file

@ -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,
)

72
litellm/litellm_core_utils/streaming_chunk_builder_utils.py
View file

@ -148,6 +148,7 @@ class _ToolCallChunk(TypedDict):
class _UsageBearingChunk(TypedDict, total=False):
usage: Usage | None
_hidden_params: Mapping[str, str]
choices: ReadOnly[Sequence[StreamingChoices | Mapping[str, object]]]
class _UsageSummary(TypedDict):
@ -921,21 +922,22 @@ class ChunkProcessor:
prompt_tokens_details = attach_cache_creation_token_details(prompt_tokens_details, cache_creation_token_details)
completion_tokens = self._reset_anthropic_cursor_completion_tokens(
recovered_completion_tokens: Final = self._reset_anthropic_cursor_completion_tokens(
chunks=chunks,
completion_tokens=completion_tokens,
completion_usage_updates=completion_usage_updates,
)
cursor_was_reset: Final = recovered_completion_tokens != completion_tokens
return UsagePerChunk(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
completion_tokens=recovered_completion_tokens,
cache_creation_input_tokens=cache_creation_input_tokens,
cache_read_input_tokens=cache_read_input_tokens,
server_tool_use=server_tool_use,
web_search_requests=web_search_requests,
google_maps_grounding_requests=google_maps_grounding_requests,
completion_tokens_details=completion_tokens_details,
completion_tokens_details=None if cursor_was_reset else completion_tokens_details,
prompt_tokens_details=prompt_tokens_details,
cost=cost,
inference_geo=self._last_provider_pricing_field(chunks, "inference_geo"),
@ -960,6 +962,30 @@ class ChunkProcessor:
]
return values[-1] if values else None
@staticmethod
def _finish_reason_of_choice(choice: object) -> str | None:
match choice:
case StreamingChoices(finish_reason=reason) | Choices(finish_reason=reason):
return reason
case {"finish_reason": str() as reason}:
return reason
case _:
return None
@staticmethod
def _chunk_choices(chunk: "_UsageBearingChunk | ModelResponse | ModelResponseStream") -> Sequence[object]:
if isinstance(chunk, dict):
return chunk.get("choices", ())
return getattr(chunk, "choices", ())
@staticmethod
def _saw_finish_reason(chunks: Sequence["_UsageBearingChunk | ModelResponse"]) -> bool:
return any(
ChunkProcessor._finish_reason_of_choice(choice) is not None
for chunk in chunks
for choice in ChunkProcessor._chunk_choices(chunk)
)
@staticmethod
def _reset_anthropic_cursor_completion_tokens(
chunks: Sequence["_UsageBearingChunk | ModelResponse"],
@ -970,18 +996,18 @@ class ChunkProcessor:
See the ``completion_usage_updates`` comment in
``_calculate_usage_per_chunk``. The accumulated value is NOT a stale
cursor when either it is > 1 (definitely not a placeholder) or we saw
>= 2 completion-bearing usage events (positive evidence ``message_delta``
arrived). Otherwise — the only completion update we ever saw was the
Anthropic ``message_start`` cursor (=1) — reset to 0 so
``calculate_usage()``'s ``or token_counter(text=...)`` fallback estimates
from the actually-received completion text instead of trusting the
placeholder. Gated on ``custom_llm_provider == "anthropic"`` so the
heuristic (which encodes Anthropic's specific message_start SSE shape)
does not silently affect other providers that may legitimately report
``completion_tokens=1`` from a single usage event.
cursor when we saw >= 2 completion-bearing usage events or any chunk
carried a ``finish_reason`` (positive evidence ``message_delta``
arrived). Otherwise the only completion update we ever saw was the
Anthropic ``message_start`` cursor, a small placeholder whose magnitude
varies per request (1 and 8 both observed live), so reset to 0 and let
``calculate_usage()``'s ``or token_counter(...)`` fallback estimate from
the actually-received text and reasoning instead. Gated on
``custom_llm_provider == "anthropic"`` so the heuristic (which encodes
Anthropic's specific message_start SSE shape) does not silently affect
other providers that legitimately report usage from a single event.
"""
saw_non_cursor_completion: Final = completion_tokens > 1 or completion_usage_updates >= 2
saw_non_cursor_completion: Final = completion_usage_updates >= 2 or ChunkProcessor._saw_finish_reason(chunks)
if saw_non_cursor_completion:
return completion_tokens
@ -995,7 +1021,7 @@ class ChunkProcessor:
if isinstance(hp, dict):
custom_llm_provider = hp.get("custom_llm_provider")
if custom_llm_provider == "anthropic" and completion_tokens == 1:
if custom_llm_provider == "anthropic":
return 0
return completion_tokens
@ -1039,10 +1065,13 @@ class ChunkProcessor:
returned_usage.prompt_tokens = 0
returned_usage.completion_tokens = (
completion_tokens
or token_counter(
model=model,
text=completion_output,
count_response_tokens=True, # count_response_tokens is a Flag to tell token counter this is a response, No need to add extra tokens we do for input messages
or (
token_counter(
model=model,
text=completion_output,
count_response_tokens=True, # count_response_tokens is a Flag to tell token counter this is a response, No need to add extra tokens we do for input messages
)
+ (reasoning_tokens or 0)
)
)
returned_usage.total_tokens = returned_usage.prompt_tokens + returned_usage.completion_tokens
@ -1066,15 +1095,16 @@ class ChunkProcessor:
returned_usage.completion_tokens_details = completion_tokens_details
if reasoning_tokens is not None:
capped_reasoning_tokens: Final = min(max(0, reasoning_tokens), returned_usage.completion_tokens)
if returned_usage.completion_tokens_details is None:
returned_usage.completion_tokens_details = CompletionTokensDetailsWrapper(
reasoning_tokens=reasoning_tokens
reasoning_tokens=capped_reasoning_tokens,
text_tokens=returned_usage.completion_tokens - capped_reasoning_tokens,
)
elif (
returned_usage.completion_tokens_details is not None
and returned_usage.completion_tokens_details.reasoning_tokens is None
):
capped_reasoning_tokens: Final = min(max(0, reasoning_tokens), returned_usage.completion_tokens)
returned_usage.completion_tokens_details.reasoning_tokens = capped_reasoning_tokens
if returned_usage.completion_tokens_details.text_tokens is None:
returned_usage.completion_tokens_details.text_tokens = (

4
litellm/llms/anthropic/chat/guardrail_translation/handler.py
View file

@ -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

5
litellm/llms/anthropic/chat/handler.py
View file

@ -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

27
litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py
View file

@ -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
)

7
litellm/llms/azure/responses/transformation.py
View file

@ -49,12 +49,7 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
return BaseAzureLLM._base_validate_azure_environment(headers=headers, litellm_params=litellm_params)
def get_stripped_model_name(self, model: str) -> str:
# if "responses/" is in the model name, remove it
if "responses/" in model:
model = model.replace("responses/", "")
if "o_series" in model:
model = model.replace("o_series/", "")
return model
return model.replace("responses/", "").replace("o_series/", "").replace("azure_ai/", "")
def _handle_reasoning_item(self, item: dict[str, Any]) -> dict[str, Any]:
"""

22
litellm/llms/azure_ai/common_utils.py
View file

@ -9,6 +9,7 @@ from litellm.types.llms.openai import AllMessageValues
from litellm.types.router import GenericLiteLLMParams
AzureAIApiKeyHeader = Literal["Authorization", "api-key", "Api-Key", "Ocp-Apim-Subscription-Key"]
AZURE_OPENAI_V1_HOST_SUFFIXES: Final = (".services.ai.azure.com", ".openai.azure.com")
def is_foundry_model_inference_base(api_base: str) -> bool:
@ -19,11 +20,13 @@ def is_foundry_model_inference_base(api_base: str) -> bool:
return "/openai/deployments" not in parsed.path
def api_key_header_for_base(api_base: str | None) -> AzureAIApiKeyHeader:
def is_azure_openai_v1_host(api_base: str | None) -> bool:
host: Final = urlparse(api_base).hostname if api_base else None
if host and (host.endswith(".services.ai.azure.com") or host.endswith(".openai.azure.com")):
return "api-key"
return "Authorization"
return host is not None and host.endswith(AZURE_OPENAI_V1_HOST_SUFFIXES)
def api_key_header_for_base(api_base: str | None) -> AzureAIApiKeyHeader:
return "api-key" if is_azure_openai_v1_host(api_base) else "Authorization"
def get_azure_ai_entra_token(litellm_params: Mapping[str, object] | None = None) -> str | None:
@ -70,6 +73,17 @@ def get_azure_ai_auth_headers(
AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: Final = "azure_model_router_selected_model"
def azure_ai_supports_native_responses(model: str | None, api_base: str | None) -> bool:
resolved_base: Final = AzureFoundryModelInfo.get_api_base(api_base)
if resolved_base is not None and not is_azure_openai_v1_host(resolved_base):
return False
if model is None:
return True
if "claude" in model.lower():
return False
return AzureFoundryModelInfo.get_azure_ai_route(model) == "default"
class AzureFoundryModelInfo(BaseLLMModelInfo):
"""Model info for Azure AI / Azure Foundry models."""

0
litellm/llms/azure_ai/responses/__init__.py Normal file
View file

53
litellm/llms/azure_ai/responses/transformation.py Normal file
View file

@ -0,0 +1,53 @@
from typing import Final
import httpx
from litellm.llms.azure.responses.transformation import AzureOpenAIResponsesAPIConfig
from litellm.llms.azure_ai.common_utils import (
AzureFoundryModelInfo,
api_key_header_for_base,
get_azure_ai_auth_headers,
)
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import LlmProviders
_PROJECT_PATH_PREFIX: Final = ("api", "projects")
_RESPONSES_PATH: Final = ("openai", "v1", "responses")
def _responses_url(api_base: str) -> str:
base_url: Final = httpx.URL(api_base)
segments: Final = tuple(segment for segment in base_url.path.split("/") if segment)
project_root: Final = segments[:3] if segments[:2] == _PROJECT_PATH_PREFIX else ()
return str(base_url.copy_with(path="/" + "/".join((*project_root, *_RESPONSES_PATH)), query=None))
class AzureAIResponsesAPIConfig(AzureOpenAIResponsesAPIConfig):
@property
def custom_llm_provider(self) -> LlmProviders:
return LlmProviders.AZURE_AI
def validate_environment(self, headers: dict, model: str, litellm_params: GenericLiteLLMParams | None) -> dict:
params: Final = litellm_params or GenericLiteLLMParams()
auth_headers: Final = get_azure_ai_auth_headers(
api_key=AzureFoundryModelInfo.get_api_key(params.api_key),
litellm_params=params.model_dump(),
api_key_header=api_key_header_for_base(AzureFoundryModelInfo.get_api_base(params.api_base)),
)
return { # mutable-ok: the handler updates the returned headers in place per the dict contract
**headers,
**auth_headers,
"Content-Type": "application/json",
}
def supports_native_websocket(self) -> bool:
return False
def get_complete_url(self, api_base: str | None, litellm_params: dict) -> str:
resolved_base: Final = AzureFoundryModelInfo.get_api_base(api_base)
if resolved_base is None:
raise ValueError(
"api_base is required for the Azure AI Foundry Responses API. "
"Set the api_base parameter or the AZURE_AI_API_BASE environment variable."
)
return _responses_url(resolved_base)

6
litellm/llms/base_llm/guardrail_translation/base_translation.py
View file

@ -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:

101
litellm/llms/bedrock/chat/converse_transformation.py
View file

@ -53,6 +53,7 @@ from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionAnnotation,
ChatCompletionAssistantMessage,
ChatCompletionAssistantToolCall,
ChatCompletionRedactedThinkingBlock,
ChatCompletionResponseMessage,
ChatCompletionSystemMessage,
@ -205,6 +206,84 @@ class AmazonConverseConfig(BaseConfig):
return messages_copy
@staticmethod
def _has_orphaned_tool_blocks(messages: list[AllMessageValues]) -> bool:
return any(
(m.get("role") == "assistant" and m.get("tool_calls")) or m.get("role") in ("tool", "function")
for m in messages
)
@staticmethod
def _neutralize_orphaned_tool_blocks(
messages: list[AllMessageValues], optional_params: dict
) -> list[AllMessageValues]:
if optional_params.get("tools") or not AmazonConverseConfig._has_orphaned_tool_blocks(messages):
return messages
from litellm.litellm_core_utils.prompt_templates.common_utils import (
convert_content_list_to_str,
)
def _tool_call_text(tool_call: ChatCompletionAssistantToolCall) -> str:
function = tool_call.get("function") or {}
name = function.get("name") or "unknown_tool"
arguments = function.get("arguments") or ""
call_id = tool_call.get("id")
label = f"tool call {call_id}" if call_id else "tool call"
return f"[{label}: {name}({arguments})]"
def _result_text(message: AllMessageValues) -> str:
rendered = convert_content_list_to_str(message).strip()
return rendered or "<non-text tool result omitted>"
guardrail_active: Final = "guardrailConfig" in optional_params
def _rewrite(message: AllMessageValues) -> AllMessageValues:
role = message.get("role")
tool_calls = message.get("tool_calls")
if role == "assistant" and tool_calls:
base_text: Final = convert_content_list_to_str(message)
call_texts: Final = tuple(_tool_call_text(call) for call in tool_calls)
text: Final = "\n".join(part for part in (base_text, *call_texts) if part)
return ChatCompletionAssistantMessage(role="assistant", content=text)
if role in ("tool", "function"):
tool_call_id = message.get("tool_call_id")
name = message.get("name")
label = f"tool result for {tool_call_id or name or 'unknown'}"
result_text: Final = f"[{label}: {_result_text(message)}]"
# Tool results are externally controlled, so guard them wherever they
# land in history; _convert_consecutive_user_messages_to_guarded_text
# only covers the trailing user turn.
content: Final = [{"type": "guarded_text", "text": result_text}] if guardrail_active else result_text
return ChatCompletionUserMessage(role="user", content=content)
return message
verbose_logger.warning(
"litellm.bedrock: request has tool blocks in message history but no "
"`tools=` param; neutralizing orphaned tool blocks to text so Bedrock "
"accepts the request without a toolConfig. Non-text tool-result "
"payloads are dropped. Pass `tools=` to preserve structured tool calling."
)
return [_rewrite(message) for message in messages]
@staticmethod
def _handle_orphaned_tool_blocks(messages: list[AllMessageValues], optional_params: dict) -> list[AllMessageValues]:
if litellm.bedrock_neutralize_orphaned_tool_blocks:
return AmazonConverseConfig._neutralize_orphaned_tool_blocks(messages, optional_params)
if "tools" in optional_params or not has_tool_call_blocks(messages):
return messages
if litellm.modify_params:
optional_params["tools"] = add_dummy_tool(custom_llm_provider="bedrock_converse")
return messages
raise litellm.utils.UnsupportedParamsError(
message="Bedrock doesn't support tool calling without `tools=` param specified. Pass `tools=` param OR set `litellm.modify_params = True` // `litellm_settings::modify_params: True` to add dummy tool to the request.",
model="",
llm_provider="bedrock",
)
@classmethod
def get_config(cls):
return {
@ -1609,20 +1688,6 @@ class AmazonConverseConfig(BaseConfig):
drop_params: bool = False,
litellm_params: Mapping[str, object] | None = None,
) -> CommonRequestObject:
## VALIDATE REQUEST
"""
Bedrock doesn't support tool calling without `tools=` param specified.
"""
if "tools" not in optional_params and messages is not None and has_tool_call_blocks(messages):
if litellm.modify_params:
optional_params["tools"] = add_dummy_tool(custom_llm_provider="bedrock_converse")
else:
raise litellm.UnsupportedParamsError(
message="Bedrock doesn't support tool calling without `tools=` param specified. Pass `tools=` param OR set `litellm.modify_params = True` // `litellm_settings::modify_params: True` to add dummy tool to the request.",
model="",
llm_provider="bedrock",
)
# Drop thinking param if thinking is enabled but thinking_blocks are missing
# This prevents the error: "Expected thinking or redacted_thinking, but found tool_use"
#
@ -1735,7 +1800,9 @@ class AmazonConverseConfig(BaseConfig):
messages, system_content_blocks = self._transform_system_message(messages, model=model)
# Convert last user message to guarded_text if guardrailConfig is present
messages = self._convert_consecutive_user_messages_to_guarded_text(messages, optional_params)
messages = self._convert_consecutive_user_messages_to_guarded_text(
self._handle_orphaned_tool_blocks(messages, optional_params), optional_params
)
## TRANSFORMATION ##
_data: Final[CommonRequestObject] = self._transform_request_helper(
@ -1796,7 +1863,9 @@ class AmazonConverseConfig(BaseConfig):
messages, system_content_blocks = self._transform_system_message(messages, model=model)
# Convert last user message to guarded_text if guardrailConfig is present
messages = self._convert_consecutive_user_messages_to_guarded_text(messages, optional_params)
messages = self._convert_consecutive_user_messages_to_guarded_text(
self._handle_orphaned_tool_blocks(messages, optional_params), optional_params
)
_data: Final[CommonRequestObject] = self._transform_request_helper(
model=model,

18
litellm/llms/bedrock_mantle/chat/transformation.py
View file

@ -22,6 +22,7 @@ from litellm.llms.bedrock_mantle.common_utils import (
BEDROCK_MANTLE_DEFAULT_REGION,
BedrockMantleAuthMixin,
)
from litellm.llms.openai.chat.gpt_5_transformation import is_gpt_reasoning_series_name
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllMessageValues
from litellm.types.router import GenericLiteLLMParams
@ -108,13 +109,22 @@ class BedrockMantleChatConfig(BedrockMantleAuthMixin, OpenAILikeChatConfig):
def get_supported_openai_params(self, model: str) -> list:
base_params: Final = super().get_supported_openai_params(model)
extra_params: Final = tuple(
param
for param, supported in (
("verbosity", is_gpt_reasoning_series_name(model)),
("reasoning_effort", self._supports_reasoning(model)),
)
if supported and param not in base_params
)
return [*base_params, *extra_params] # mutable-ok: fresh list required by the inherited signature
def _supports_reasoning(self, model: str) -> bool:
try:
if litellm.supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider):
if "reasoning_effort" not in base_params:
base_params.append("reasoning_effort")
return litellm.supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider)
except Exception as e:
verbose_logger.debug("BedrockMantleChatConfig: error checking reasoning support: %s", e)
return base_params
return False
def get_model_response_iterator(
self,

6
litellm/llms/dashscope/chat/transformation.py
View file

@ -12,6 +12,12 @@ from ...openai.chat.gpt_transformation import OpenAIGPTConfig
class DashScopeChatConfig(OpenAIGPTConfig):
def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: base class contract returns a list
return [ # mutable-ok: base class contract returns a list
*super().get_supported_openai_params(model=model),
"reasoning_effort",
]
def remove_cache_control_flag_from_messages_and_tools(
self,
model: str,

4
litellm/llms/openai/chat/gpt_5_transformation.py
View file

@ -4,9 +4,9 @@ from typing import Final
import litellm
from litellm.utils import (
_is_explicitly_disabled_factory,
_supports_factory,
declared_value_factory,
is_explicitly_disabled_factory,
)
from .gpt_transformation import OpenAIGPTConfig
@ -192,7 +192,7 @@ class OpenAIGPT5Config(OpenAIGPTConfig):
Use this for opt-out checks where unknown models should be allowed through.
"""
return _is_explicitly_disabled_factory(
return is_explicitly_disabled_factory(
model=cls._model_map_lookup_name(model),
custom_llm_provider=None,
key=f"supports_{level}_reasoning_effort",

12
litellm/llms/openai/chat/guardrail_translation/handler.py
View file

@ -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]:

15
litellm/llms/openai/responses/guardrail_translation/handler.py
View file

@ -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

37
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
View file

@ -79,6 +79,7 @@ from litellm.utils import (
CustomStreamWrapper,
ModelResponse,
is_base64_encoded,
is_explicitly_disabled_factory,
supports_reasoning,
)
@ -866,6 +867,14 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
else:
raise _unsupported_reasoning_effort(reasoning_effort)
@staticmethod
def _supports_minimal_thinking_level(model: str) -> bool:
lowered: Final = model.lower()
is_gemini3flash: Final = "gemini-3" in lowered and "flash" in lowered
return is_gemini3flash and not is_explicitly_disabled_factory(
model=model, custom_llm_provider=None, key="supports_minimal_reasoning_effort"
)
@staticmethod
def _map_reasoning_effort_to_thinking_level(
reasoning_effort: str,
@ -880,13 +889,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
Returns:
GeminiThinkingConfig with thinkingLevel and includeThoughts
"""
# Check if this is gemini-3-flash which supports MINIMAL thinking level
# Covers gemini-3-flash, gemini-3-flash-preview, gemini-3.1-flash, gemini-3.1-flash-lite-preview,
# gemini-3.5-flash, and any future 3.x-flash variants.
is_gemini3flash: Final = model and ("flash" in model.lower() and "gemini-3" in model.lower())
supports_minimal: Final = bool(model) and VertexGeminiConfig._supports_minimal_thinking_level(model)
is_gemini31pro: Final = model and ("gemini-3.1-pro-preview" in model.lower())
if reasoning_effort == "minimal":
if is_gemini3flash:
if supports_minimal:
return {"thinkingLevel": "minimal", "includeThoughts": True}
else:
return {"thinkingLevel": "low", "includeThoughts": True}
@ -899,18 +906,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
return {"thinkingLevel": "high", "includeThoughts": True}
elif reasoning_effort == "high":
return {"thinkingLevel": "high", "includeThoughts": True}
elif reasoning_effort == "disable":
# Gemini 3 cannot fully disable thinking, so we use "minimal" for gemini-3-flash-preview, "low" for others
if is_gemini3flash:
return {"thinkingLevel": "minimal", "includeThoughts": False}
else:
return {"thinkingLevel": "low", "includeThoughts": False}
elif reasoning_effort == "none":
# For gemini-3-flash-preview, use "minimal" instead of "low"
if is_gemini3flash:
return {"thinkingLevel": "minimal", "includeThoughts": False}
else:
return {"thinkingLevel": "low", "includeThoughts": False}
elif reasoning_effort in ("disable", "none"):
return {
"thinkingLevel": "minimal" if supports_minimal else "low",
"includeThoughts": False,
}
else:
raise _unsupported_reasoning_effort(reasoning_effort)
@ -977,8 +977,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
params["includeThoughts"] = True
# Follow provider defaults unless explicitly opted into legacy behavior.
if litellm.enable_gemini_default_thinking_level_low is True:
is_gemini3flash: Final = "gemini-3" in model.lower() and "flash" in model.lower()
params["thinkingLevel"] = "minimal" if is_gemini3flash else "low"
params["thinkingLevel"] = (
"minimal" if VertexGeminiConfig._supports_minimal_thinking_level(model) else "low"
)
else:
# Thinking disabled
params["includeThoughts"] = False

1160
litellm/model_prices_and_context_window_backup.json
View file

File diff suppressed because it is too large Load diff

2
litellm/models/verification_token.py
View file

@ -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

112
litellm/ocr/input.py
View file

@ -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)

7
litellm/ocr/legacy.py
View file

@ -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

2
litellm/ocr/main.py
View file

@ -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

51
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -102,6 +102,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials import (
UpstreamCredentialProvider,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import (
prepare_mcp_client,
raise_public,
raise_token_exchange_challenge,
raise_user_oauth_challenge,
@ -2804,6 +2805,8 @@ class MCPServerManager:
headers=headers,
server_label=server.name or server.server_name or server.alias or server.server_id,
relays_upstream_auth=server.is_client_forwarded_token,
auth_type=server.auth_type,
upstream_token_header=server.upstream_token_header,
)
tool_func.__name__ = prefixed_tool_name
tool_func.__doc__ = description
@ -4259,15 +4262,20 @@ class MCPServerManager:
user_api_key_auth=user_api_key_auth,
extra_headers=extra_headers,
)
return MCPClient(
server_url=server_url,
transport_type=transport,
auth_type=resolved_server.auth_type,
timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT),
extra_headers=extra_headers,
resolved_auth=resolved_auth,
sampling_callback=sampling_cb,
elicitation_callback=elicitation_cb,
return await prepare_mcp_client(
resolved_server,
MCPClient(
server_url=server_url,
transport_type=transport,
auth_type=resolved_server.auth_type,
timeout=(
resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT
),
extra_headers=extra_headers,
resolved_auth=resolved_auth,
sampling_callback=sampling_cb,
elicitation_callback=elicitation_cb,
),
)
# Create SigV4 auth if configured
@ -4297,17 +4305,20 @@ class MCPServerManager:
else AuthResolution.no_auth
)
record_auth_resolution(server.server_id, legacy_source)
return MCPClient(
server_url=server_url,
transport_type=transport,
auth_type=resolved_server.auth_type,
auth_value=auth_value,
auth_header_name=auth_header_name,
timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT),
extra_headers=extra_headers,
aws_auth=aws_auth,
sampling_callback=sampling_cb,
elicitation_callback=elicitation_cb,
return await prepare_mcp_client(
resolved_server,
MCPClient(
server_url=server_url,
transport_type=transport,
auth_type=resolved_server.auth_type,
auth_value=auth_value,
auth_header_name=auth_header_name,
timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT),
extra_headers=extra_headers,
aws_auth=aws_auth,
sampling_callback=sampling_cb,
elicitation_callback=elicitation_cb,
),
)
async def _get_tools_from_server(

16
litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py
View file

@ -54,7 +54,7 @@ from litellm.llms.custom_httpx.http_handler import (
from litellm.proxy._experimental.mcp_server.tool_registry import (
global_mcp_tool_registry,
)
from litellm.types.mcp import credential_redirect_hook, custom_credential_slot
from litellm.types.mcp import MCPAuthType, credential_redirect_hook, custom_credential_slot
class _OpenAPIJSONSchema(TypedDict, total=False):
@ -471,6 +471,8 @@ def create_tool_function(
headers: dict[str, str] | None = None,
server_label: str | None = None,
relays_upstream_auth: bool = False,
auth_type: MCPAuthType = None,
upstream_token_header: str | None = None,
):
"""Create a tool function for an OpenAPI operation.
@ -503,6 +505,18 @@ def create_tool_function(
by using **kwargs instead of named parameters.
"""
effective_headers: Final = _merge_openapi_tool_request_headers(headers)
if auth_type is not None:
from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import (
raise_public,
validate_static_credential,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok
match validate_static_credential(auth_type, effective_headers, upstream_token_header, headers or ()):
case Error(error):
raise_public(error)
case Ok():
pass
# Build URL from base_url and path
url = base_url + path

101
litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py
View file

@ -13,15 +13,17 @@ from __future__ import annotations
import base64
import os
from collections.abc import Iterable, Mapping
from typing import TYPE_CHECKING, Final, Literal, NoReturn
from fastapi import HTTPException
from pydantic import SecretStr
from typing_extensions import assert_never
from litellm.experimental_mcp_client.client import strip_auth_scheme, to_basic_credentials
from litellm.experimental_mcp_client.client import MCPClient, strip_auth_scheme, to_basic_credentials
from litellm.proxy._experimental.mcp_server.exceptions import MCPServerURLCredentialsError
from litellm.proxy._experimental.mcp_server.oauth_utils import resolve_upstream_resource
from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok, Result
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
DEFAULT_CREDENTIAL_HEADER,
ApiKeyConfig,
@ -39,7 +41,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
Subject,
TokenExchangeConfig,
)
from litellm.types.mcp import DEFAULT_SUBJECT_TOKEN_TYPE, MCPAuth
from litellm.types.mcp import DEFAULT_SUBJECT_TOKEN_TYPE, MCPAuth, MCPAuthType, MCPTransport
if TYPE_CHECKING:
from litellm.proxy._types import UserAPIKeyAuth
@ -79,7 +81,7 @@ def to_server_spec(server: MCPServer) -> ServerSpec | None:
BYOK is the per-user source of the ``api_key`` mode; its scheme rides on ``auth_type`` just
like a shared key, but the value is per-user and not migrated yet, so a BYOK server defers
to v1 regardless of ``auth_type`` (this guard is the seam the BYOK arm replaces later).
to v1 for its static schemes. Declared OBO always stays with the exchange arm.
Dispatches on the declared ``auth_type``. The match is exhaustive over ``MCPAuthType`` with
an ``assert_never`` tail, so a newly added auth mode fails the type gate here until it is
@ -90,8 +92,8 @@ def to_server_spec(server: MCPServer) -> ServerSpec | None:
modes ``true_passthrough`` / ``oauth_delegate`` (``PassthroughConfig``); delegated/passthrough
oauth2 and SigV4 return None and stay on v1.
"""
if server.is_byok:
return None # per-user BYOK source not migrated yet -> defer to v1 (any auth_type)
if server.is_byok and server.auth_type != MCPAuth.oauth2_token_exchange:
return None # per-user BYOK source not migrated yet -> defer to v1
resource: Final = server.url or server.server_id
auth_type: Final = server.auth_type
match auth_type:
@ -165,21 +167,9 @@ def _client_credentials_spec(server: MCPServer, resource: str) -> ServerSpec:
)
def _token_exchange_spec(server: MCPServer, resource: str) -> ServerSpec | None:
"""Build a token_exchange (OBO) spec, or defer (None) when it is not OBO-configured.
An OBO server with ``client_id``/``client_secret`` is owned by the v2 arm even if the
``token_exchange_endpoint``/``token_url`` is absent: a missing endpoint then fails closed (412) at
the exchanger rather than silently deferring to v1 and connecting unauthenticated, since the
gateway must not guess the IdP or fall back to a weaker source. Without client credentials there is
nothing to own, so the server stays on v1 (parity-safe). ``profile`` selects the wire dialect
(``rfc8693`` default, ``entra_obo`` for Microsoft Entra On-Behalf-Of); an unrecognized value
normalizes to ``rfc8693`` so a bad config value cannot crash spec-building. ``audience`` is
forwarded only when the operator set it; a missing one is omitted, not derived.
"""
def _token_exchange_spec(server: MCPServer, resource: str) -> ServerSpec:
"""Keep declared OBO owned by the resolver, including incomplete client configuration."""
endpoint: Final = server.token_exchange_endpoint or server.effective_token_url
if not server.client_id or not server.client_secret:
return None
profile: Final[Literal["rfc8693", "entra_obo"]] = (
"entra_obo" if server.token_exchange_profile == "entra_obo" else "rfc8693"
)
@ -193,7 +183,7 @@ def _token_exchange_spec(server: MCPServer, resource: str) -> ServerSpec | None:
token_exchange_endpoint=endpoint,
audience=server.audience,
client_id=server.client_id,
client_secret=SecretStr(server.client_secret),
client_secret=SecretStr(server.client_secret) if server.client_secret else None,
token_endpoint_auth_method=server.token_endpoint_auth_method,
scopes=tuple(server.scopes or ()),
),
@ -397,3 +387,74 @@ def raise_token_exchange_challenge(
detail="Unauthorized",
headers={"WWW-Authenticate": www_authenticate},
)
_STATIC_MODES: Final = frozenset(
(MCPAuth.api_key, MCPAuth.bearer_token, MCPAuth.basic, MCPAuth.token, MCPAuth.authorization)
)
def _usable_credential_value(auth_type: MCPAuthType, name: str, value: str) -> bool:
if not value:
return False
if auth_type == MCPAuth.api_key and name != "authorization":
return True
if value.lower() in ("bearer", "basic", "token", "apikey"):
return False
if auth_type == MCPAuth.api_key:
api_scheme: Final = value.split(None, 1)[0]
if api_scheme.lower() in ("bearer", "token", "apikey"):
api_credential: Final = strip_auth_scheme(value, api_scheme).strip()
return api_credential.lower() != api_scheme.lower()
if auth_type in (MCPAuth.bearer_token, MCPAuth.token):
scheme: Final = "Bearer" if auth_type == MCPAuth.bearer_token else "token"
credential: Final = strip_auth_scheme(value, scheme).strip()
return bool(credential) and credential.lower() != scheme.lower()
if auth_type == MCPAuth.basic:
parts: Final = value.split(None, 1)
if len(parts) != 2 or parts[0].lower() != "basic":
return False
try:
decoded: Final = base64.b64decode(parts[1], validate=True).strip()
return b":" in decoded
except ValueError:
return False
return True
def validate_static_credential(
auth_type: MCPAuthType,
headers: Mapping[str, str],
upstream_token_header: str | None = None,
static_header_names: Iterable[str] = (),
) -> Result[None, CredError]:
if auth_type not in _STATIC_MODES:
return Ok(None)
default_slot: Final = "X-API-Key" if auth_type == MCPAuth.api_key else "Authorization"
admin_chosen_slots: Final = tuple(static_header_names) if auth_type == MCPAuth.api_key else ()
slots: Final = frozenset(
name.lower()
for name in (
upstream_token_header or default_slot,
default_slot,
"Authorization",
*admin_chosen_slots,
)
)
values: Final = tuple((name.lower(), value.strip()) for name, value in headers.items() if name.lower() in slots)
if any(_usable_credential_value(auth_type, name, value) for name, value in values):
return Ok(None)
return Error(CredError.of_misconfigured(f"{auth_type} requires a usable upstream credential"))
async def prepare_mcp_client(server: MCPServer, client: MCPClient) -> MCPClient:
if server.auth_type not in _STATIC_MODES or client.transport_type == MCPTransport.stdio:
return client
request: Final = await client.prepare_request_auth()
match validate_static_credential(
server.auth_type, request.headers, server.upstream_token_header, server.static_headers or ()
):
case Error(error):
raise_public(error)
case Ok():
return client

44
litellm/proxy/_types.py
View file

@ -844,6 +844,9 @@ class LiteLLMRoutes(enum.Enum):
)
self_managed_routes = [
# update_team resolves proxy/org/team admin itself and filters team admins
# through the team_admin_editable_team_fields setting
"/team/update",
"/team/member_add",
"/team/member_delete",
"/management/v1/teams/{team_id}/members/bulk_delete",
@ -2004,6 +2007,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 +2115,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 +3049,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 +3728,7 @@ class AllCallbacks(LiteLLMPydanticObjectBase):
"AWS_ACCESS_KEY_ID",
"AWS_SECRET_ACCESS_KEY",
"AWS_REGION_NAME",
"S3_LOG_PROMPTS_ONLY",
],
)
@ -4451,6 +4470,29 @@ class TeamInfoMember(Member):
user_alias: str | None = None
class TeamEditUnrestricted(BaseModel):
kind: Literal["unrestricted"] = "unrestricted"
class TeamEditAsTeamAdmin(BaseModel):
kind: Literal["team_admin"] = "team_admin"
editable_fields: tuple[str, ...]
class TeamEditAsTeamAdminDisabled(BaseModel):
kind: Literal["team_admin_disabled"] = "team_admin_disabled"
class TeamEditNone(BaseModel):
kind: Literal["none"] = "none"
TeamEditAccess = Annotated[
TeamEditUnrestricted | TeamEditAsTeamAdmin | TeamEditAsTeamAdminDisabled | TeamEditNone,
Field(discriminator="kind"),
]
class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable):
members_with_roles: tuple[TeamInfoMember, ...] = ()
team_member_budget_table: LiteLLM_BudgetTableFull | None = None
@ -4462,6 +4504,8 @@ 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
caller_edit_access: TeamEditAccess = Field(default_factory=TeamEditNone)
class TeamInfoResponseObject(TypedDict):

15
litellm/proxy/anthropic_endpoints/endpoints.py
View file

@ -8,7 +8,6 @@ from fastapi import APIRouter, Depends, HTTPException, Request, Response
from fastapi.responses import JSONResponse
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.anthropic_interface.exceptions import AnthropicErrorResponse, AnthropicExceptionMapping
from litellm.integrations.custom_guardrail import ModifyResponseException
from litellm.llms.anthropic.experimental_pass_through.context_management import (
@ -22,13 +21,16 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_request_processing import (
ProxyBaseLLMRequestProcessing,
create_response,
log_llm_api_exception,
proxy_exception_from_http_exception,
resolve_litellm_call_id,
)
from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
from litellm.proxy.common_utils.openai_error_payload import (
error_status_code,
openai_error_param,
openai_error_type,
with_litellm_call_id,
)
from litellm.types.utils import TokenCountResponse
@ -218,10 +220,12 @@ async def anthropic_response(
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict, original_exception=e, request_data=base_llm_response_processor.data
)
verbose_proxy_logger.exception("litellm.proxy.proxy_server.anthropic_response(): Exception occured - %s", e)
log_llm_api_exception(e, base_llm_response_processor.litellm_call_id)
if isinstance(e, ProxyException):
return _anthropic_error_json_response(e, request)
return _anthropic_error_json_response(
with_litellm_call_id(e, base_llm_response_processor.litellm_call_id), request
)
# Extract model_id from request metadata (same as success path)
litellm_metadata: Final = data.get("litellm_metadata", {}) or {}
@ -231,7 +235,7 @@ async def anthropic_response(
# Get headers
headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=user_api_key_dict,
call_id=data.get("litellm_call_id", ""),
call_id=base_llm_response_processor.litellm_call_id,
model_id=model_id,
version=version,
response_cost=0,
@ -288,6 +292,7 @@ async def count_tokens(
"""
from litellm.proxy.proxy_server import token_counter as internal_token_counter
litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))
try:
request_data: Final = await _read_request_body(request=request)
data: Final[dict] = {**request_data}
@ -339,7 +344,7 @@ async def count_tokens(
detail=detail,
)
except Exception as e:
verbose_proxy_logger.exception("litellm.proxy.anthropic_endpoints.count_tokens(): Exception occurred - %s", e)
log_llm_api_exception(e, litellm_call_id)
raise HTTPException(status_code=500, detail={"error": f"Internal server error: {e}"})

166
litellm/proxy/auth/fallback_budget.py Normal file
View file

@ -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)

2
litellm/proxy/auth/team_grants.py
View file

@ -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,

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

@ -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,

30
litellm/proxy/batches_endpoints/endpoints.py
View file

@ -7,6 +7,7 @@
import asyncio
import os
from collections.abc import Mapping
from types import MappingProxyType
from typing import Any, Final, cast
from fastapi import APIRouter, Depends, HTTPException, Path, Request, Response
@ -17,7 +18,11 @@ from litellm.batches.main import CancelBatchRequest, RetrieveBatchRequest
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.batches_endpoints.common_utils import validate_batch_list_limit
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.proxy.common_request_processing import (
ProxyBaseLLMRequestProcessing,
log_llm_api_exception,
request_litellm_call_id,
)
from litellm.proxy.common_utils.callback_utils import sanitize_openai_provider_metadata
from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
from litellm.proxy.common_utils.openai_endpoint_utils import (
@ -383,8 +388,9 @@ async def create_batch(
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
)
verbose_proxy_logger.exception("litellm.proxy.proxy_server.create_batch(): Exception occured - %s", e)
raise handle_exception_on_proxy(e)
litellm_call_id: Final = request_litellm_call_id(data)
log_llm_api_exception(e, litellm_call_id)
raise handle_exception_on_proxy(e, litellm_call_id)
@router.get(
@ -674,8 +680,9 @@ async def retrieve_batch(
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
)
verbose_proxy_logger.exception("litellm.proxy.proxy_server.retrieve_batch(): Exception occured - %s", e)
raise handle_exception_on_proxy(e)
litellm_call_id: Final = request_litellm_call_id(data)
log_llm_api_exception(e, litellm_call_id)
raise handle_exception_on_proxy(e, litellm_call_id)
@router.get(
@ -725,6 +732,7 @@ async def list_batches(
)
verbose_proxy_logger.debug("GET /v1/batches after=%s limit=%s", after, limit)
data: Mapping[str, object] = MappingProxyType({})
try:
if llm_router is None:
raise HTTPException(
@ -854,10 +862,11 @@ async def list_batches(
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict,
original_exception=e,
request_data={"after": after, "limit": limit},
request_data={**data, "after": after, "limit": limit},
)
verbose_proxy_logger.error("litellm.proxy.proxy_server.retrieve_batch(): Exception occured - %s", e)
raise handle_exception_on_proxy(e)
litellm_call_id: Final = request_litellm_call_id(data)
log_llm_api_exception(e, litellm_call_id)
raise handle_exception_on_proxy(e, litellm_call_id)
@router.post(
@ -1079,8 +1088,9 @@ async def cancel_batch(
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
)
verbose_proxy_logger.exception("litellm.proxy.proxy_server.create_batch(): Exception occured - %s", e)
raise handle_exception_on_proxy(e)
litellm_call_id: Final = request_litellm_call_id(data)
log_llm_api_exception(e, litellm_call_id)
raise handle_exception_on_proxy(e, litellm_call_id)
######################################################################

103
litellm/proxy/common_request_processing.py
View file

@ -7,7 +7,18 @@ from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequen
from datetime import datetime
from functools import lru_cache
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, Protocol, TypeAlias, TypeVar, overload
from typing import (
TYPE_CHECKING,
Any,
Final,
Literal,
NamedTuple,
Protocol,
TypeAlias,
TypeVar,
overload,
runtime_checkable,
)
import anyio
import httpx
@ -34,7 +45,11 @@ from litellm.constants import (
UNSAFE_PROXY_RESPONSE_HEADERS,
)
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket, is_expected_client_error
from litellm.litellm_core_utils.core_helpers import (
get_or_create_metadata_bucket,
independent_snapshot,
is_expected_client_error,
)
from litellm.litellm_core_utils.dd_tracing import NullTracer, tracer
from litellm.litellm_core_utils.get_supported_openai_params import (
get_supported_openai_params,
@ -1452,7 +1467,19 @@ def _has_attribute_error_in_chain(exc: Exception) -> bool:
_CLIENT_DISCONNECT_DETAIL: Final = "Client disconnected the request"
def _log_llm_api_exception(e: Exception, litellm_call_id: str | None) -> None:
@runtime_checkable
class _CarriesLitellmCallId(Protocol):
litellm_call_id: str | None
def request_litellm_call_id(data: Mapping[str, object]) -> str | None:
logging_obj: Final = data.get("litellm_logging_obj")
logged_id: Final = logging_obj.litellm_call_id if isinstance(logging_obj, _CarriesLitellmCallId) else None
call_id: Final = logged_id or data.get("litellm_call_id")
return call_id if isinstance(call_id, str) else None
def log_llm_api_exception(e: Exception, litellm_call_id: str | None) -> None:
if getattr(e, "status_code", None) == 499 and getattr(e, "detail", None) == _CLIENT_DISCONNECT_DETAIL:
verbose_proxy_logger.info(
"litellm.proxy.proxy_server._handle_llm_api_exception(): client disconnected, "
@ -1532,6 +1559,10 @@ class ProxyBaseLLMRequestProcessing:
def __init__(self, data: dict):
self.data = data
@property
def litellm_call_id(self) -> str | None:
return request_litellm_call_id(self.data)
@staticmethod
def _merge_passthrough_streaming_headers(
response_headers: httpx.Headers | dict | None,
@ -2062,6 +2093,13 @@ class ProxyBaseLLMRequestProcessing:
) -> tuple[dict, LiteLLMLoggingObj]:
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
configured_fallbacks: Final = (
self._configured_fallbacks(llm_router=llm_router, user_api_key_dict=user_api_key_dict)
if llm_router is not None and not self.data.get("disable_fallbacks")
else None
)
pristine: Final = independent_snapshot(self.data) if configured_fallbacks else None
try:
return await self.common_processing_pre_call_logic(
request=request,
@ -2080,14 +2118,19 @@ class ProxyBaseLLMRequestProcessing:
llm_router=llm_router,
)
except ProxyRateLimitError as original_exc:
original_model: Final = self.data.get("model")
if not original_model or not llm_router or self.data.get("disable_fallbacks"):
rate_limited_data: Final = self.data
original_model: Final = rate_limited_data.get("model")
if (
pristine is None
or not configured_fallbacks
or rate_limited_data.get("disable_fallbacks")
or not isinstance(original_model, str)
):
raise
fallback_models: Final = self._resolve_fallback_models(
model=original_model,
llm_router=llm_router,
user_api_key_dict=user_api_key_dict,
fallbacks=configured_fallbacks,
)
if not fallback_models:
raise
@ -2102,6 +2145,7 @@ class ProxyBaseLLMRequestProcessing:
for fallback_model in fallback_models:
if fallback_model == original_model:
continue
self.data = independent_snapshot(pristine)
self.data["model"] = fallback_model
try:
return await self.common_processing_pre_call_logic(
@ -2123,39 +2167,30 @@ class ProxyBaseLLMRequestProcessing:
except ProxyRateLimitError:
continue
except BaseException:
self.data["model"] = original_model
self.data = rate_limited_data
raise
self.data["model"] = original_model
self.data = rate_limited_data
raise original_exc
def _resolve_fallback_models(
self,
model: str,
llm_router: Router,
user_api_key_dict: UserAPIKeyAuth,
) -> list | None:
from litellm.router_utils.fallback_event_handlers import get_fallback_model_group
fallbacks = None
@staticmethod
def _configured_fallbacks(llm_router: Router, user_api_key_dict: UserAPIKeyAuth) -> list | None:
key_router_settings: Final = user_api_key_dict.router_settings
if isinstance(key_router_settings, dict) and "fallbacks" in key_router_settings:
fallbacks = key_router_settings["fallbacks"]
key_fallbacks: Final = key_router_settings.get("fallbacks") if isinstance(key_router_settings, dict) else None
fallbacks: Final = key_fallbacks if key_fallbacks is not None else llm_router.fallbacks
return fallbacks if isinstance(fallbacks, list) and fallbacks else None
if fallbacks is None:
fallbacks = llm_router.fallbacks
if not fallbacks:
return None
@staticmethod
def _resolve_fallback_models(model: str, fallbacks: list) -> list | None:
from litellm.router_utils.fallback_event_handlers import get_fallback_model_group
fallback_model_group, generic_fallback_idx = get_fallback_model_group(
fallbacks=fallbacks,
model_group=model,
)
if fallback_model_group is None and generic_fallback_idx is not None:
fallback_model_group = fallbacks[generic_fallback_idx]["*"]
return fallback_model_group
if fallback_model_group is not None:
return fallback_model_group
return fallbacks[generic_fallback_idx]["*"] if generic_fallback_idx is not None else None
@staticmethod
def _get_model_id_from_response(hidden_params: Mapping[str, object], data: Mapping[str, object]) -> str:
@ -3429,11 +3464,7 @@ class ProxyBaseLLMRequestProcessing:
version: str | None = None,
):
"""Raises ProxyException (OpenAI API compatible) if an exception is raised"""
logging_obj: Final[LiteLLMLoggingObj | None] = self.data.get("litellm_logging_obj", None)
_log_llm_api_exception(
e,
(logging_obj.litellm_call_id if logging_obj is not None else None) or self.data.get("litellm_call_id"),
)
log_llm_api_exception(e, self.litellm_call_id)
# Allow callbacks to transform the error response
transformed_exception: Final = await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict,
@ -3463,9 +3494,7 @@ class ProxyBaseLLMRequestProcessing:
custom_headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=user_api_key_dict,
call_id=(
_litellm_logging_obj.litellm_call_id if _litellm_logging_obj else self.data.get("litellm_call_id")
),
call_id=self.litellm_call_id,
model_id=model_id,
version=version,
response_cost=0,

23
litellm/proxy/common_utils/openai_error_payload.py
View file

@ -9,6 +9,9 @@ from typing import Final
from fastapi import status
from litellm.constants import STRINGIFIED_NONE
from litellm.proxy._types import ProxyException
LITELLM_CALL_ID_HEADER: Final = "x-litellm-call-id"
_OPENAI_ERROR_TYPE_BY_STATUS: Final[Mapping[int, str]] = MappingProxyType(
{
@ -52,3 +55,23 @@ def openai_error_param(exc: object) -> str | None:
serializes as JSON ``null``."""
carried: Final = attribute_of(exc, "param")
return carried if isinstance(carried, str) and carried != STRINGIFIED_NONE else None
def litellm_call_id_headers(litellm_call_id: str | None) -> dict[str, str] | None: # mutable-ok: ProxyException.headers
if litellm_call_id is None:
return None
return {LITELLM_CALL_ID_HEADER: litellm_call_id} # mutable-ok: ProxyException mutates its headers dict
def with_litellm_call_id(exc: ProxyException, litellm_call_id: str | None) -> ProxyException:
"""The same error object, answering with ``x-litellm-call-id`` when it was raised without one."""
if litellm_call_id is not None:
exc.headers.setdefault(LITELLM_CALL_ID_HEADER, litellm_call_id)
return exc
def headers_with_litellm_call_id(headers: Mapping[str, str] | None, litellm_call_id: str) -> Mapping[str, str]:
"""``headers`` plus ``x-litellm-call-id``, keeping the value they already carry under that name."""
if headers is None:
return MappingProxyType({LITELLM_CALL_ID_HEADER: litellm_call_id})
return MappingProxyType({LITELLM_CALL_ID_HEADER: litellm_call_id, **headers})

2
litellm/proxy/common_utils/proxy_rate_limit_error.py
View file

@ -11,7 +11,7 @@ exception types:
an upstream LLM provider returns 429.
* :class:`fastapi.HTTPException` (status 429) — raised directly by proxy hooks
such as ``parallel_request_limiter``, ``dynamic_rate_limiter``,
``batch_rate_limiter``, ``max_budget_limiter``, ``max_iterations_limiter``,
``batch_rate_limiter``, ``max_iterations_limiter``,
etc.
* :class:`litellm.llms.base_llm.chat.transformation.BaseLLMException` (status
429) — raised by some provider transports.

1
litellm/proxy/db/create_views.py
View file

@ -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,

10
litellm/proxy/db/db_spend_update_writer.py
View file

@ -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),
},
)

6
litellm/proxy/guardrails/guardrail_endpoints.py
View file

@ -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)

9
litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
View file

@ -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]:

2
litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/__init__.py
View file

@ -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,
)

6
litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py
View file

@ -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

77
litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py
View file

@ -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(

1
litellm/proxy/guardrails/guardrail_initializers.py
View file

@ -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

4
litellm/proxy/hooks/__init__.py
View file

@ -4,7 +4,6 @@ from typing import Final, Literal
from . import *
from .cache_control_check import _PROXY_CacheControlCheck
from .litellm_skills import SkillsInjectionHook
from .max_budget_limiter import _PROXY_MaxBudgetLimiter
from .max_budget_per_session_limiter import _PROXY_MaxBudgetPerSessionHandler
from .max_iterations_limiter import _PROXY_MaxIterationsHandler
from .parallel_request_limiter import _PROXY_MaxParallelRequestsHandler
@ -18,7 +17,6 @@ from .sensitive_data_routing import _PROXY_SensitiveDataRoutingHandler
# transitively through `enterprise.enterprise_hooks` can resolve `PROXY_HOOKS`
# and `get_proxy_hook` from this partially-initialized module without circling.
PROXY_HOOKS: Final = {
"max_budget_limiter": _PROXY_MaxBudgetLimiter,
"parallel_request_limiter": _PROXY_MaxParallelRequestsHandler_v3,
"cache_control_check": _PROXY_CacheControlCheck,
"responses_id_security": ResponsesIDSecurity,
@ -35,7 +33,7 @@ if os.getenv("LEGACY_MULTI_INSTANCE_RATE_LIMITING", "false").lower() == "true":
def get_proxy_hook(
hook_name: Literal["max_budget_limiter", "managed_files", "parallel_request_limiter", "cache_control_check"] | str,
hook_name: Literal["managed_files", "parallel_request_limiter", "cache_control_check"] | str,
):
"""
Factory method to get a proxy hook instance by name

84
litellm/proxy/hooks/max_budget_limiter.py
View file

@ -1,84 +0,0 @@
from typing import Final
from fastapi import HTTPException
from litellm import verbose_logger
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
from litellm.exceptions import RateLimitType
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit
class _PROXY_MaxBudgetLimiter(CustomLogger):
# Class variables or attributes
def __init__(self):
pass
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: str,
):
try:
verbose_proxy_logger.debug("Inside Max Budget Limiter Pre-Call Hook")
max_budget: Final = user_api_key_dict.user_max_budget
user_id: Final = user_api_key_dict.user_id
if max_budget is None or user_id is None:
return
from litellm.proxy.proxy_server import general_settings
if (
user_api_key_dict.team_id is not None
and general_settings.get("apply_user_budget_to_team_keys") is not True
):
return
# The reservation path admits at the strict-`<` boundary and
# atomically pre-fills the same counter we'd read here. Re-checking
# with `>=` would reject a request the reservation already admitted
# when the reservation fills the counter to exactly max_budget.
# Imported lazily to avoid a circular import via proxy.utils.
from litellm.proxy.spend_tracking.budget_reservation import (
get_reserved_counter_keys,
)
user_counter_key: Final = f"spend:user:{user_id}"
if user_counter_key in get_reserved_counter_keys(user_api_key_dict.budget_reservation):
return
from litellm.proxy.proxy_server import get_current_spend
curr_spend: Final = await get_current_spend(
counter_key=user_counter_key,
fallback_spend=user_api_key_dict.user_spend or 0.0,
)
verbose_proxy_logger.debug(
"MaxBudgetLimiter: user_id=%s, spend=%.6f, max=%.6f",
user_id,
curr_spend,
max_budget,
)
# CHECK IF REQUEST ALLOWED
if curr_spend >= max_budget:
resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(data.get("model") if data else None)
raise ProxyRateLimitError(
detail="Max budget limit reached.",
rate_limit_type=RateLimitType.BUDGET,
model=resolved_model,
llm_provider=llm_provider,
)
except HTTPException as e:
raise e
except Exception as e:
verbose_logger.exception(
"litellm.proxy.hooks.max_budget_limiter.py::async_pre_call_hook(): Exception occured - %s", e
)

58
litellm/proxy/hooks/model_max_budget_limiter.py
View file

@ -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

78
litellm/proxy/hooks/parallel_request_limiter_v3.py
View file

@ -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(

27
litellm/proxy/image_endpoints/endpoints.py
View file

@ -1,6 +1,5 @@
import asyncio
import io
import traceback
from collections.abc import Sequence
from typing import Final, get_type_hints
@ -9,19 +8,23 @@ from fastapi import APIRouter, Depends, File, HTTPException, Request, Response,
from fastapi.responses import ORJSONResponse
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.prompt_templates.common_utils import (
get_str_from_messages,
)
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
from litellm.proxy.common_request_processing import (
ProxyBaseLLMRequestProcessing,
log_llm_api_exception,
resolve_litellm_call_id,
)
from litellm.proxy.common_utils.http_parsing_utils import (
coerce_numeric_form_fields,
numeric_form_fields,
)
from litellm.proxy.common_utils.openai_error_payload import (
error_status_code,
litellm_call_id_headers,
openai_error_param,
openai_error_type,
)
@ -91,11 +94,12 @@ async def image_generation(
version,
)
data = {}
litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))
data = {"litellm_call_id": litellm_call_id}
try:
# Use orjson to parse JSON data, orjson speeds up requests significantly
body: Final = await request.body()
data = orjson.loads(body)
data = orjson.loads(body) | data
# Include original request and headers in the data
data = await add_litellm_data_to_request(
@ -153,9 +157,7 @@ async def image_generation(
response = await llm_call
### ALERTING ###
asyncio.create_task(
proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success")
)
asyncio.create_task(proxy_logging_obj.update_request_status(litellm_call_id=litellm_call_id, status="success"))
### CALL HOOKS ### - modify outgoing data (guardrails, otel, etc.)
response = await proxy_logging_obj.post_call_success_hook(
@ -168,7 +170,7 @@ async def image_generation(
cache_key: Final = hidden_params.get("cache_key", None) or ""
api_base: Final = hidden_params.get("api_base", None) or ""
response_cost: Final = hidden_params.get("response_cost", None) or ""
litellm_call_id: Final = hidden_params.get("litellm_call_id", None) or ""
response_call_id: Final = hidden_params.get("litellm_call_id", None) or ""
fastapi_response.headers.update(
ProxyBaseLLMRequestProcessing.get_custom_headers(
@ -179,7 +181,7 @@ async def image_generation(
version=version,
response_cost=response_cost,
model_region=getattr(user_api_key_dict, "allowed_model_region", ""),
call_id=litellm_call_id,
call_id=response_call_id,
request_data=data,
hidden_params=hidden_params,
)
@ -200,13 +202,13 @@ async def image_generation(
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
)
verbose_proxy_logger.error("litellm.proxy.proxy_server.image_generation(): Exception occured - %s", e)
verbose_proxy_logger.debug(traceback.format_exc())
log_llm_api_exception(e, litellm_call_id)
if isinstance(e, HTTPException):
raise ProxyException(
message=getattr(e, "message", str(e)),
type=openai_error_type(e, error_status_code(e, status.HTTP_400_BAD_REQUEST)),
param=openai_error_param(e),
headers=litellm_call_id_headers(litellm_call_id),
code=error_status_code(e, status.HTTP_400_BAD_REQUEST),
)
else:
@ -215,6 +217,7 @@ async def image_generation(
message=getattr(e, "message", error_msg),
type=openai_error_type(e, error_status_code(e, 500)),
param=openai_error_param(e),
headers=litellm_call_id_headers(litellm_call_id),
openai_code=getattr(e, "code", None),
code=error_status_code(e, 500),
)

10
litellm/proxy/litellm_pre_call_utils.py
View file

@ -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

51
litellm/proxy/management_endpoints/common_utils.py
View file

@ -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:

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

@ -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

191
litellm/proxy/management_endpoints/team_admin_field_permissions.py Normal file
View file

@ -0,0 +1,191 @@
"""Proxy-wide allow-list of team-settings fields a team admin may change on /team/update."""
from collections.abc import Mapping
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final, Literal, TypeAlias
from fastapi import HTTPException
from pydantic import TypeAdapter, ValidationError
from typing_extensions import assert_never
from litellm._logging import verbose_proxy_logger
from litellm.models.team import LiteLLM_TeamTable
from litellm.proxy._types import (
LiteLLM_ManagementEndpoint_MetadataFields,
LiteLLM_ManagementEndpoint_MetadataFields_Premium,
UpdateTeamRequest,
)
TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING: Final = "team_admin_editable_team_fields"
# TODO(LIT-5722): add the remaining team settings one per PR, each with its value-diff tests and dashboard field
SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS: Final[frozenset[str]] = frozenset({"tpm_limit"})
_FIELD_LIST: Final = TypeAdapter(list[str])
_JSON_OBJECT: Final = TypeAdapter(dict[str, object])
_EMPTY: Final[Mapping[str, object]] = MappingProxyType({})
_METADATA_FOLDED_FIELDS: Final[frozenset[str]] = frozenset(
(*LiteLLM_ManagementEndpoint_MetadataFields, *LiteLLM_ManagementEndpoint_MetadataFields_Premium)
)
_SYSTEM_MANAGED_METADATA_KEYS: Final[frozenset[str]] = frozenset({"team_member_budget_id"})
_NOT_COLUMNS: Final[frozenset[str]] = frozenset({"team_id", "metadata"})
_SETTINGS_LOCATION: Final = "Settings > UI > Team admin editable fields"
@dataclass(frozen=True, slots=True)
class TeamAdminEditAllowed:
request: UpdateTeamRequest
kind: Literal["allowed"] = "allowed"
@dataclass(frozen=True, slots=True)
class TeamAdminEditingDisabled:
kind: Literal["disabled"] = "disabled"
@dataclass(frozen=True, slots=True)
class TeamAdminFieldNotPermitted:
field: str
kind: Literal["field_not_permitted"] = "field_not_permitted"
TeamAdminEditVerdict: TypeAlias = TeamAdminEditAllowed | TeamAdminEditingDisabled | TeamAdminFieldNotPermitted
def resolve_team_admin_editable_fields(
general_settings: Mapping[str, object],
supported: frozenset[str],
) -> frozenset[str]:
raw: Final = general_settings.get(TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING)
if raw is None:
return frozenset()
try:
configured: Final = frozenset(_FIELD_LIST.validate_python(raw))
except ValidationError:
verbose_proxy_logger.warning(
"%s must be a list of field names; ignoring %r", TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING, raw
)
return frozenset()
unsupported: Final = configured - supported
if unsupported:
verbose_proxy_logger.warning(
"%s ignores unsupported field(s) %s; supported: %s",
TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING,
sorted(unsupported),
sorted(supported),
)
return configured & supported
def _as_object(value: object) -> Mapping[str, object]:
try:
return _JSON_OBJECT.validate_json(value) if isinstance(value, str) else _JSON_OBJECT.validate_python(value)
except ValidationError:
return _EMPTY
def _stored_metadata(existing: Mapping[str, object]) -> Mapping[str, object]:
return _as_object(existing.get("metadata"))
def _submitted_metadata(
data: UpdateTeamRequest, submitted: Mapping[str, object], existing: Mapping[str, object]
) -> Mapping[str, object]:
"""Metadata as it would be stored: the caller's dict (or the stored one) with top-level folded fields laid over."""
base: Final = (
_as_object(submitted.get("metadata")) if "metadata" in data.model_fields_set else _stored_metadata(existing)
)
folded: Final = data.model_fields_set & _METADATA_FOLDED_FIELDS
return MappingProxyType({key: submitted[key] if key in folded else base[key] for key in base.keys() | folded})
def _metadata_changes(
data: UpdateTeamRequest, submitted: Mapping[str, object], existing: Mapping[str, object]
) -> frozenset[str]:
merged: Final = _submitted_metadata(data, submitted, existing)
stored: Final = _stored_metadata(existing)
return frozenset(
key if key in _METADATA_FOLDED_FIELDS else "metadata"
for key in (merged.keys() | stored.keys()) - _SYSTEM_MANAGED_METADATA_KEYS
if merged.get(key) != stored.get(key)
)
def _stored_model_aliases(existing_row: LiteLLM_TeamTable) -> Mapping[str, object]:
table: Final = existing_row.litellm_model_table
return _as_object(_JSON_OBJECT.validate_json(table.model_dump_json()).get("model_aliases")) if table else _EMPTY
def _column_changed(
field: str, submitted: Mapping[str, object], existing: Mapping[str, object], existing_row: LiteLLM_TeamTable
) -> bool:
if field == "model_aliases":
return _as_object(submitted.get(field)) != _stored_model_aliases(existing_row)
if field in LiteLLM_TeamTable.model_fields:
return submitted.get(field) != existing.get(field)
return True
def changed_team_fields(data: UpdateTeamRequest, existing_row: LiteLLM_TeamTable) -> frozenset[str]:
"""Logical field names whose stored value the request would change.
Request and stored row are compared as JSON values so both sides share one representation. Fields the
server folds into metadata are attributed to their own name whether they arrive top-level or inside
``metadata``; anything else in ``metadata`` is attributed to ``metadata``. Fields with no stored
counterpart on the team row count as changed whenever they are sent.
"""
submitted: Final = _JSON_OBJECT.validate_json(data.model_dump_json(exclude_unset=True))
existing: Final = _JSON_OBJECT.validate_json(existing_row.model_dump_json())
column_fields: Final = frozenset(data.model_fields_set) - _NOT_COLUMNS - _METADATA_FOLDED_FIELDS
column_changes: Final = frozenset(
field for field in column_fields if _column_changed(field, submitted, existing, existing_row)
)
return column_changes | _metadata_changes(data, submitted, existing)
def _only_changes(data: UpdateTeamRequest, changed: frozenset[str]) -> UpdateTeamRequest:
"""The request without the values it resends unchanged, which would otherwise still trigger derived writes
such as a resent budget_duration pushing budget_reset_at back."""
sent: Final = frozenset(data.model_fields_set)
via_metadata: Final = frozenset({"metadata"}) if changed - sent else frozenset()
kept: Final = frozenset({"team_id"}) | (changed & sent) | via_metadata
return UpdateTeamRequest.model_validate(data.model_dump(include=MappingProxyType({field: True for field in kept})))
def team_admin_edit_verdict(
data: UpdateTeamRequest,
existing: LiteLLM_TeamTable,
permitted: frozenset[str],
) -> TeamAdminEditVerdict:
if not permitted:
return TeamAdminEditingDisabled()
changed: Final = changed_team_fields(data, existing)
blocked: Final = sorted(changed - permitted)
if blocked:
return TeamAdminFieldNotPermitted(field=blocked[0])
return TeamAdminEditAllowed(request=_only_changes(data, changed))
def team_admin_request_or_raise(verdict: TeamAdminEditVerdict) -> UpdateTeamRequest:
match verdict:
case TeamAdminEditAllowed(request=request):
return request
case TeamAdminEditingDisabled():
raise HTTPException(
status_code=403,
detail=(
"Team admins on this proxy cannot edit team settings. "
f"Ask a proxy admin to enable fields under {_SETTINGS_LOCATION}."
),
)
case TeamAdminFieldNotPermitted(field=field):
raise HTTPException(
status_code=403,
detail=(
f"Team admins on this proxy do not have permission to update '{field}'. "
f"Ask a proxy admin to add it under {_SETTINGS_LOCATION}."
),
)
case _:
assert_never(verdict)

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

@ -18,12 +18,23 @@ from collections.abc import Iterable, Mapping, Sequence
from collections.abc import Set as AbstractSet
from datetime import datetime, timezone
from types import MappingProxyType
from typing import TYPE_CHECKING, Annotated, Final, NamedTuple, NoReturn, Protocol, TypeVar, cast
from typing import (
TYPE_CHECKING,
Annotated,
Final,
Literal,
NamedTuple,
NoReturn,
Protocol,
TypeAlias,
TypeVar,
cast,
)
import fastapi
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
from pydantic import BaseModel, JsonValue
from typing_extensions import ReadOnly, TypedDict
from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict, assert_never
import litellm
from litellm._logging import verbose_proxy_logger
@ -38,6 +49,7 @@ from litellm.proxy._types import (
DeleteTeamRequest,
LiteLLM_AuditLogs,
LiteLLM_DeletedTeamTable,
Litellm_EntityType,
LiteLLM_ManagementEndpoint_MetadataFields,
LiteLLM_ManagementEndpoint_MetadataFields_Premium,
LiteLLM_ModelTable,
@ -61,6 +73,11 @@ from litellm.proxy._types import (
SpecialProxyStrings,
TeamAccessGroupModelGrant,
TeamAddMemberResponse,
TeamEditAccess,
TeamEditAsTeamAdmin,
TeamEditAsTeamAdminDisabled,
TeamEditNone,
TeamEditUnrestricted,
TeamInfoMember,
TeamInfoResponseObject,
TeamInfoResponseObjectTeamTable,
@ -95,6 +112,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 +129,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,
@ -116,6 +138,12 @@ from litellm.proxy.management_endpoints.router_weights import validate_router_se
from litellm.proxy.management_endpoints.tag_management_endpoints import (
get_daily_activity,
)
from litellm.proxy.management_endpoints.team_admin_field_permissions import (
SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS,
resolve_team_admin_editable_fields,
team_admin_edit_verdict,
team_admin_request_or_raise,
)
from litellm.proxy.management_helpers.access_group_team_sync import (
TEAM_ADVISORY_LOCK_SQL,
AccessGroupSyncTx,
@ -177,6 +205,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
@ -432,32 +461,70 @@ async def _refresh_cached_team(
)
async def _can_manage_team(
TeamAccessRole: TypeAlias = Literal["proxy_admin", "org_admin", "team_admin"]
def _raise_team_access_denied() -> NoReturn:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail="You do not have access to this team",
)
async def _resolve_team_access(
team_obj: LiteLLM_TeamTable,
user_api_key_dict: UserAPIKeyAuth,
) -> bool:
"""True for a proxy admin, an admin of this team, or an org admin for the team's organization."""
) -> TeamAccessRole | None:
"""Strongest role the caller holds over ``team_obj``, or None when they hold none.
Org admin outranks team admin so a caller holding both keeps unrestricted edits.
"""
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
return True
return "proxy_admin"
if await _is_user_org_admin_for_team(user_api_key_dict=user_api_key_dict, team_obj=team_obj):
return "org_admin"
if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_obj):
return True
return "team_admin"
return await _is_user_org_admin_for_team(user_api_key_dict=user_api_key_dict, team_obj=team_obj)
return None
async def _verify_team_access(
team_obj: LiteLLM_TeamTable,
user_api_key_dict: UserAPIKeyAuth,
) -> None:
"""Raise HTTPException(403) unless the caller can manage the given team."""
if await _can_manage_team(team_obj=team_obj, user_api_key_dict=user_api_key_dict):
return
"""Raise 403 unless the caller is a proxy admin, an org admin for the team's org, or a team admin."""
if await _resolve_team_access(team_obj=team_obj, user_api_key_dict=user_api_key_dict) is None:
_raise_team_access_denied()
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail="You do not have access to this team",
)
_GENERAL_SETTINGS: Final = TypeAdapter(dict[str, object])
def _general_settings() -> Mapping[str, object]:
from litellm.proxy.proxy_server import general_settings
return _GENERAL_SETTINGS.validate_python(general_settings)
def _caller_edit_access(role: TeamAccessRole | None, general_settings: Mapping[str, object]) -> TeamEditAccess:
"""What the caller may change on /team/update, reported on /team/info so the dashboard never re-derives it."""
match role:
case "proxy_admin" | "org_admin":
return TeamEditUnrestricted()
case "team_admin":
permitted: Final = resolve_team_admin_editable_fields(
general_settings, SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS
)
if not permitted:
return TeamEditAsTeamAdminDisabled()
return TeamEditAsTeamAdmin(editable_fields=tuple(sorted(permitted)))
case None:
return TeamEditNone()
case _:
assert_never(role)
class TeamMemberBudgetHandler:
@ -1170,6 +1237,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 +1353,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 +1415,7 @@ async def new_team(
general_settings,
litellm_proxy_admin_name,
llm_router,
premium_user,
prisma_client,
user_api_key_cache,
)
@ -1321,6 +1446,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 +2106,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 +2158,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,22 +2197,36 @@ 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}
)
if existing_team_row is None:
# Non-proxy-admins get the same 403 as an access denial so /team/update
# cannot be used to probe which team ids exist
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
_raise_team_access_denied()
raise HTTPException(
status_code=404,
detail={"error": f"Team not found, passed team_id={data.team_id}"},
)
# Verify caller has access to manage this team
await _verify_team_access(
team_obj=LiteLLM_TeamTable.model_validate(existing_team_row.model_dump()),
user_api_key_dict=user_api_key_dict,
)
existing_team: Final = LiteLLM_TeamTable.model_validate(existing_team_row.model_dump())
access_role: Final = await _resolve_team_access(team_obj=existing_team, user_api_key_dict=user_api_key_dict)
if access_role is None:
_raise_team_access_denied()
if access_role == "team_admin":
data = team_admin_request_or_raise( # rebind-ok: resent values must not reach the derived writes below
team_admin_edit_verdict(
data=data,
existing=existing_team,
permitted=resolve_team_admin_editable_fields(
_general_settings(), SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS
),
)
)
await validate_router_settings_weights(
data.router_settings,
@ -2188,6 +2330,7 @@ async def update_team(
org_id=org_id_to_check,
user_api_key_cache=user_api_key_cache,
prisma_client=prisma_client,
include_budget_table=True,
)
if org_table is not None:
await _check_org_team_limits(
@ -2204,8 +2347,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 +4623,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:
@ -4507,10 +4657,9 @@ async def team_info(
)
team_table: Final = LiteLLM_TeamTable.model_validate(team_info.model_dump())
await validate_membership(user_api_key_dict=user_api_key_dict, team_table=team_table)
access_role: Final = await _resolve_team_access(team_obj=team_table, user_api_key_dict=user_api_key_dict)
organization_models: Final[list[str] | None] = (
_parent_organization_models(team_info)
if await _can_manage_team(team_obj=team_table, user_api_key_dict=user_api_key_dict)
else None
_parent_organization_models(team_info) if access_role is not None else None
)
## GET ALL KEYS ##
@ -4573,6 +4722,13 @@ 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,
),
"caller_edit_access": _caller_edit_access(access_role, _general_settings()),
}
)

17
litellm/proxy/ocr_endpoints/endpoints.py
View file

@ -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,

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

@ -1180,9 +1180,8 @@ async def bedrock_proxy_route(
endpoint_func: Final = create_pass_through_route(
endpoint=endpoint,
target=str(prepped.url),
custom_headers=prepped.headers,
custom_headers=_upstream_headers_for_bedrock_agent_runtime_route(request, user_api_key_dict, prepped.headers),
is_streaming_request=is_streaming_request,
_forward_headers=True,
) # dynamically construct pass-through endpoint based on incoming path
setattr(request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, data)
# SigV4 signs an exact payload; pass-through must send prepped.body, not json.dumps
@ -2001,6 +2000,9 @@ _ANTHROPIC_UPSTREAM_CREDENTIAL_HEADERS: Final = frozenset({"authorization", "x-a
_HEADERS_NEVER_FORWARDED_TO_ANTHROPIC: Final = frozenset({"content-length", "host", "accept-encoding"}) | (
SpecialHeaders.litellm_credential_header_names() - _ANTHROPIC_UPSTREAM_CREDENTIAL_HEADERS
)
_HEADERS_NEVER_FORWARDED_TO_BEDROCK: Final = (
frozenset({"content-length", "host", "accept-encoding"}) | SpecialHeaders.litellm_credential_header_names()
)
_MAPPED_ROUTE_CALLER_KEY_HEADER: Final = "litellm_user_api_key"
@ -2099,6 +2101,17 @@ def _upstream_headers_for_anthropic_route(
return MappingProxyType({**caller_headers, **(proxy_auth_header or {})})
def _upstream_headers_for_bedrock_agent_runtime_route(
request: Request, user_api_key_dict: UserAPIKeyAuth, signed_headers: Mapping[str, object]
) -> Mapping[str, object]:
caller_headers: Final = _caller_headers_without_litellm_secrets(
request,
user_api_key_dict,
_HEADERS_NEVER_FORWARDED_TO_BEDROCK | frozenset(name.lower() for name in signed_headers),
)
return MappingProxyType({**caller_headers, **signed_headers})
async def _prepare_vertex_auth_headers(
request: Request,
vertex_credentials: VertexPassThroughCredentials | None,

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

@ -72,7 +72,9 @@ from litellm.proxy.auth.auth_utils import request_dispatched_to_pass_through_end
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_request_processing import (
ProxyBaseLLMRequestProcessing,
log_llm_api_exception,
open_sse_before_first_byte,
resolve_litellm_call_id,
)
from litellm.proxy.common_utils.http_parsing_utils import (
_read_request_body,
@ -80,6 +82,7 @@ from litellm.proxy.common_utils.http_parsing_utils import (
)
from litellm.proxy.common_utils.openai_error_payload import (
error_status_code,
litellm_call_id_headers,
openai_error_param,
openai_error_type,
)
@ -196,14 +199,15 @@ async def chat_completion_pass_through_endpoint(
version,
)
data = {}
litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))
data = {"litellm_call_id": litellm_call_id}
try:
body: Final = await request.body()
body_str: Final = body.decode()
try:
data = ast.literal_eval(body_str)
data = ast.literal_eval(body_str) | data
except Exception:
data = json.loads(body_str)
data = json.loads(body_str) | data
data["adapter_id"] = adapter_id
@ -290,9 +294,7 @@ async def chat_completion_pass_through_endpoint(
response_cost: Final = hidden_params.get("response_cost", None) or ""
### ALERTING ###
asyncio.create_task(
proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success")
)
asyncio.create_task(proxy_logging_obj.update_request_status(litellm_call_id=litellm_call_id, status="success"))
verbose_proxy_logger.debug("final response: %s", response)
@ -313,12 +315,13 @@ async def chat_completion_pass_through_endpoint(
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
)
verbose_proxy_logger.exception("litellm.proxy.proxy_server.completion(): Exception occured - %s", e)
log_llm_api_exception(e, litellm_call_id)
error_msg: Final = f"{e}"
raise ProxyException(
message=getattr(e, "message", error_msg),
type=openai_error_type(e, error_status_code(e, 500)),
param=openai_error_param(e),
headers=litellm_call_id_headers(litellm_call_id),
code=error_status_code(e, 500),
)
@ -609,6 +612,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(
@ -985,6 +989,7 @@ async def pass_through_request(
headers=headers,
forward_headers=forward_headers,
)
upstream_headers: Final = _with_trace_context(headers, parent_span=user_api_key_dict.parent_otel_span)
requested_query_params: dict | None = query_params or dict(request.query_params)
@ -1018,7 +1023,7 @@ async def pass_through_request(
verbose_proxy_logger.debug(
"Pass through endpoint sending request to \nURL %s\nheaders: %s\nbody: %s\n",
url,
headers,
upstream_headers,
_parsed_body,
)
@ -1256,7 +1261,7 @@ async def pass_through_request(
additional_args={
"complete_input_dict": _parsed_body,
"api_base": str(logging_url),
"headers": headers,
"headers": upstream_headers,
},
)
stream = HttpPassThroughEndpointHelpers._update_stream_param_based_on_request_body(
@ -1273,7 +1278,7 @@ async def pass_through_request(
request=request,
async_client=async_client,
url=url,
headers=headers,
headers=upstream_headers,
requested_query_params=requested_query_params,
stream=True,
)
@ -1285,7 +1290,7 @@ async def pass_through_request(
request.method,
url,
params=requested_query_params,
headers=headers,
headers=upstream_headers,
content=state_raw_body,
)
if state_raw_body is not None
@ -1293,7 +1298,7 @@ async def pass_through_request(
request.method,
url,
params=requested_query_params,
headers=headers,
headers=upstream_headers,
json=_parsed_body,
)
)
@ -1370,7 +1375,7 @@ async def pass_through_request(
raw_body_request: Final = async_client.build_request(
request.method,
url,
headers=headers,
headers=upstream_headers,
params=requested_query_params,
content=state_raw_body,
)
@ -1380,7 +1385,7 @@ async def pass_through_request(
request=request,
async_client=async_client,
url=url,
headers=headers,
headers=upstream_headers,
requested_query_params=requested_query_params,
_parsed_body=_parsed_body,
forward_multipart=is_multipart,
@ -2157,6 +2162,17 @@ def _upstream_close_to_relay(task_results: Iterable[object]) -> Close | None:
return upstream_close
_WEBSOCKET_FORWARDED_HEADERS: Final = frozenset(("authorization", "x-api-key", "x-goog-user-project"))
def _with_trace_context(headers: Mapping[str, str], parent_span: object) -> dict[str, str]:
try:
from litellm.integrations.otel.plumbing.context import inject_trace_context
except ImportError:
return dict(headers) # mutable-ok: matches inject_trace_context's carrier return type
return inject_trace_context(headers, parent_span=parent_span)
async def websocket_passthrough_request(
websocket: WebSocket,
target: str,
@ -2199,20 +2215,15 @@ async def websocket_passthrough_request(
await websocket.accept()
verbose_proxy_logger.debug("WebSocket passthrough (%s): WebSocket connection accepted", endpoint)
# Prepare headers for the upstream connection
upstream_headers: Final = custom_headers.copy()
if forward_headers:
# Forward relevant headers from the incoming request
incoming_headers: Final = dict(websocket.headers)
for header_name, header_value in incoming_headers.items():
# Only forward certain headers to avoid conflicts
if header_name.lower() in [
"authorization",
"x-api-key",
"x-goog-user-project",
]:
upstream_headers[header_name] = header_value
forwarded_headers: Final = { # mutable-ok: one-shot upstream header dict, read as a Mapping
**custom_headers,
**{
header_name: header_value
for header_name, header_value in websocket.headers.items()
if forward_headers and header_name.lower() in _WEBSOCKET_FORWARDED_HEADERS
},
}
upstream_headers: Final = _with_trace_context(forwarded_headers, parent_span=user_api_key_dict.parent_otel_span)
# Initialize logging object similar to HTTP passthrough
team_callbacks: Final = _resolve_team_callback_wiring(

162
litellm/proxy/proxy_server.py
View file

@ -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
@ -349,7 +350,10 @@ from litellm.proxy.common_request_processing import (
_is_azure_model_router_request,
_should_return_raw_model_name,
create_response,
log_llm_api_exception,
open_sse_before_first_byte,
request_litellm_call_id,
resolve_litellm_call_id,
ttft_keepalive_interval,
)
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
@ -388,6 +392,11 @@ from litellm.proxy.common_utils.model_listing_utils import (
from litellm.proxy.common_utils.openai_endpoint_utils import (
remove_sensitive_info_from_deployment,
)
from litellm.proxy.common_utils.openai_error_payload import (
headers_with_litellm_call_id,
litellm_call_id_headers,
with_litellm_call_id,
)
from litellm.proxy.common_utils.periodic_reload_schedule import (
MODEL_COST_MAP_RELOAD_PARAM_NAME,
clear_reload_interval,
@ -680,6 +689,9 @@ from litellm.proxy.types_utils.utils import get_instance_fn
from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
router as ui_crud_endpoints_router,
)
from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
sync_ui_settings_to_general_settings,
)
from litellm.proxy.ui_crud_endpoints.user_banner_endpoints import (
router as user_banner_endpoints_router,
)
@ -1744,10 +1756,6 @@ class _SSOConfigRow(Protocol):
sso_settings: MutableMapping[str, object]
class _UISettingsRow(Protocol):
ui_settings: Mapping[str, object] | str | None
class _InvitationLinkRow(Protocol):
user_id: str
expires_at: datetime
@ -6161,6 +6169,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 +6631,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)
@ -7401,7 +7411,12 @@ class ProxyConfig:
Returns what the reconcile saw, captured before the lock is released so a
caller's verdict cannot be corrupted by the next reconcile's own in-flight
window. See ReconcileOutcome.
Also re-reads the UI settings that back runtime flags. That runs before the lock, so a
setting written through one pod reaches the others without waiting on a model reconcile.
"""
await sync_ui_settings_to_general_settings(prisma_client)
async with MODEL_RECONCILE_LOCK:
return await self._add_deployment_locked(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj)
@ -9637,35 +9652,12 @@ class ProxyStartupEvent:
@classmethod
async def _sync_ui_settings_to_general_settings(cls):
"""
Load persisted UI settings from the database and sync runtime flags
into general_settings so they take effect immediately after startup.
"""
try:
import json
from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
_RUNTIME_GENERAL_SETTINGS_FLAGS,
)
if prisma_client is None:
return
db_record: Final[_UISettingsRow | None] = cast( # cast-ok: prisma Json stub is `str`, runtime is a dict
"_UISettingsRow | None",
await UISettingsRepository(prisma_client).table.find_unique(where={"id": "ui_settings"}),
)
if db_record and db_record.ui_settings:
raw: Final = db_record.ui_settings
ui_settings: Final = json.loads(raw) if isinstance(raw, str) else dict(raw)
flags_to_sync: Final = {k: ui_settings[k] for k in _RUNTIME_GENERAL_SETTINGS_FLAGS if k in ui_settings}
if flags_to_sync:
general_settings.update(flags_to_sync)
verbose_proxy_logger.info(
"Synced UI settings to general_settings on startup: %s",
list(flags_to_sync.keys()),
)
except Exception as e:
verbose_proxy_logger.debug("UI settings sync on startup skipped or failed: %s", e)
"""Apply the persisted UI settings to general_settings before this pod serves traffic."""
if prisma_client is None:
return
applied: Final = await sync_ui_settings_to_general_settings(prisma_client)
if applied:
verbose_proxy_logger.info("Synced UI settings to general_settings on startup: %s", list(applied))
@classmethod
async def _load_heuristic_v1_tuning_baselines(
@ -11299,12 +11291,14 @@ async def completion(
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
)
verbose_proxy_logger.exception("litellm.proxy.proxy_server.completion(): Exception occured - %s", e)
litellm_call_id: Final = request_litellm_call_id(data)
log_llm_api_exception(e, litellm_call_id)
error_msg: Final = f"{e}"
raise ProxyException(
message=getattr(e, "message", error_msg),
type=getattr(e, "type", "None"),
param=getattr(e, "param", "None"),
headers=litellm_call_id_headers(litellm_call_id),
openai_code=getattr(e, "code", None),
code=getattr(e, "status_code", 500),
)
@ -11461,11 +11455,12 @@ async def moderations(
```
"""
global proxy_logging_obj
data: dict = {}
litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))
data: dict = {"litellm_call_id": litellm_call_id}
try:
# Use orjson to parse JSON data, orjson speeds up requests significantly
body: Final = await request.body()
data = orjson.loads(body)
data = orjson.loads(body) | data
# Include original request and headers in the data
data = await add_litellm_data_to_request(
@ -11502,9 +11497,7 @@ async def moderations(
response: Final = await llm_call
### ALERTING ###
asyncio.create_task(
proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success")
)
asyncio.create_task(proxy_logging_obj.update_request_status(litellm_call_id=litellm_call_id, status="success"))
### RESPONSE HEADERS ###
hidden_params: Final = getattr(response, "_hidden_params", {}) or {}
@ -11530,14 +11523,15 @@ async def moderations(
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
)
verbose_proxy_logger.exception("litellm.proxy.proxy_server.moderations(): Exception occured - %s", e)
log_llm_api_exception(e, litellm_call_id)
if isinstance(e, ProxyException):
raise
raise with_litellm_call_id(e, litellm_call_id)
if isinstance(e, HTTPException):
raise ProxyException(
message=getattr(e, "message", str(e)),
type=getattr(e, "type", "None"),
param=getattr(e, "param", "None"),
headers=litellm_call_id_headers(litellm_call_id),
code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST),
)
else:
@ -11546,6 +11540,7 @@ async def moderations(
message=getattr(e, "message", error_msg),
type=getattr(e, "type", "None"),
param=getattr(e, "param", "None"),
headers=litellm_call_id_headers(litellm_call_id),
code=getattr(e, "status_code", 500),
)
@ -11583,11 +11578,12 @@ async def audio_speech(
https://platform.openai.com/docs/api-reference/audio/createSpeech
"""
global proxy_logging_obj
data: dict = {}
litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))
data: dict = {"litellm_call_id": litellm_call_id}
try:
# Use orjson to parse JSON data, orjson speeds up requests significantly
body: Final = await request.body()
data = orjson.loads(body)
data = orjson.loads(body) | data
# Include original request and headers in the data
data = await add_litellm_data_to_request(
@ -11620,9 +11616,7 @@ async def audio_speech(
response: Final = await llm_call
### ALERTING ###
asyncio.create_task(
proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success")
)
asyncio.create_task(proxy_logging_obj.update_request_status(litellm_call_id=litellm_call_id, status="success"))
### RESPONSE HEADERS ###
hidden_params: Final = getattr(response, "_hidden_params", {}) or {}
@ -11630,7 +11624,7 @@ async def audio_speech(
cache_key: Final = hidden_params.get("cache_key", None) or ""
api_base: Final = hidden_params.get("api_base", None) or ""
response_cost: Final = hidden_params.get("response_cost", None) or ""
litellm_call_id: Final = hidden_params.get("litellm_call_id", None) or ""
response_call_id: Final = hidden_params.get("litellm_call_id", None) or ""
custom_headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=user_api_key_dict,
@ -11641,7 +11635,7 @@ async def audio_speech(
response_cost=response_cost,
model_region=getattr(user_api_key_dict, "allowed_model_region", ""),
fastest_response_batch_completion=None,
call_id=litellm_call_id,
call_id=response_call_id,
request_data=data,
hidden_params=hidden_params,
)
@ -11677,14 +11671,20 @@ async def audio_speech(
original_exception=e,
request_data=data,
)
verbose_proxy_logger.error("litellm.proxy.proxy_server.audio_speech(): Exception occured - %s", e)
verbose_proxy_logger.debug(traceback.format_exc())
if isinstance(e, (ProxyException, HTTPException)):
raise e
log_llm_api_exception(e, litellm_call_id)
if isinstance(e, ProxyException):
raise with_litellm_call_id(e, litellm_call_id)
if isinstance(e, HTTPException):
raise HTTPException(
status_code=e.status_code,
detail=e.detail,
headers=headers_with_litellm_call_id(e.headers, litellm_call_id),
)
raise ProxyException(
message=getattr(e, "message", f"{e}"),
type=getattr(e, "type", "None"),
param=getattr(e, "param", "None"),
headers=litellm_call_id_headers(litellm_call_id),
openai_code=getattr(e, "code", None),
code=getattr(e, "status_code", 500),
)
@ -11712,11 +11712,12 @@ async def audio_transcriptions(
https://platform.openai.com/docs/api-reference/audio/createTranscription?lang=curl
"""
global proxy_logging_obj
data: dict = {}
litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))
data: dict = {"litellm_call_id": litellm_call_id}
try:
# Use orjson to parse JSON data, orjson speeds up requests significantly
form_data: Final = await get_form_data(request)
data = {key: value for key, value in form_data.items() if key != "file"}
data = {key: value for key, value in form_data.items() if key != "file"} | data
# Include original request and headers in the data
data = await add_litellm_data_to_request(
@ -11783,9 +11784,7 @@ async def audio_transcriptions(
file_object.close() # close the file read in by io library
### ALERTING ###
asyncio.create_task(
proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success")
)
asyncio.create_task(proxy_logging_obj.update_request_status(litellm_call_id=litellm_call_id, status="success"))
### RESPONSE HEADERS ###
hidden_params: Final = getattr(response, "_hidden_params", {}) or {}
@ -11793,7 +11792,7 @@ async def audio_transcriptions(
cache_key: Final = hidden_params.get("cache_key", None) or ""
api_base: Final = hidden_params.get("api_base", None) or ""
response_cost: Final = hidden_params.get("response_cost", None) or ""
litellm_call_id: Final = hidden_params.get("litellm_call_id", None) or ""
response_call_id: Final = hidden_params.get("litellm_call_id", None) or ""
additional_headers: Final[dict] = hidden_params.get("additional_headers", {}) or {}
fastapi_response.headers.update(
@ -11805,7 +11804,7 @@ async def audio_transcriptions(
version=version,
response_cost=response_cost,
model_region=getattr(user_api_key_dict, "allowed_model_region", ""),
call_id=litellm_call_id,
call_id=response_call_id,
request_data=data,
hidden_params=hidden_params,
**additional_headers,
@ -11827,12 +11826,13 @@ async def audio_transcriptions(
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
)
verbose_proxy_logger.exception("litellm.proxy.proxy_server.audio_transcription(): Exception occured - %s", e)
log_llm_api_exception(e, litellm_call_id)
if isinstance(e, HTTPException):
raise ProxyException(
message=getattr(e, "message", str(e.detail)),
type=getattr(e, "type", "None"),
param=getattr(e, "param", "None"),
headers=litellm_call_id_headers(litellm_call_id),
code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST),
)
else:
@ -11841,6 +11841,7 @@ async def audio_transcriptions(
message=getattr(e, "message", error_msg),
type=getattr(e, "type", "None"),
param=getattr(e, "param", "None"),
headers=litellm_call_id_headers(litellm_call_id),
openai_code=getattr(e, "code", None),
code=getattr(e, "status_code", 500),
)
@ -12859,7 +12860,6 @@ from litellm.repositories.table_repositories import (
InvitationLinkRepository,
PromptRepository,
SSOConfigRepository,
UISettingsRepository,
)
from litellm.repositories.team_repository import TeamRepository
from litellm.repositories.user_repository import UserRepository
@ -15537,18 +15537,34 @@ async def model_group_info(
from litellm.proxy.utils import get_available_models_for_user
# Get available models for the user
all_models_str: Final = await get_available_models_for_user(
user_api_key_dict=user_api_key_dict,
llm_router=llm_router,
general_settings=general_settings,
user_model=user_model,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
team_id=None,
include_model_access_groups=False,
only_model_access_groups=False,
return_wildcard_routes=False,
user_api_key_cache=user_api_key_cache,
is_proxy_admin: Final = user_api_key_dict.user_role in (
LitellmUserRoles.PROXY_ADMIN,
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
)
all_models_str: Final = (
get_complete_model_list(
key_models=(),
team_models=(),
proxy_model_list=llm_router.get_model_names(),
user_model=user_model,
infer_model_from_keys=general_settings.get("infer_model_from_keys", False),
return_wildcard_routes=False,
llm_router=llm_router,
)
if is_proxy_admin
else await get_available_models_for_user(
user_api_key_dict=user_api_key_dict,
llm_router=llm_router,
general_settings=general_settings,
user_model=user_model,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
team_id=None,
include_model_access_groups=False,
only_model_access_groups=False,
return_wildcard_routes=False,
user_api_key_cache=user_api_key_cache,
)
)
model_groups: list[ModelGroupInfoProxy] = _get_model_group_info(
llm_router=llm_router, all_models_str=all_models_str, model_group=model_group

23
litellm/proxy/rerank_endpoints/endpoints.py
View file

@ -7,12 +7,16 @@ import orjson
from fastapi import APIRouter, Depends, HTTPException, Request, Response, status
from fastapi.responses import ORJSONResponse
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.proxy.common_request_processing import (
ProxyBaseLLMRequestProcessing,
log_llm_api_exception,
resolve_litellm_call_id,
)
from litellm.proxy.common_utils.openai_error_payload import (
error_status_code,
litellm_call_id_headers,
openai_error_param,
openai_error_type,
)
@ -54,10 +58,11 @@ async def rerank(
version,
)
data = {}
litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))
data = {"litellm_call_id": litellm_call_id}
try:
body: Final = await request.body()
data = orjson.loads(body)
data = orjson.loads(body) | data
# Include original request and headers in the data
data = await add_litellm_data_to_request(
@ -82,9 +87,7 @@ async def rerank(
response: Final = await llm_call
### ALERTING ###
asyncio.create_task(
proxy_logging_obj.update_request_status(litellm_call_id=data.get("litellm_call_id", ""), status="success")
)
asyncio.create_task(proxy_logging_obj.update_request_status(litellm_call_id=litellm_call_id, status="success"))
### RESPONSE HEADERS ###
hidden_params: Final = getattr(response, "_hidden_params", {}) or {}
@ -95,7 +98,7 @@ async def rerank(
fastapi_response.headers.update(
ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=user_api_key_dict,
call_id=hidden_params.get("litellm_call_id", None) or data.get("litellm_call_id", None),
call_id=hidden_params.get("litellm_call_id", None) or litellm_call_id,
model_id=model_id,
cache_key=cache_key,
api_base=api_base,
@ -113,12 +116,13 @@ async def rerank(
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
)
verbose_proxy_logger.error("litellm.proxy.proxy_server.rerank(): Exception occured - %s", e)
log_llm_api_exception(e, litellm_call_id)
if isinstance(e, HTTPException):
raise ProxyException(
message=getattr(e, "message", str(e)),
type=openai_error_type(e, error_status_code(e, status.HTTP_400_BAD_REQUEST)),
param=openai_error_param(e),
headers=litellm_call_id_headers(litellm_call_id),
code=error_status_code(e, status.HTTP_400_BAD_REQUEST),
)
else:
@ -127,5 +131,6 @@ async def rerank(
message=getattr(e, "message", error_msg),
type=openai_error_type(e, error_status_code(e, 500)),
param=openai_error_param(e),
headers=litellm_call_id_headers(litellm_call_id),
code=error_status_code(e, 500),
)

2
litellm/proxy/schema.prisma
View file

@ -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("{}")

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

@ -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,

89
litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py
View file

@ -14,7 +14,7 @@ from typing import (
from urllib.parse import urlparse
from fastapi import APIRouter, Body, Depends, File, HTTPException, UploadFile
from pydantic import ConfigDict, JsonValue, ValidationError, create_model
from pydantic import ConfigDict, JsonValue, TypeAdapter, ValidationError, create_model
from pydantic.fields import FieldInfo
from typing_extensions import NotRequired, ReadOnly, TypedDict
@ -29,6 +29,10 @@ from litellm.proxy.config_resolvers.sso import (
SSO_SECRET_FIELDS,
resolve_sso_config,
)
from litellm.proxy.management_endpoints.team_admin_field_permissions import (
SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS,
TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING,
)
from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled
from litellm.proxy.utils import invalidate_config_param
from litellm.repositories.config_repository import ConfigRepository
@ -212,6 +216,9 @@ class UIThemeSettingsResponse(SettingsResponse):
"""Response model for UI theme settings"""
_TEAM_ADMIN_FIELD_ENUM: Final = tuple(sorted(SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS))
class UISettings(BaseModel):
"""Configuration for UI-specific flags"""
@ -304,6 +311,18 @@ class UISettings(BaseModel):
description="If true, shows the Chat page in the UI sidebar, letting users chat with an LLM and connect their own MCP server credentials via OAuth.",
)
team_admin_editable_team_fields: Sequence[str] = Field(
default=(),
description=(
"Team settings fields a team admin may change on the teams they administer. "
"Empty means team admins cannot edit team settings at all. "
"Proxy admins and org admins are not affected."
),
json_schema_extra={ # mutable-ok: pydantic only merges json_schema_extra when it is a plain dict
"items": {"type": "string", "enum": [*_TEAM_ADMIN_FIELD_ENUM]}, # mutable-ok: nested in the dict above
},
)
class UISettingsResponse(SettingsResponse):
"""Response model for UI settings"""
@ -326,6 +345,7 @@ ALLOWED_UI_SETTINGS_FIELDS: Final = {
"disable_custom_api_keys",
"disable_key_generate_for_org_admin",
"enable_chat_ui",
TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING,
}
ENABLE_PTU_COST_ATTRIBUTION_UI_SETTING: Final = "enable_ptu_cost_attribution"
@ -360,6 +380,7 @@ _RUNTIME_GENERAL_SETTINGS_FLAGS: Final = [
"disable_vector_stores_for_internal_users",
"allow_vector_stores_for_team_admins",
"disable_key_generate_for_org_admin",
TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING,
]
# Extension point: packages outside OSS (e.g. litellm_enterprise) can
@ -1457,6 +1478,42 @@ async def get_ui_settings_cached() -> dict[str, JsonValue]:
return ui_settings
_UI_SETTINGS_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
def apply_runtime_general_settings_flags(ui_settings: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]:
"""Copy the UI settings that gate runtime behavior into ``general_settings``. Returns what was applied."""
from litellm.proxy.proxy_server import general_settings
flags: Final = {k: ui_settings[k] for k in _RUNTIME_GENERAL_SETTINGS_FLAGS if k in ui_settings}
if flags:
general_settings.update(flags)
return MappingProxyType(flags)
async def sync_ui_settings_to_general_settings(prisma_client: object) -> Mapping[str, JsonValue]:
"""Re-read the persisted UI settings and apply the runtime flags to ``general_settings``.
Runs on startup and on every periodic config reload: the PATCH handler only updates the pod
that served it, so every other pod needs its own read to pick up a change without a restart.
Never raises. A read that fails leaves this pod on the flags it already had.
"""
try:
db_record: Final = await _ui_settings_db(UISettingsRepository(prisma_client)).find_unique(
where={"id": "ui_settings"}
)
stored: Final = (db_record.ui_settings if db_record else None) or "{}"
parsed: Final = (
_UI_SETTINGS_OBJECT.validate_json(stored)
if isinstance(stored, str)
else _UI_SETTINGS_OBJECT.validate_python(stored)
)
except Exception as e:
verbose_proxy_logger.warning("Could not refresh UI settings from the database: %s", e)
return MappingProxyType({})
return apply_runtime_general_settings_flags(parsed)
@router.get(
"/get/ui_settings",
tags=["UI Settings"],
@ -1485,13 +1542,7 @@ async def get_ui_settings():
# Sanitize any unexpected keys from persisted config before returning
ui_settings: Final = {k: v for k, v in parsed.items() if k in ALLOWED_UI_SETTINGS_FIELDS}
# Sync runtime flags into general_settings so the proxy picks them up
# at runtime (covers server restart scenarios).
_flags_to_sync: Final = {k: ui_settings[k] for k in _RUNTIME_GENERAL_SETTINGS_FLAGS if k in ui_settings}
if _flags_to_sync:
from litellm.proxy.proxy_server import general_settings
general_settings.update(_flags_to_sync)
apply_runtime_general_settings_flags(ui_settings)
# Refresh DualCache so other code paths (e.g. /user/filter/ui) see fresh values
from litellm.proxy.proxy_server import user_api_key_cache
@ -1571,6 +1622,20 @@ async def update_ui_settings(
except ValidationError as e:
raise HTTPException(status_code=422, detail=e.errors())
unsupported_team_fields: Final = sorted(
frozenset(settings.team_admin_editable_team_fields) - SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS
)
if unsupported_team_fields:
raise HTTPException(
status_code=400,
detail={ # mutable-ok: HTTPException detail must be a plain dict for FastAPI JSON serialization
"error": (
f"{TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING} does not support {unsupported_team_fields}. "
f"Supported fields: {sorted(SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS)}."
)
},
)
# Only include fields the caller actually sent (not Pydantic defaults).
settings_dict: Final[Mapping[str, JsonValue]] = settings.model_dump(exclude_unset=True)
@ -1616,13 +1681,7 @@ async def update_ui_settings(
},
)
# Sync runtime flags to general_settings so the proxy picks them up
# at runtime (general_settings is checked in pre-call utils).
_flags_to_sync: Final = {k: ui_settings[k] for k in _RUNTIME_GENERAL_SETTINGS_FLAGS if k in ui_settings}
if _flags_to_sync:
from litellm.proxy.proxy_server import general_settings
general_settings.update(_flags_to_sync)
apply_runtime_general_settings_flags(ui_settings)
# Invalidate + set DualCache so subsequent reads see the new values immediately
from litellm.proxy.proxy_server import user_api_key_cache

21
litellm/proxy/utils.py
View file

@ -38,7 +38,11 @@ from litellm.proxy._types import (
SpendLogsMetadata,
SpendLogsPayload,
)
from litellm.proxy.common_utils.openai_error_payload import openai_error_param
from litellm.proxy.common_utils.openai_error_payload import (
litellm_call_id_headers,
openai_error_param,
with_litellm_call_id,
)
from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.proxy.model_listing import ModelInfoResponse
@ -164,7 +168,6 @@ from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrai
)
from litellm.proxy.hooks import PROXY_HOOKS, get_proxy_hook
from litellm.proxy.hooks.cache_control_check import _PROXY_CacheControlCheck
from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter
from litellm.proxy.hooks.parallel_request_limiter import (
_PROXY_MaxParallelRequestsHandler,
)
@ -982,7 +985,6 @@ class ProxyLogging:
dual_cache=DualCache(default_in_memory_ttl=1) # ping redis cache every 1s
)
self.max_parallel_request_limiter = _PROXY_MaxParallelRequestsHandler(self.internal_usage_cache)
self.max_budget_limiter = _PROXY_MaxBudgetLimiter()
self.cache_control_check = _PROXY_CacheControlCheck()
self.alerting: list[str] | None = None
self.alerting_threshold: float = 300 # default to 5 min. threshold
@ -3052,7 +3054,7 @@ class ProxyLogging:
if litellm_logging_obj is None:
from litellm._uuid import uuid
request_data["litellm_call_id"] = str(uuid.uuid4())
request_data.setdefault("litellm_call_id", str(uuid.uuid4()))
user_api_key_logged_metadata: Final = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(
user_api_key_dict=user_api_key_dict
)
@ -3580,7 +3582,7 @@ class ProxyLogging:
caps: Final = ProxyLogging._callback_capabilities()
post_call_pipelines: Final = _streamable_post_call_pipelines(request_data, user_api_key_dict)
# Fast path: no real overrides. Internal proxy CustomLogger callbacks
# (e.g. _PROXY_MaxBudgetLimiter, ManagedFiles) inherit the default
# (e.g. _PROXY_CacheControlCheck, ManagedFiles) inherit the default
# ``async for chunk: yield chunk`` body, so wrapping the iterator
# through each of them adds N pass-through trampolines per chunk for
# zero behavior change. Skip the chain entirely and stream through.
@ -4340,6 +4342,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 +4782,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,
@ -7659,7 +7663,7 @@ def _recreate_writer_on_read_only_transaction(prisma_client: "PrismaClient | Non
asyncio.create_task(prisma_client.recreate_read_only_writer(reason="postgres_read_only_transaction"))
def handle_exception_on_proxy(e: Exception) -> ProxyException:
def handle_exception_on_proxy(e: Exception, litellm_call_id: str | None = None) -> ProxyException:
"""
Returns an Exception as ProxyException, this ensures all exceptions are OpenAI API compatible
"""
@ -7671,20 +7675,23 @@ def handle_exception_on_proxy(e: Exception) -> ProxyException:
_recreate_writer_on_read_only_transaction(prisma_client)
headers: Final = litellm_call_id_headers(litellm_call_id)
if isinstance(e, HTTPException):
return ProxyException(
message=getattr(e, "detail", f"error({e})"),
type=ProxyErrorTypes.internal_server_error,
param=openai_error_param(e),
headers=headers,
code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR),
)
elif isinstance(e, ProxyException):
return e
return with_litellm_call_id(e, litellm_call_id)
_status_code: Final = getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR)
return ProxyException(
message=str(e),
type=ProxyErrorTypes.internal_server_error,
param=openai_error_param(e),
headers=headers,
code=_status_code,
)

34
litellm/responses/main.py
View file

@ -482,18 +482,27 @@ class _AsyncPromptManagementOutcome:
def _resolve_responses_api_provider_config(
model: str, custom_llm_provider: str, model_info: object
model: str, custom_llm_provider: str, model_info: object, api_base: str | None
) -> BaseResponsesAPIConfig | None:
provider_config: Final = ProviderConfigManager.get_provider_responses_api_config(
model=model, provider=custom_llm_provider
model=model, provider=custom_llm_provider, api_base=api_base
)
if provider_config is not None or not _deployment_passes_through_responses(model_info):
return provider_config
return OpenAILikeResponsesConfig()
def _api_base_kwarg(kwargs: Mapping[str, object]) -> str | None:
api_base: Final = kwargs.get("api_base")
return api_base if isinstance(api_base, str) else None
def _will_bridge_to_chat_completions(
model: str, custom_llm_provider: str | None, use_chat_completions_api: bool, model_info: object
model: str,
custom_llm_provider: str | None,
use_chat_completions_api: bool,
model_info: object,
api_base: str | None,
) -> bool:
"""``_bridges_to_chat_completions`` for callers running before the provider config is resolved.
@ -507,7 +516,7 @@ def _will_bridge_to_chat_completions(
if custom_llm_provider is None:
return True
return _bridges_to_chat_completions(
_resolve_responses_api_provider_config(normalized_model[0], custom_llm_provider, model_info),
_resolve_responses_api_provider_config(normalized_model[0], custom_llm_provider, model_info, api_base),
use_chat_completions_api or normalized_model[1],
)
@ -618,6 +627,7 @@ async def aresponses(
custom_llm_provider,
bool(kwargs.get("use_chat_completions_api")),
kwargs.get("model_info"),
_api_base_kwarg(kwargs),
),
):
(
@ -783,7 +793,11 @@ def _apply_prompt_management_to_responses_call(
with _prompt_management_sees_a_provisional_message_list(
kwargs,
bridged=_will_bridge_to_chat_completions(
model, custom_llm_provider, use_chat_completions_api, kwargs.get("model_info")
model,
custom_llm_provider,
use_chat_completions_api,
kwargs.get("model_info"),
_api_base_kwarg(kwargs),
),
):
(
@ -1237,7 +1251,7 @@ def responses(
responses_api_provider_config = None
else:
responses_api_provider_config = _resolve_responses_api_provider_config(
model, custom_llm_provider, deployment_model_info
model, custom_llm_provider, deployment_model_info, litellm_params.api_base
)
if (
@ -1496,6 +1510,7 @@ def delete_responses(
ProviderConfigManager.get_provider_responses_api_config(
model=None,
provider=custom_llm_provider,
api_base=litellm_params.api_base,
)
)
@ -1667,6 +1682,7 @@ def get_responses(
ProviderConfigManager.get_provider_responses_api_config(
model=None,
provider=custom_llm_provider,
api_base=litellm_params.api_base,
)
)
@ -1811,6 +1827,7 @@ def list_input_items(
ProviderConfigManager.get_provider_responses_api_config(
model=None,
provider=custom_llm_provider,
api_base=litellm_params.api_base,
)
)
@ -1960,6 +1977,7 @@ def cancel_responses(
ProviderConfigManager.get_provider_responses_api_config(
model=None,
provider=custom_llm_provider,
api_base=litellm_params.api_base,
)
)
@ -2132,6 +2150,7 @@ def compact_responses(
ProviderConfigManager.get_provider_responses_api_config(
model=model,
provider=custom_llm_provider,
api_base=litellm_params.api_base,
)
)
@ -2270,14 +2289,15 @@ async def _aresponses_websocket(
custom_llm_provider=_custom_llm_provider,
)
resolved_api_base: Final = dynamic_api_base or litellm_params.api_base or litellm.api_base or None
responses_api_provider_config: BaseResponsesAPIConfig | None = None
if _custom_llm_provider is not None:
responses_api_provider_config = ProviderConfigManager.get_provider_responses_api_config(
model=resolved_model,
provider=litellm.LlmProviders(_custom_llm_provider),
api_base=resolved_api_base,
)
resolved_api_base: Final = dynamic_api_base or litellm_params.api_base or litellm.api_base or None
resolved_api_key: Final = (
dynamic_api_key
or litellm_params.api_key

243
litellm/router.py
View file

@ -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:

9
litellm/router_utils/add_retry_fallback_headers.py
View file

@ -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]:

21
litellm/router_utils/fallback_event_handlers.py
View file

@ -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:

23
litellm/router_utils/router_callbacks/track_deployment_metrics.py
View file

@ -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,

13
litellm/rust_bridge/_native.pyi
View file

@ -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",

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