From caf37c8b6f21bfc84752baaafc568deeaf3b1996 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 20:36:07 -0700 Subject: [PATCH 01/23] refactor(rust): share settings lookup and layer merge through core-utils Settings sources beyond HTTP (media fetch, Azure Document Intelligence, Vertex, timeouts) need the same env lookup and precedence merge, so move them out of litellm-http into core_utils::settings. Lookup readers name the Python idiom they mirror: get keeps a present empty value like os.getenv(X, fallback), truthy drops it like an `or` chain, enabled only switches on for "true". SSL_CERT_FILE now reads through truthy, matching Python's `if ssl_cert_file and ...` check. Co-Authored-By: Claude Opus 5 --- litellm-rust/Cargo.lock | 2 + litellm-rust/crates/core-utils/src/lib.rs | 1 + .../crates/core-utils/src/settings.rs | 144 ++++++++++++++++++ litellm-rust/crates/http/Cargo.toml | 1 + litellm-rust/crates/http/src/settings.rs | 50 +++--- litellm-rust/crates/python-bridge/Cargo.toml | 1 + litellm-rust/crates/python-bridge/src/http.rs | 5 +- 7 files changed, 174 insertions(+), 30 deletions(-) create mode 100644 litellm-rust/crates/core-utils/src/settings.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index d4b32659ba1..83cdbc6a782 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2137,6 +2137,7 @@ version = "0.1.0" dependencies = [ "http 1.4.2", "hyper-util", + "litellm-core-utils", "reqwest 0.12.28", "rstest", "rustls 0.23.42", @@ -2187,6 +2188,7 @@ dependencies = [ "litellm-auth-gcp", "litellm-callbacks-legacy", "litellm-core", + "litellm-core-utils", "litellm-host-python", "litellm-http", "litellm-llms", diff --git a/litellm-rust/crates/core-utils/src/lib.rs b/litellm-rust/crates/core-utils/src/lib.rs index fcb232d8980..ceb0e9eb3f2 100644 --- a/litellm-rust/crates/core-utils/src/lib.rs +++ b/litellm-rust/crates/core-utils/src/lib.rs @@ -6,4 +6,5 @@ pub mod params; pub mod prompt_templates; pub mod secret_redaction; pub mod serde_compat; +pub mod settings; pub mod url_utils; diff --git a/litellm-rust/crates/core-utils/src/settings.rs b/litellm-rust/crates/core-utils/src/settings.rs new file mode 100644 index 00000000000..59c76ce3015 --- /dev/null +++ b/litellm-rust/crates/core-utils/src/settings.rs @@ -0,0 +1,144 @@ +use std::str::FromStr; + +pub trait Lookup { + fn get(&self, name: &str) -> Option; + + fn truthy(&self, name: &str) -> Option { + self.get(name).filter(|value| !value.is_empty()) + } + + fn enabled(&self, name: &str) -> Option { + self.get(name) + .is_some_and(|value| value.trim().eq_ignore_ascii_case("true")) + .then_some(true) + } + + fn parsed(&self, name: &str) -> Option + where + Self: Sized, + { + self.get(name).and_then(|value| value.trim().parse().ok()) + } +} + +impl Option> Lookup for F { + fn get(&self, name: &str) -> Option { + self(name) + } +} + +pub struct ProcessEnvironment; + +impl Lookup for ProcessEnvironment { + fn get(&self, name: &str) -> Option { + std::env::var(name).ok() + } +} + +pub trait Layer: Default { + fn or(self, lower: Self) -> Self; +} + +pub fn merge(highest_precedence_first: impl IntoIterator) -> L { + highest_precedence_first + .into_iter() + .reduce(L::or) + .unwrap_or_default() +} + +#[cfg(test)] +mod tests { + use super::*; + + fn env_of(values: &'static [(&'static str, &'static str)]) -> impl Fn(&str) -> Option { + move |name| { + values + .iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| value.to_string()) + } + } + + #[test] + fn get_keeps_a_present_empty_value_like_os_getenv_with_a_fallback() { + let env = env_of(&[("EMPTY", "")]); + assert_eq!(env.get("EMPTY"), Some(String::new())); + assert_eq!(env.get("ABSENT"), None); + } + + #[test] + fn truthy_drops_an_empty_value_like_a_python_or_chain() { + let env = env_of(&[("EMPTY", ""), ("SET", "value")]); + assert_eq!(env.truthy("EMPTY"), None); + assert_eq!(env.truthy("SET").as_deref(), Some("value")); + } + + #[test] + fn enabled_only_switches_on_for_true_and_never_forces_off() { + let env = env_of(&[ + ("LOWER", "true"), + ("PADDED", " True "), + ("OFF", "false"), + ("ONE", "1"), + ]); + assert_eq!(env.enabled("LOWER"), Some(true)); + assert_eq!(env.enabled("PADDED"), Some(true)); + assert_eq!(env.enabled("OFF"), None); + assert_eq!(env.enabled("ONE"), None); + assert_eq!(env.enabled("ABSENT"), None); + } + + #[test] + fn parsed_trims_and_skips_values_that_do_not_parse() { + let env = env_of(&[("PADDED", " 45 "), ("WORD", "soon"), ("FRACTION", "0.5")]); + assert_eq!(env.parsed::("PADDED"), Some(45)); + assert_eq!(env.parsed::("WORD"), None); + assert_eq!(env.parsed::("FRACTION"), Some(0.5)); + assert_eq!(env.parsed::("ABSENT"), None); + } + + #[derive(Debug, Default, PartialEq)] + struct Pair { + first: Option, + second: Option, + } + + impl Layer for Pair { + fn or(self, lower: Self) -> Self { + Self { + first: self.first.or(lower.first), + second: self.second.or(lower.second), + } + } + } + + #[test] + fn merge_takes_each_field_from_the_highest_layer_that_sets_it() { + let merged = merge([ + Pair { + first: Some(1), + second: None, + }, + Pair { + first: Some(2), + second: Some(2), + }, + Pair { + first: Some(3), + second: Some(3), + }, + ]); + assert_eq!( + merged, + Pair { + first: Some(1), + second: Some(2), + } + ); + } + + #[test] + fn merging_no_layers_yields_the_empty_layer() { + assert_eq!(merge(Vec::::new()), Pair::default()); + } +} diff --git a/litellm-rust/crates/http/Cargo.toml b/litellm-rust/crates/http/Cargo.toml index 0ac09a9d155..b0dc7693840 100644 --- a/litellm-rust/crates/http/Cargo.toml +++ b/litellm-rust/crates/http/Cargo.toml @@ -7,6 +7,7 @@ repository.workspace = true [dependencies] http.workspace = true +litellm-core-utils.workspace = true hyper-util.workspace = true reqwest.workspace = true rustls.workspace = true diff --git a/litellm-rust/crates/http/src/settings.rs b/litellm-rust/crates/http/src/settings.rs index 8ac7ef92568..43c7f6223d2 100644 --- a/litellm-rust/crates/http/src/settings.rs +++ b/litellm-rust/crates/http/src/settings.rs @@ -3,6 +3,8 @@ use std::{ time::Duration, }; +use litellm_core_utils::settings::{Layer, Lookup, merge}; + #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub enum SslVerify { Enabled, @@ -45,38 +47,35 @@ pub struct HttpSettingsLayer { } impl HttpSettingsLayer { - pub fn from_environment(env: &(dyn Fn(&str) -> Option + Sync)) -> Self { - let enabled = |name: &str| { - env(name) - .is_some_and(|value| value.trim().eq_ignore_ascii_case("true")) - .then_some(true) - }; - let number = |name: &str| env(name).and_then(|value| value.trim().parse::().ok()); + pub fn from_environment(env: &impl Lookup) -> Self { let seconds = |name: &str, default: u32| { - Duration::from_secs(u64::from(number(name).unwrap_or(default))) + Duration::from_secs(u64::from(env.parsed::(name).unwrap_or(default))) }; Self { - ssl_verify: env("SSL_VERIFY").map(|value| SslVerify::parse(&value)), - ssl_cert_file: env("SSL_CERT_FILE").map(PathBuf::from), - ssl_certificate: env("SSL_CERTIFICATE").map(PathBuf::from), - ssl_security_level: env("SSL_SECURITY_LEVEL"), - ssl_ecdh_curve: env("SSL_ECDH_CURVE"), + ssl_verify: env.get("SSL_VERIFY").map(|value| SslVerify::parse(&value)), + ssl_cert_file: env.truthy("SSL_CERT_FILE").map(PathBuf::from), + ssl_certificate: env.get("SSL_CERTIFICATE").map(PathBuf::from), + ssl_security_level: env.get("SSL_SECURITY_LEVEL"), + ssl_ecdh_curve: env.get("SSL_ECDH_CURVE"), force_ipv4: None, - http2: enabled("LITELLM_HTTP2"), - aiohttp_trust_env: enabled("AIOHTTP_TRUST_ENV"), - disable_aiohttp_trust_env: enabled("DISABLE_AIOHTTP_TRUST_ENV"), - disable_aiohttp_transport: enabled("DISABLE_AIOHTTP_TRANSPORT"), - user_agent: env("LITELLM_USER_AGENT"), - tcp_keepalive: enabled("AIOHTTP_SO_KEEPALIVE").map(|_| TcpKeepalive { + http2: env.enabled("LITELLM_HTTP2"), + aiohttp_trust_env: env.enabled("AIOHTTP_TRUST_ENV"), + disable_aiohttp_trust_env: env.enabled("DISABLE_AIOHTTP_TRUST_ENV"), + disable_aiohttp_transport: env.enabled("DISABLE_AIOHTTP_TRANSPORT"), + user_agent: env.get("LITELLM_USER_AGENT"), + tcp_keepalive: env.enabled("AIOHTTP_SO_KEEPALIVE").map(|_| TcpKeepalive { idle: seconds("AIOHTTP_TCP_KEEPIDLE", 60), interval: seconds("AIOHTTP_TCP_KEEPINTVL", 30), - retries: number("AIOHTTP_TCP_KEEPCNT").unwrap_or(5), + retries: env.parsed("AIOHTTP_TCP_KEEPCNT").unwrap_or(5), }), - pool_idle_timeout: number("AIOHTTP_KEEPALIVE_TIMEOUT") + pool_idle_timeout: env + .parsed::("AIOHTTP_KEEPALIVE_TIMEOUT") .map(|timeout| Duration::from_secs(u64::from(timeout))), } } +} +impl Layer for HttpSettingsLayer { fn or(self, lower: Self) -> Self { Self { ssl_verify: self.ssl_verify.or(lower.ssl_verify), @@ -139,10 +138,7 @@ impl HttpSettings { pub fn from_layers( highest_precedence_first: impl IntoIterator, ) -> Self { - let merged = highest_precedence_first - .into_iter() - .reduce(HttpSettingsLayer::or) - .unwrap_or_default(); + let merged = merge(highest_precedence_first); let defaults = Self::default(); let http2 = merged.http2.unwrap_or(defaults.http2); Self { @@ -190,9 +186,7 @@ mod tests { None } - fn env_of( - values: &'static [(&'static str, &'static str)], - ) -> impl Fn(&str) -> Option + Sync { + fn env_of(values: &'static [(&'static str, &'static str)]) -> impl Fn(&str) -> Option { move |name| { values .iter() diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index c66701548d1..8d31855f2fa 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -20,6 +20,7 @@ bytes.workspace = true litellm-auth.workspace = true litellm-callbacks-legacy.workspace = true litellm-core.workspace = true +litellm-core-utils.workspace = true litellm-auth-gcp.workspace = true litellm-http.workspace = true litellm-llms.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/http.rs b/litellm-rust/crates/python-bridge/src/http.rs index d174dccaa56..1fc3e4a60f1 100644 --- a/litellm-rust/crates/python-bridge/src/http.rs +++ b/litellm-rust/crates/python-bridge/src/http.rs @@ -4,6 +4,7 @@ use std::{ sync::{Arc, LazyLock, Mutex, PoisonError}, }; +use litellm_core_utils::settings::ProcessEnvironment; use litellm_http::{ HttpClientConfig, HttpClientPool, HttpSettings, HttpSettingsLayer, Resolution, SslVerify, Unsupported, @@ -29,7 +30,7 @@ pub(crate) fn call_config( ) -> PyResult { let settings = HttpSettings::from_layers([ for_call(call_ssl_verify(kwargs)?, asynchronous), - HttpSettingsLayer::from_environment(&|name| std::env::var(name).ok()), + HttpSettingsLayer::from_environment(&ProcessEnvironment), configured(&PythonSettings::Http.read(py)?)?, ]) .without_missing_files(&|path: &Path| path.exists()); @@ -232,7 +233,7 @@ user_agent='litellm/9.9.9', Python::initialize(); Python::attach(|py| { let settings = HttpSettings::from_layers([ - HttpSettingsLayer::from_environment(&|name| { + HttpSettingsLayer::from_environment(&|name: &str| { (name == "LITELLM_USER_AGENT").then(|| "operator/1".to_string()) }), configured(&python_settings(py, "")).unwrap(), From a41885e48e57aed9e70afb138719a16db1860765 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 20:41:07 -0700 Subject: [PATCH 02/23] refactor(rust): read proxy env vars through the settings lookup reqwest and hyper each read HTTP(S)_PROXY, ALL_PROXY and NO_PROXY from the process on their own, so tests could not inject them and the pooled client key ignored proxy changes. EnvironmentProxies now reads them through Lookup with the same precedence hyper used, the resolved config carries them (empty when the transport does not trust the env), and both the provider clients and the media fetcher build from that one value. Co-Authored-By: Claude Opus 5 --- litellm-rust/crates/http/src/config.rs | 41 +++++++-- litellm-rust/crates/http/src/pool.rs | 62 ++++++++++++- litellm-rust/crates/http/src/proxy.rs | 91 ++++++++++++++++++- litellm-rust/crates/http/src/settings.rs | 9 ++ .../crates/llms/src/custom_httpx/media.rs | 15 +-- 5 files changed, 191 insertions(+), 27 deletions(-) diff --git a/litellm-rust/crates/http/src/config.rs b/litellm-rust/crates/http/src/config.rs index 10f28b44eec..bf8ecef85a8 100644 --- a/litellm-rust/crates/http/src/config.rs +++ b/litellm-rust/crates/http/src/config.rs @@ -6,6 +6,7 @@ use std::{ use crate::{ error::Error, + proxy::EnvironmentProxies, settings::{HttpSettings, SslVerify, TcpKeepalive}, tls::{CipherSelection, KeyExchangeGroup, Tls12CipherSuite, Unsupported}, }; @@ -26,7 +27,7 @@ pub struct HttpClientConfig { pub force_ipv4: bool, pub http2: bool, pub user_agent: Option, - pub trust_proxy_env: bool, + pub proxies: EnvironmentProxies, pub connect_timeout: Duration, pub tcp_keepalive: Option, pub pool_idle_timeout: Duration, @@ -72,7 +73,11 @@ impl From<&HttpSettings> for Resolution { force_ipv4: settings.force_ipv4, http2: settings.http2, user_agent: settings.user_agent.clone(), - trust_proxy_env: settings.trust_proxy_env, + proxies: if settings.trust_proxy_env { + settings.proxies.clone() + } else { + EnvironmentProxies::default() + }, connect_timeout: settings.connect_timeout, tcp_keepalive: settings.tcp_keepalive, pool_idle_timeout: settings.pool_idle_timeout, @@ -111,11 +116,11 @@ impl TryFrom<&HttpClientConfig> for reqwest::ClientBuilder { Some(agent) => with_protocol.user_agent(agent), None => with_protocol, }; - Ok(if config.trust_proxy_env { - with_agent - } else { - with_agent.no_proxy() - }) + Ok(config + .proxies + .reqwest_proxies() + .into_iter() + .fold(with_agent.no_proxy(), reqwest::ClientBuilder::proxy)) } } @@ -227,6 +232,25 @@ mod tests { ); } + fn proxies() -> EnvironmentProxies { + EnvironmentProxies::from_environment(&|name: &str| { + (name == "HTTPS_PROXY").then(|| "http://proxy.corp:3128".to_string()) + }) + } + + #[test] + fn proxies_are_dropped_when_the_transport_does_not_trust_the_environment() { + let settings = HttpSettings { + trust_proxy_env: false, + proxies: proxies(), + ..HttpSettings::default() + }; + assert_eq!( + Resolution::from(&settings).config.proxies, + EnvironmentProxies::default() + ); + } + #[test] fn connection_settings_carry_over_unchanged() { let keepalive = TcpKeepalive { @@ -240,6 +264,7 @@ mod tests { http2: true, user_agent: Some("litellm/1.0".into()), trust_proxy_env: true, + proxies: proxies(), connect_timeout: Duration::from_secs(7), tcp_keepalive: Some(keepalive), pool_idle_timeout: Duration::from_secs(45), @@ -256,7 +281,7 @@ mod tests { force_ipv4: true, http2: true, user_agent: Some("litellm/1.0".into()), - trust_proxy_env: true, + proxies: proxies(), connect_timeout: Duration::from_secs(7), tcp_keepalive: Some(keepalive), pool_idle_timeout: Duration::from_secs(45), diff --git a/litellm-rust/crates/http/src/pool.rs b/litellm-rust/crates/http/src/pool.rs index 330d6de29e8..ee47e5dc52a 100644 --- a/litellm-rust/crates/http/src/pool.rs +++ b/litellm-rust/crates/http/src/pool.rs @@ -6,7 +6,7 @@ use std::{ use reqwest::dns::Resolve; -use crate::{config::HttpClientConfig, error::Error}; +use crate::{config::HttpClientConfig, error::Error, proxy::EnvironmentProxies}; #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] pub enum ClientVariant { @@ -52,7 +52,7 @@ impl HttpClientPool { let effective = match variant { ClientVariant::Media => HttpClientConfig { client_certificate: None, - trust_proxy_env: false, + proxies: EnvironmentProxies::default(), ..config.clone() }, ClientVariant::UnpinnedMedia => HttpClientConfig { @@ -138,6 +138,13 @@ mod tests { } } + fn proxied_through(proxy: &str) -> EnvironmentProxies { + let proxy = proxy.to_owned(); + EnvironmentProxies::from_environment(&move |name: &str| { + (name == "HTTP_PROXY").then(|| proxy.clone()) + }) + } + async fn serve( status_line: &'static str, ) -> (SocketAddr, Arc, Arc>>) { @@ -202,6 +209,50 @@ mod tests { assert_eq!(connections.load(Ordering::SeqCst), 3); } + #[tokio::test] + async fn provider_clients_route_through_the_resolved_proxy_not_the_process_environment() { + let (proxy, connections, requests) = serve("HTTP/1.1 204 No Content").await; + let config = HttpClientConfig { + proxies: proxied_through(&format!("http://user:secret@{proxy}")), + ..config("a") + }; + let response = get( + &pool(), + &config, + ClientVariant::Provider, + "http://upstream.invalid/v1/ocr", + ) + .await; + assert_eq!(response.status(), 204); + assert_eq!(connections.load(Ordering::SeqCst), 1); + let request = requests.lock().unwrap().concat(); + assert!(request.starts_with("GET http://upstream.invalid/v1/ocr HTTP/1.1")); + assert!(request.contains("proxy-authorization: Basic dXNlcjpzZWNyZXQ=")); + } + + #[tokio::test] + async fn no_proxy_hosts_bypass_the_resolved_proxy() { + let (upstream, _, _) = serve("HTTP/1.1 204 No Content").await; + let (proxy, proxy_connections, _) = serve("HTTP/1.1 502 Bad Gateway").await; + let config = HttpClientConfig { + proxies: EnvironmentProxies::from_environment(&move |name: &str| match name { + "HTTP_PROXY" => Some(format!("http://{proxy}")), + "NO_PROXY" => Some("127.0.0.1".into()), + _ => None, + }), + ..config("a") + }; + let response = get( + &pool(), + &config, + ClientVariant::Provider, + &format!("http://{upstream}/v1/ocr"), + ) + .await; + assert_eq!(response.status(), 204); + assert_eq!(proxy_connections.load(Ordering::SeqCst), 0); + } + #[tokio::test] async fn expired_clients_are_rebuilt() { let (address, connections, _) = serve("HTTP/1.1 204 No Content").await; @@ -220,9 +271,12 @@ mod tests { let (address, connections, _) = serve("HTTP/1.1 204 No Content").await; let pool = HttpClientPool::new(Arc::new(FixedResolver(address))); let url = format!("http://media.invalid:{}/doc", address.port()); - for trust_proxy_env in [true, false] { + for proxies in [ + proxied_through("http://proxy.invalid:3128"), + EnvironmentProxies::default(), + ] { let config = HttpClientConfig { - trust_proxy_env, + proxies, ..config("a") }; get(&pool, &config, ClientVariant::Media, &url).await; diff --git a/litellm-rust/crates/http/src/proxy.rs b/litellm-rust/crates/http/src/proxy.rs index 4dc4bf778b8..e51ce3141e5 100644 --- a/litellm-rust/crates/http/src/proxy.rs +++ b/litellm-rust/crates/http/src/proxy.rs @@ -1,15 +1,98 @@ use hyper_util::client::proxy::matcher::Matcher; +use litellm_core_utils::settings::Lookup; -pub struct EnvironmentProxies(Matcher); +#[derive(Clone, Debug, Default, PartialEq, Eq, Hash)] +pub struct EnvironmentProxies { + all: String, + http: String, + https: String, + no: String, +} impl EnvironmentProxies { - pub fn from_environment() -> Self { - Self(Matcher::from_system()) + pub fn from_environment(env: &impl Lookup) -> Self { + if env.get("REQUEST_METHOD").is_some() { + return Self::default(); + } + let first = |upper: &str, lower: &str| { + env.get(upper) + .or_else(|| env.get(lower)) + .unwrap_or_default() + }; + Self { + all: first("ALL_PROXY", "all_proxy"), + http: first("HTTP_PROXY", "http_proxy"), + https: first("HTTPS_PROXY", "https_proxy"), + no: first("NO_PROXY", "no_proxy"), + } } pub fn apply_to(&self, url: &reqwest::Url) -> bool { + let matcher = Matcher::builder() + .all(self.all.clone()) + .http(self.http.clone()) + .https(self.https.clone()) + .no(self.no.clone()) + .build(); url.as_str() .parse::() - .is_ok_and(|uri| self.0.intercept(&uri).is_some()) + .is_ok_and(|uri| matcher.intercept(&uri).is_some()) + } + + pub(crate) fn reqwest_proxies(&self) -> Vec { + let no_proxy = reqwest::NoProxy::from_string(&self.no); + [ + reqwest::Proxy::http(self.http.as_str()), + reqwest::Proxy::https(self.https.as_str()), + reqwest::Proxy::all(self.all.as_str()), + ] + .into_iter() + .filter_map(Result::ok) + .map(|proxy| proxy.no_proxy(no_proxy.clone())) + .collect() + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + fn env_of(values: &'static [(&'static str, &'static str)]) -> impl Fn(&str) -> Option { + move |name| { + values + .iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| value.to_string()) + } + } + + fn url(value: &str) -> reqwest::Url { + reqwest::Url::parse(value).unwrap() + } + + #[rstest] + #[case::http_only(&[("HTTP_PROXY", "http://proxy:3128")], "http://api.test/", true)] + #[case::http_proxy_skips_https(&[("HTTP_PROXY", "http://proxy:3128")], "https://api.test/", false)] + #[case::all_covers_https(&[("ALL_PROXY", "http://proxy:3128")], "https://api.test/", true)] + #[case::lowercase(&[("https_proxy", "http://proxy:3128")], "https://api.test/", true)] + #[case::no_proxy_bypass(&[("HTTPS_PROXY", "http://proxy:3128"), ("NO_PROXY", "api.test")], "https://api.test/", false)] + #[case::cgi_ignores_everything(&[("HTTPS_PROXY", "http://proxy:3128"), ("REQUEST_METHOD", "GET")], "https://api.test/", false)] + #[case::uppercase_wins_even_when_empty(&[("HTTPS_PROXY", ""), ("https_proxy", "http://proxy:3128")], "https://api.test/", false)] + fn proxies_follow_the_injected_environment( + #[case] env: &'static [(&'static str, &'static str)], + #[case] target: &str, + #[case] expected: bool, + ) { + let proxies = EnvironmentProxies::from_environment(&env_of(env)); + assert_eq!(proxies.apply_to(&url(target)), expected); + } + + #[test] + fn an_empty_environment_proxies_nothing() { + let proxies = EnvironmentProxies::from_environment(&env_of(&[])); + assert_eq!(proxies, EnvironmentProxies::default()); + assert!(proxies.reqwest_proxies().is_empty()); } } diff --git a/litellm-rust/crates/http/src/settings.rs b/litellm-rust/crates/http/src/settings.rs index 43c7f6223d2..a6397f1e8e3 100644 --- a/litellm-rust/crates/http/src/settings.rs +++ b/litellm-rust/crates/http/src/settings.rs @@ -5,6 +5,8 @@ use std::{ use litellm_core_utils::settings::{Layer, Lookup, merge}; +use crate::proxy::EnvironmentProxies; + #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub enum SslVerify { Enabled, @@ -44,6 +46,7 @@ pub struct HttpSettingsLayer { pub user_agent: Option, pub tcp_keepalive: Option, pub pool_idle_timeout: Option, + pub proxies: Option, } impl HttpSettingsLayer { @@ -71,6 +74,8 @@ impl HttpSettingsLayer { pool_idle_timeout: env .parsed::("AIOHTTP_KEEPALIVE_TIMEOUT") .map(|timeout| Duration::from_secs(u64::from(timeout))), + proxies: Some(EnvironmentProxies::from_environment(env)) + .filter(|proxies| *proxies != EnvironmentProxies::default()), } } } @@ -95,6 +100,7 @@ impl Layer for HttpSettingsLayer { user_agent: self.user_agent.or(lower.user_agent), tcp_keepalive: self.tcp_keepalive.or(lower.tcp_keepalive), pool_idle_timeout: self.pool_idle_timeout.or(lower.pool_idle_timeout), + proxies: self.proxies.or(lower.proxies), } } } @@ -110,6 +116,7 @@ pub struct HttpSettings { pub http2: bool, pub user_agent: Option, pub trust_proxy_env: bool, + pub proxies: EnvironmentProxies, pub connect_timeout: Duration, pub tcp_keepalive: Option, pub pool_idle_timeout: Duration, @@ -127,6 +134,7 @@ impl Default for HttpSettings { http2: false, user_agent: None, trust_proxy_env: true, + proxies: EnvironmentProxies::default(), connect_timeout: Duration::from_secs(10), tcp_keepalive: None, pool_idle_timeout: Duration::from_secs(120), @@ -160,6 +168,7 @@ impl HttpSettings { pool_idle_timeout: merged .pool_idle_timeout .unwrap_or(defaults.pool_idle_timeout), + proxies: merged.proxies.unwrap_or_default(), ..defaults } } diff --git a/litellm-rust/crates/llms/src/custom_httpx/media.rs b/litellm-rust/crates/llms/src/custom_httpx/media.rs index 572e7f12e54..059d0a05010 100644 --- a/litellm-rust/crates/llms/src/custom_httpx/media.rs +++ b/litellm-rust/crates/llms/src/custom_httpx/media.rs @@ -7,7 +7,7 @@ use std::{ time::Duration, }; -use litellm_http::{ClientVariant, EnvironmentProxies, HttpClientConfig, HttpClientPool}; +use litellm_http::{ClientVariant, HttpClientConfig, HttpClientPool}; use reqwest::{ Url, dns::{Addrs, Name, Resolve, Resolving}, @@ -102,12 +102,8 @@ impl MediaFetcher { config: &HttpClientConfig, url_policy: UrlPolicy, ) -> Result { - let uses_proxy: ProxyMatch = if config.trust_proxy_env { - let proxies = EnvironmentProxies::from_environment(); - Arc::new(move |url| proxies.apply_to(url)) - } else { - Arc::new(|_| false) - }; + let proxies = config.proxies.clone(); + let uses_proxy: ProxyMatch = Arc::new(move |url| proxies.apply_to(url)); Self::with_resolution( pool, config, @@ -443,10 +439,7 @@ mod tests { url_policy: UrlPolicy, uses_proxy: bool, ) -> MediaFetcher { - let direct = HttpClientConfig { - trust_proxy_env: false, - ..Resolution::from(&HttpSettings::default()).config - }; + let direct = Resolution::from(&HttpSettings::default()).config; MediaFetcher::with_resolution( &HttpClientPool::new(Arc::new(LoopbackDnsResolver(pinned_address))), &direct, From d77c144c6cb8b22aa8687c46ef0889df500fc96d Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 20:46:36 -0700 Subject: [PATCH 03/23] refactor(rust): split custom_httpx into litellm-http and the OCR handler custom_httpx mirrored a Python module that mixes transport plumbing with OCR orchestration. The transport half (media fetcher, transport errors, request and header helpers) now lives in litellm-http next to the pool, TLS, proxies and settings, and the OCR request handler moves to base_llm/ocr/handler.rs. Drops the unused deserialize_optional_param and stale dead_code allows. Co-Authored-By: Claude Opus 5 --- litellm-rust/Cargo.lock | 1 + litellm-rust/crates/core/AGENTS.md | 7 +-- litellm-rust/crates/core/Cargo.toml | 2 +- .../core/src/audio_transcription/error.rs | 4 +- .../core/src/audio_transcription/handler.rs | 20 +++----- .../core/src/audio_transcription/prepare.rs | 2 +- .../core/src/chat_completions/common_utils.rs | 2 +- .../crates/core/src/chat_completions/error.rs | 4 +- .../core/src/chat_completions/handler.rs | 32 ++++-------- .../core/src/chat_completions/prepare.rs | 6 +-- .../crates/core/src/chat_completions/tests.rs | 22 +++----- .../crates/core/src/messages/common_utils.rs | 6 +-- .../crates/core/src/messages/error.rs | 4 +- .../crates/core/src/messages/handler.rs | 6 +-- .../crates/core/src/messages/tests.rs | 4 +- litellm-rust/crates/core/src/ocr/client.rs | 5 +- litellm-rust/crates/core/src/ocr/handler.rs | 10 ++-- .../crates/core/src/ocr/provider_config.rs | 4 +- litellm-rust/crates/core/src/ocr/route.rs | 5 +- .../crates/core/src/responses/error.rs | 4 +- .../crates/core/src/responses/websocket.rs | 34 +++++-------- litellm-rust/crates/core/tests/ocr.rs | 25 ++++----- litellm-rust/crates/core/tests/ocr/support.rs | 7 +-- litellm-rust/crates/http/Cargo.toml | 5 ++ litellm-rust/crates/http/src/lib.rs | 3 ++ .../src/custom_httpx => http/src}/media.rs | 17 ++++--- .../http_handler.rs => http/src/request.rs} | 20 -------- .../custom_httpx => http/src}/transport.rs | 13 ++--- litellm-rust/crates/llms/AGENTS.md | 2 +- litellm-rust/crates/llms/Cargo.toml | 2 +- .../ocr/cohere_parse_transformation.rs | 4 +- .../document_intelligence/transformation.rs | 51 ++++++++----------- .../llms/src/azure_ai/ocr/transformation.rs | 7 ++- .../crates/llms/src/base_llm/ocr/document.rs | 24 ++++----- .../crates/llms/src/base_llm/ocr/error.rs | 8 ++- .../ocr/handler.rs} | 24 ++++----- .../crates/llms/src/base_llm/ocr/mod.rs | 1 + .../llms/src/base_llm/ocr/transformation.rs | 8 ++- .../llms/src/cohere/ocr/transformation.rs | 23 ++++----- .../crates/llms/src/custom_httpx/mod.rs | 4 -- litellm-rust/crates/llms/src/lib.rs | 1 - .../llms/src/mistral/ocr/transformation.rs | 18 +++---- .../llms/src/reducto/ocr/transformation.rs | 48 ++++++++--------- .../vertex_ai/ocr/deepseek_transformation.rs | 16 +++--- .../llms/src/vertex_ai/ocr/transformation.rs | 4 +- .../crates/python-bridge/src/errors.rs | 5 +- litellm-rust/crates/python-bridge/src/http.rs | 2 +- .../python-bridge/src/routes/messages/host.rs | 2 +- .../python-bridge/src/routes/ocr/errors.rs | 7 ++- .../python-bridge/src/routes/ocr/mod.rs | 2 +- 50 files changed, 216 insertions(+), 321 deletions(-) rename litellm-rust/crates/{llms/src/custom_httpx => http/src}/media.rs (97%) rename litellm-rust/crates/{llms/src/custom_httpx/http_handler.rs => http/src/request.rs} (93%) rename litellm-rust/crates/{llms/src/custom_httpx => http/src}/transport.rs (88%) rename litellm-rust/crates/llms/src/{custom_httpx/llm_http_handler.rs => base_llm/ocr/handler.rs} (94%) delete mode 100644 litellm-rust/crates/llms/src/custom_httpx/mod.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 83cdbc6a782..5fbddcaffcf 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2141,6 +2141,7 @@ dependencies = [ "reqwest 0.12.28", "rstest", "rustls 0.23.42", + "serde_json", "thiserror 2.0.19", "tokio", "webpki-roots", diff --git a/litellm-rust/crates/core/AGENTS.md b/litellm-rust/crates/core/AGENTS.md index 449c3e647f7..0c8a747019d 100644 --- a/litellm-rust/crates/core/AGENTS.md +++ b/litellm-rust/crates/core/AGENTS.md @@ -5,10 +5,11 @@ litellm-core is the LiteLLM SDK in Rust — it makes the LLM call. Each top-leve Each crate mirrors one top-level Python package, so a Rust path reads as its Python path with the crate name in place of the package directory. Dependencies only point down: - `litellm-types` mirrors `litellm/types/`: pure serde data, no I/O -- `litellm-core-utils` mirrors `litellm/litellm_core_utils/`: pure helpers (provider resolution, prompt factory, call arguments), no network I/O -- `litellm-llms` mirrors `litellm/llms/`: `base_llm//transformation.rs`, `//transformation.rs`, and `custom_httpx/` (HTTP helpers and the OCR request handler) +- `litellm-core-utils` mirrors `litellm/litellm_core_utils/`: pure helpers (provider resolution, prompt factory, call arguments, settings lookup and layer merge), no network I/O +- `litellm-http` is Rust-only and route-neutral: settings resolution, the pooled `reqwest` clients, TLS, proxies, the SSRF-safe media fetcher, request and header helpers, and transport errors. Python's `litellm/llms/custom_httpx/` is split by responsibility instead of mirrored: its transport half lives here, its OCR handler in `litellm-llms` +- `litellm-llms` mirrors `litellm/llms/`: `base_llm//transformation.rs`, `//transformation.rs`, and `base_llm/ocr/handler.rs` (the OCR request handler) - `litellm-core` mirrors the route packages (`litellm/ocr/`, `litellm/messages/`, ...): entrypoints, route request types, provider dispatch, the route machine, and hooks -A route module owns the call entrypoint, route request types (`*Request<'a>`), credential fallback, provider dispatch, and the handler glue that runs a provider config. Provider code never imports from core; when it needs the caller's hooks mid-call it goes through `litellm_llms::custom_httpx::llm_http_handler::CallHooks`, which each route implements over its host. Import every item from its canonical path. Never re-export another crate's items or give an item a second public path; the only re-export allowed is a private submodule surfacing its item at its module root (`mod error; pub use error::Error;`). Handlers belong in core or llms, never in a host crate +A route module owns the call entrypoint, route request types (`*Request<'a>`), credential fallback, provider dispatch, and the handler glue that runs a provider config. Provider code never imports from core; when it needs the caller's hooks mid-call it goes through `litellm_llms::base_llm::ocr::handler::CallHooks`, which each route implements over its host. Import every item from its canonical path. Never re-export another crate's items or give an item a second public path; the only re-export allowed is a private submodule surfacing its item at its module root (`mod error; pub use error::Error;`). Handlers belong in core or llms, never in a host crate Not here: serving HTTP (axum routes, extractors), config file reading, rollout state, databases, or callback execution of any kind. Core runs each route as a machine that yields host operations and call events; which integrations consume those events is the host's business. diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index ab04fb8d4ae..69ae8004d46 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -15,6 +15,7 @@ futures-util.workspace = true base64.workspace = true litellm-auth.workspace = true litellm-auth-aws.workspace = true +litellm-http.workspace = true litellm-llms.workspace = true moka.workspace = true mime_guess = "2.0.5" @@ -36,7 +37,6 @@ veil.workspace = true [dev-dependencies] litellm-auth-gcp.workspace = true -litellm-http.workspace = true litellm-llms = { workspace = true, features = ["test-support"] } rstest.workspace = true rstest_reuse.workspace = true diff --git a/litellm-rust/crates/core/src/audio_transcription/error.rs b/litellm-rust/crates/core/src/audio_transcription/error.rs index 39b08e882f5..122cbab358f 100644 --- a/litellm-rust/crates/core/src/audio_transcription/error.rs +++ b/litellm-rust/crates/core/src/audio_transcription/error.rs @@ -20,9 +20,9 @@ pub enum Error { #[error(transparent)] Auth(#[from] litellm_auth::Error), #[error(transparent)] - Transport(#[from] litellm_llms::custom_httpx::transport::Error), + Transport(#[from] litellm_http::transport::Error), #[error(transparent)] - Headers(#[from] litellm_llms::custom_httpx::http_handler::HeaderError), + Headers(#[from] litellm_http::request::HeaderError), #[error(transparent)] Aws(#[from] litellm_auth_aws::Error), } diff --git a/litellm-rust/crates/core/src/audio_transcription/handler.rs b/litellm-rust/crates/core/src/audio_transcription/handler.rs index 0704f9391b0..503cc922966 100644 --- a/litellm-rust/crates/core/src/audio_transcription/handler.rs +++ b/litellm-rust/crates/core/src/audio_transcription/handler.rs @@ -1,4 +1,4 @@ -use litellm_llms::custom_httpx::http_handler::{http_request, truncate_error_body}; +use litellm_http::request::{http_request, truncate_error_body}; use serde_json::Value; use super::{Error, client::http_client}; @@ -18,23 +18,17 @@ pub async fn execute_audio_transcription_provider_call( request_builder = request_builder.timeout(duration); } let response = http_request(request_builder).await.map_err(|error| { - Error::Transport(litellm_llms::custom_httpx::transport::Error::Network( - error.to_string(), - )) + Error::Transport(litellm_http::transport::Error::Network(error.to_string())) })?; let status = response.status(); let text = response.text().await.map_err(|error| { - Error::Transport(litellm_llms::custom_httpx::transport::Error::Network( - error.to_string(), - )) + Error::Transport(litellm_http::transport::Error::Network(error.to_string())) })?; if !status.is_success() { - return Err(Error::Transport( - litellm_llms::custom_httpx::transport::Error::Http { - status: status.as_u16(), - body: truncate_error_body(&text), - }, - )); + return Err(Error::Transport(litellm_http::transport::Error::Http { + status: status.as_u16(), + body: truncate_error_body(&text), + })); } let response_json = serde_json::from_str(&text) .map_err(|error| Error::InvalidResponse(format!("invalid audio response JSON: {error}")))?; diff --git a/litellm-rust/crates/core/src/audio_transcription/prepare.rs b/litellm-rust/crates/core/src/audio_transcription/prepare.rs index 193122db733..829617d26bd 100644 --- a/litellm-rust/crates/core/src/audio_transcription/prepare.rs +++ b/litellm-rust/crates/core/src/audio_transcription/prepare.rs @@ -1,10 +1,10 @@ use litellm_core_utils::get_llm_provider_logic::{CustomLlmProvider, get_custom_llm_provider}; +use litellm_http::request::{has_header, string_headers}; use litellm_llms::{ base_llm::audio_transcription::transformation::{ AudioTranscriptionAuth, BaseAudioTranscriptionConfig, }, bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG, - custom_httpx::http_handler::{has_header, string_headers}, }; use super::Error; diff --git a/litellm-rust/crates/core/src/chat_completions/common_utils.rs b/litellm-rust/crates/core/src/chat_completions/common_utils.rs index cc9459793df..4ed39a90366 100644 --- a/litellm-rust/crates/core/src/chat_completions/common_utils.rs +++ b/litellm-rust/crates/core/src/chat_completions/common_utils.rs @@ -1,8 +1,8 @@ +use litellm_http::request::string_headers as shared_string_headers; use litellm_llms::{ anthropic::chat::transformation::ANTHROPIC_CHAT_COMPLETIONS_CONFIG, base_llm::chat::transformation::BaseConfig, bedrock::chat::converse_transformation::BEDROCK_CHAT_COMPLETIONS_CONFIG, - custom_httpx::http_handler::string_headers as shared_string_headers, }; use serde_json::{Map, Value}; diff --git a/litellm-rust/crates/core/src/chat_completions/error.rs b/litellm-rust/crates/core/src/chat_completions/error.rs index 39b08e882f5..122cbab358f 100644 --- a/litellm-rust/crates/core/src/chat_completions/error.rs +++ b/litellm-rust/crates/core/src/chat_completions/error.rs @@ -20,9 +20,9 @@ pub enum Error { #[error(transparent)] Auth(#[from] litellm_auth::Error), #[error(transparent)] - Transport(#[from] litellm_llms::custom_httpx::transport::Error), + Transport(#[from] litellm_http::transport::Error), #[error(transparent)] - Headers(#[from] litellm_llms::custom_httpx::http_handler::HeaderError), + Headers(#[from] litellm_http::request::HeaderError), #[error(transparent)] Aws(#[from] litellm_auth_aws::Error), } diff --git a/litellm-rust/crates/core/src/chat_completions/handler.rs b/litellm-rust/crates/core/src/chat_completions/handler.rs index 034408bdf17..b73d4838760 100644 --- a/litellm-rust/crates/core/src/chat_completions/handler.rs +++ b/litellm-rust/crates/core/src/chat_completions/handler.rs @@ -1,7 +1,5 @@ -use litellm_llms::{ - base_llm::chat::transformation::{ChatCompletionsAuth, ProviderChatResponseData}, - custom_httpx::http_handler::{http_request, truncate_error_body}, -}; +use litellm_http::request::{http_request, truncate_error_body}; +use litellm_llms::base_llm::chat::transformation::{ChatCompletionsAuth, ProviderChatResponseData}; use litellm_types::utils::ChatCompletionsResponse; use serde_json::Value; @@ -34,30 +32,22 @@ pub(super) async fn execute_chat_completions_provider_call( // so the host can still serve it. Everything else here, a timeout // above all, may have reached the provider and been answered. if err.is_connect() || err.is_builder() { - Error::Transport(litellm_llms::custom_httpx::transport::Error::Connect( - err.to_string(), - )) + Error::Transport(litellm_http::transport::Error::Connect(err.to_string())) } else { - Error::Transport(litellm_llms::custom_httpx::transport::Error::Network( - err.to_string(), - )) + Error::Transport(litellm_http::transport::Error::Network(err.to_string())) } })?; let status = response.status(); let text = response.text().await.map_err(|err| { - Error::Transport(litellm_llms::custom_httpx::transport::Error::Network( - err.to_string(), - )) + Error::Transport(litellm_http::transport::Error::Network(err.to_string())) })?; if !status.is_success() { - return Err(Error::Transport( - litellm_llms::custom_httpx::transport::Error::Http { - status: status.as_u16(), - body: truncate_error_body(&text), - }, - )); + return Err(Error::Transport(litellm_http::transport::Error::Http { + status: status.as_u16(), + body: truncate_error_body(&text), + })); } let body: Value = serde_json::from_str(&text).map_err(|err| { @@ -82,9 +72,7 @@ pub(super) async fn execute_chat_completions_provider_call( pub(super) fn as_response_error(err: Error) -> Error { match err { already @ (Error::InvalidResponse(_) - | Error::Transport(litellm_llms::custom_httpx::transport::Error::Http { - .. - })) => already, + | Error::Transport(litellm_http::transport::Error::Http { .. })) => already, other => Error::InvalidResponse(other.to_string()), } } diff --git a/litellm-rust/crates/core/src/chat_completions/prepare.rs b/litellm-rust/crates/core/src/chat_completions/prepare.rs index d408ea6574e..d0aa1e88011 100644 --- a/litellm-rust/crates/core/src/chat_completions/prepare.rs +++ b/litellm-rust/crates/core/src/chat_completions/prepare.rs @@ -1,8 +1,6 @@ use litellm_core_utils::get_llm_provider_logic::{CustomLlmProvider, get_custom_llm_provider}; -use litellm_llms::{ - base_llm::chat::transformation::{BaseConfig, ChatCompletionsAuth}, - custom_httpx::http_handler::has_header, -}; +use litellm_http::request::has_header; +use litellm_llms::base_llm::chat::transformation::{BaseConfig, ChatCompletionsAuth}; use litellm_types::llms::openai::ChatMessage; use serde_json::Value; diff --git a/litellm-rust/crates/core/src/chat_completions/tests.rs b/litellm-rust/crates/core/src/chat_completions/tests.rs index cbc4995ce0d..dcaa3397add 100644 --- a/litellm-rust/crates/core/src/chat_completions/tests.rs +++ b/litellm-rust/crates/core/src/chat_completions/tests.rs @@ -265,7 +265,7 @@ fn rejects_non_string_extra_headers() { call.extra_headers = Some(Map::from_iter([("x-trace".to_string(), json!(7))])); assert_eq!( decline(call), - Error::Headers(litellm_llms::custom_httpx::http_handler::HeaderError { + Error::Headers(litellm_http::request::HeaderError { context: "chat completions", name: "x-trace".to_string(), actual: "number", @@ -771,10 +771,7 @@ mod round_trip { assert!( matches!( err, - Error::Transport(litellm_llms::custom_httpx::transport::Error::Http { - status: 429, - .. - }) + Error::Transport(litellm_http::transport::Error::Http { status: 429, .. }) ), "expected a 429, got {err:?}" ); @@ -801,7 +798,7 @@ mod round_trip { assert!( matches!( err, - Error::Transport(litellm_llms::custom_httpx::transport::Error::Connect(_)) + Error::Transport(litellm_http::transport::Error::Connect(_)) ), "expected a pre-send connect failure, got {err:?}" ); @@ -825,16 +822,11 @@ mod round_trip { } // An upstream status is already unambiguous, so it survives intact. assert!(matches!( - as_response_error(Error::Transport( - litellm_llms::custom_httpx::transport::Error::Http { - status: 500, - body: "boom".to_string() - } - )), - Error::Transport(litellm_llms::custom_httpx::transport::Error::Http { + as_response_error(Error::Transport(litellm_http::transport::Error::Http { status: 500, - .. - }) + body: "boom".to_string() + })), + Error::Transport(litellm_http::transport::Error::Http { status: 500, .. }) )); } } diff --git a/litellm-rust/crates/core/src/messages/common_utils.rs b/litellm-rust/crates/core/src/messages/common_utils.rs index ec392324784..dcefa3ebffc 100644 --- a/litellm-rust/crates/core/src/messages/common_utils.rs +++ b/litellm-rust/crates/core/src/messages/common_utils.rs @@ -1,11 +1,9 @@ -pub(super) use litellm_llms::custom_httpx::http_handler::{ - has_bearer_auth, has_header, truncate_error_body, -}; +use litellm_http::request::string_headers as shared_string_headers; +pub(super) use litellm_http::request::{has_bearer_auth, has_header, truncate_error_body}; use litellm_llms::{ anthropic::experimental_pass_through::messages::transformation::ANTHROPIC_MESSAGES_CONFIG, azure_ai::anthropic::messages_transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG, base_llm::anthropic_messages::transformation::BaseAnthropicMessagesConfig, - custom_httpx::http_handler::string_headers as shared_string_headers, }; use serde_json::{Map, Value}; diff --git a/litellm-rust/crates/core/src/messages/error.rs b/litellm-rust/crates/core/src/messages/error.rs index 71bb748c50d..51fb764032c 100644 --- a/litellm-rust/crates/core/src/messages/error.rs +++ b/litellm-rust/crates/core/src/messages/error.rs @@ -15,9 +15,9 @@ pub enum Error { #[error(transparent)] Auth(#[from] litellm_auth::Error), #[error(transparent)] - Transport(#[from] litellm_llms::custom_httpx::transport::Error), + Transport(#[from] litellm_http::transport::Error), #[error(transparent)] - Headers(#[from] litellm_llms::custom_httpx::http_handler::HeaderError), + Headers(#[from] litellm_http::request::HeaderError), } impl From for Error { diff --git a/litellm-rust/crates/core/src/messages/handler.rs b/litellm-rust/crates/core/src/messages/handler.rs index 22e2c398ff7..fe7e8bb4b80 100644 --- a/litellm-rust/crates/core/src/messages/handler.rs +++ b/litellm-rust/crates/core/src/messages/handler.rs @@ -1,9 +1,7 @@ use std::time::Duration; -use litellm_llms::{ - base_llm::anthropic_messages::transformation::BaseAnthropicMessagesConfig, - custom_httpx::{http_handler::http_request, transport::Error as TransportError}, -}; +use litellm_http::{request::http_request, transport::Error as TransportError}; +use litellm_llms::base_llm::anthropic_messages::transformation::BaseAnthropicMessagesConfig; use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse; use serde_json::Value; diff --git a/litellm-rust/crates/core/src/messages/tests.rs b/litellm-rust/crates/core/src/messages/tests.rs index 55d8ead8e8b..057b42a316c 100644 --- a/litellm-rust/crates/core/src/messages/tests.rs +++ b/litellm-rust/crates/core/src/messages/tests.rs @@ -82,7 +82,7 @@ fn string_headers_rejects_non_string_values() { let err = string_headers(Some(headers)).expect_err("non-string header rejected"); assert_eq!( err, - Error::Headers(litellm_llms::custom_httpx::http_handler::HeaderError { + Error::Headers(litellm_http::request::HeaderError { context: "messages", name: "x-count".to_string(), actual: "number", @@ -432,7 +432,7 @@ async fn messages_maps_provider_error_status_to_http_error() { assert!(matches!( err, - Error::Transport(litellm_llms::custom_httpx::transport::Error::Http { status: 401, .. }) + Error::Transport(litellm_http::transport::Error::Http { status: 401, .. }) )); } diff --git a/litellm-rust/crates/core/src/ocr/client.rs b/litellm-rust/crates/core/src/ocr/client.rs index c7b4751bd9e..e635f93a294 100644 --- a/litellm-rust/crates/core/src/ocr/client.rs +++ b/litellm-rust/crates/core/src/ocr/client.rs @@ -1,6 +1,5 @@ -use litellm_llms::{ - base_llm::ocr::{error::Error, transformation::LiteLLMOcrResponse}, - custom_httpx::llm_http_handler::OcrClient, +use litellm_llms::base_llm::ocr::{ + error::Error, handler::OcrClient, transformation::LiteLLMOcrResponse, }; use crate::ocr::{ diff --git a/litellm-rust/crates/core/src/ocr/handler.rs b/litellm-rust/crates/core/src/ocr/handler.rs index bbf9cfa0e02..f49976de043 100644 --- a/litellm-rust/crates/core/src/ocr/handler.rs +++ b/litellm-rust/crates/core/src/ocr/handler.rs @@ -1,12 +1,10 @@ use futures_util::future::BoxFuture; use litellm_auth::SecretValue; use litellm_host::event::{MachineEvent, RawResponse, RequestContext, WireRequest}; -use litellm_llms::{ - base_llm::ocr::{ - error::Error, - transformation::{LiteLLMOcrResponse, PreparedOcrRequest}, - }, - custom_httpx::llm_http_handler::{CallHooks, OcrClient}, +use litellm_llms::base_llm::ocr::{ + error::Error, + handler::{CallHooks, OcrClient}, + transformation::{LiteLLMOcrResponse, PreparedOcrRequest}, }; use serde_json::Value; diff --git a/litellm-rust/crates/core/src/ocr/provider_config.rs b/litellm-rust/crates/core/src/ocr/provider_config.rs index 14b34ea4564..ee9ba76928d 100644 --- a/litellm-rust/crates/core/src/ocr/provider_config.rs +++ b/litellm-rust/crates/core/src/ocr/provider_config.rs @@ -7,13 +7,13 @@ use litellm_llms::{ }, base_llm::ocr::{ error::Error, + handler::{self, CallHooks, OcrClient}, transformation::{ BaseOcrConfig, LiteLLMOcrResponse, OcrCredentialInputs, OcrDocument, PreparedOcrRequest, ResolvedOcrCredentials, }, }, cohere::ocr::transformation::CohereParseConfig, - custom_httpx::llm_http_handler::{self, CallHooks, OcrClient}, mistral::ocr::transformation::MistralOcrConfig, reducto::ocr::transformation::{ReductoParseLegacyConfig, ReductoParseV3Config}, vertex_ai::ocr::{ @@ -116,7 +116,7 @@ impl OcrConfigKind { request: &PreparedOcrRequest, hooks: &dyn CallHooks, ) -> Result { - with_config!(self, config => llm_http_handler::ocr(&config, client, request, hooks).await) + with_config!(self, config => handler::ocr(&config, client, request, hooks).await) } } diff --git a/litellm-rust/crates/core/src/ocr/route.rs b/litellm-rust/crates/core/src/ocr/route.rs index bfc8c5ca965..26c9ac27102 100644 --- a/litellm-rust/crates/core/src/ocr/route.rs +++ b/litellm-rust/crates/core/src/ocr/route.rs @@ -6,9 +6,8 @@ use litellm_host::{ machine::{HostChannel, HostTokenProvider, MachineFault, RouteMachine, TokenRoute}, route::Route, }; -use litellm_llms::{ - base_llm::ocr::{error::Error, transformation::LiteLLMOcrResponse}, - custom_httpx::llm_http_handler::OcrClient, +use litellm_llms::base_llm::ocr::{ + error::Error, handler::OcrClient, transformation::LiteLLMOcrResponse, }; use super::handler::perform_ocr_request; diff --git a/litellm-rust/crates/core/src/responses/error.rs b/litellm-rust/crates/core/src/responses/error.rs index 677db2e08de..1c940d8ed9b 100644 --- a/litellm-rust/crates/core/src/responses/error.rs +++ b/litellm-rust/crates/core/src/responses/error.rs @@ -11,7 +11,7 @@ pub enum Error { #[error(transparent)] Auth(#[from] litellm_auth::Error), #[error(transparent)] - Transport(#[from] litellm_llms::custom_httpx::transport::Error), + Transport(#[from] litellm_http::transport::Error), #[error(transparent)] - Headers(#[from] litellm_llms::custom_httpx::http_handler::HeaderError), + Headers(#[from] litellm_http::request::HeaderError), } diff --git a/litellm-rust/crates/core/src/responses/websocket.rs b/litellm-rust/crates/core/src/responses/websocket.rs index ccf4aa75149..f57ba65a6fb 100644 --- a/litellm-rust/crates/core/src/responses/websocket.rs +++ b/litellm-rust/crates/core/src/responses/websocket.rs @@ -95,9 +95,7 @@ impl ResponsesWebSocketConnection { timeout: Option, ) -> Result { let mut request = url.into_client_request().map_err(|error| { - Error::Transport(litellm_llms::custom_httpx::transport::Error::Network( - error.to_string(), - )) + Error::Transport(litellm_http::transport::Error::Network(error.to_string())) })?; for (name, value) in headers { let header_name = name @@ -110,7 +108,7 @@ impl ResponsesWebSocketConnection { let connect = connect_upstream(request); let result = match timeout { Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| { - Error::Transport(litellm_llms::custom_httpx::transport::Error::Network( + Error::Transport(litellm_http::transport::Error::Network( "Responses WebSocket connection timed out".into(), )) })?, @@ -118,14 +116,12 @@ impl ResponsesWebSocketConnection { }; let (socket, _) = result.map_err(|error| match *error { tokio_tungstenite::tungstenite::Error::Http(response) => { - Error::Transport(litellm_llms::custom_httpx::transport::Error::Http { + Error::Transport(litellm_http::transport::Error::Http { status: response.status().as_u16(), body: String::new(), }) } - other => Error::Transport(litellm_llms::custom_httpx::transport::Error::Network( - other.to_string(), - )), + other => Error::Transport(litellm_http::transport::Error::Network(other.to_string())), })?; Ok(Self { socket: Arc::new(Mutex::new(Some(socket))), @@ -135,16 +131,12 @@ impl ResponsesWebSocketConnection { pub async fn send_text(&self, text: String) -> Result<(), Error> { let mut socket = self.socket.lock().await; let Some(socket) = socket.as_mut() else { - return Err(Error::Transport( - litellm_llms::custom_httpx::transport::Error::Network( - "Responses WebSocket is closed".into(), - ), - )); + return Err(Error::Transport(litellm_http::transport::Error::Network( + "Responses WebSocket is closed".into(), + ))); }; socket.send(Message::Text(text)).await.map_err(|error| { - Error::Transport(litellm_llms::custom_httpx::transport::Error::Network( - error.to_string(), - )) + Error::Transport(litellm_http::transport::Error::Network(error.to_string())) }) } @@ -160,9 +152,9 @@ impl ResponsesWebSocketConnection { .map_err(|error| Error::InvalidResponse(error.to_string())), Some(Ok(Message::Close(_))) | None => Ok(None), Some(Ok(_)) => Ok(None), - Some(Err(error)) => Err(Error::Transport( - litellm_llms::custom_httpx::transport::Error::Network(error.to_string()), - )), + Some(Err(error)) => Err(Error::Transport(litellm_http::transport::Error::Network( + error.to_string(), + ))), } } @@ -170,9 +162,7 @@ impl ResponsesWebSocketConnection { let mut socket = self.socket.lock().await; if let Some(socket) = socket.as_mut() { socket.close(None).await.map_err(|error| { - Error::Transport(litellm_llms::custom_httpx::transport::Error::Network( - error.to_string(), - )) + Error::Transport(litellm_http::transport::Error::Network(error.to_string())) })?; } *socket = None; diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index e7a8fc0abc1..b999c43de8b 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -6,16 +6,14 @@ use litellm_host::{ host::{Host, HostOp, HostResult}, machine::{HostFailure, Machine, MachineStep}, }; -use litellm_http::{HttpClientPool, HttpSettings, Resolution}; -use litellm_llms::{ - base_llm::ocr::{ - error::Error as OcrError, - transformation::{LiteLLMOcrResponse, OCR_RESPONSE_MAX_BYTES, OcrTransportConfig}, - }, - custom_httpx::{ - llm_http_handler::OcrClient, - media::{PublicDnsResolver, UrlPolicy}, - }, +use litellm_http::{ + HttpClientPool, HttpSettings, Resolution, + media::{PublicDnsResolver, UrlPolicy}, +}; +use litellm_llms::base_llm::ocr::{ + error::Error as OcrError, + handler::OcrClient, + transformation::{LiteLLMOcrResponse, OCR_RESPONSE_MAX_BYTES, OcrTransportConfig}, }; use rstest::rstest; use serde_json::{Value, json}; @@ -624,7 +622,7 @@ async fn read_bounded_response(response: Vec, limit: usize) -> Result { + OcrError::Transport(litellm_http::transport::Error::Http { status, body }) => { assert_eq!(status, 429); assert_eq!(body, prefix); } diff --git a/litellm-rust/crates/core/tests/ocr/support.rs b/litellm-rust/crates/core/tests/ocr/support.rs index b368a754656..974fa3d6655 100644 --- a/litellm-rust/crates/core/tests/ocr/support.rs +++ b/litellm-rust/crates/core/tests/ocr/support.rs @@ -2,9 +2,10 @@ use std::sync::{Arc, Mutex}; use futures_util::future::BoxFuture; use litellm_host::event::WireRequest; -use litellm_llms::{ - base_llm::ocr::{error::Error, transformation::LiteLLMOcrResponse}, - custom_httpx::llm_http_handler::{CallHooks, OcrClient}, +use litellm_llms::base_llm::ocr::{ + error::Error, + handler::{CallHooks, OcrClient}, + transformation::LiteLLMOcrResponse, }; use serde_json::{Value, json}; use tokio::{ diff --git a/litellm-rust/crates/http/Cargo.toml b/litellm-rust/crates/http/Cargo.toml index b0dc7693840..4f94f37a8d5 100644 --- a/litellm-rust/crates/http/Cargo.toml +++ b/litellm-rust/crates/http/Cargo.toml @@ -5,13 +5,18 @@ edition.workspace = true license.workspace = true repository.workspace = true +[features] +test-support = [] + [dependencies] http.workspace = true litellm-core-utils.workspace = true hyper-util.workspace = true reqwest.workspace = true rustls.workspace = true +serde_json.workspace = true thiserror.workspace = true +tokio.workspace = true webpki-roots.workspace = true [dev-dependencies] diff --git a/litellm-rust/crates/http/src/lib.rs b/litellm-rust/crates/http/src/lib.rs index ddbc3b63b08..c6d9959348d 100644 --- a/litellm-rust/crates/http/src/lib.rs +++ b/litellm-rust/crates/http/src/lib.rs @@ -1,9 +1,12 @@ mod config; mod error; +pub mod media; mod pool; mod proxy; +pub mod request; mod settings; mod tls; +pub mod transport; pub use config::{HttpClientConfig, Resolution, Verify}; pub use error::Error; diff --git a/litellm-rust/crates/llms/src/custom_httpx/media.rs b/litellm-rust/crates/http/src/media.rs similarity index 97% rename from litellm-rust/crates/llms/src/custom_httpx/media.rs rename to litellm-rust/crates/http/src/media.rs index 059d0a05010..ae3f55b476a 100644 --- a/litellm-rust/crates/llms/src/custom_httpx/media.rs +++ b/litellm-rust/crates/http/src/media.rs @@ -7,12 +7,13 @@ use std::{ time::Duration, }; -use litellm_http::{ClientVariant, HttpClientConfig, HttpClientPool}; use reqwest::{ Url, dns::{Addrs, Name, Resolve, Resolving}, }; +use crate::{ClientVariant, HttpClientConfig, HttpClientPool}; + #[derive(Debug, thiserror::Error)] pub enum Error { #[error("media URL rejected by network policy")] @@ -32,7 +33,7 @@ pub enum Error { #[error("media download timed out")] Timeout, #[error("{0}")] - Transport(#[from] crate::custom_httpx::transport::Error), + Transport(#[from] crate::transport::Error), } #[derive(Clone, Debug, PartialEq, Eq)] @@ -101,7 +102,7 @@ impl MediaFetcher { pool: &HttpClientPool, config: &HttpClientConfig, url_policy: UrlPolicy, - ) -> Result { + ) -> Result { let proxies = config.proxies.clone(); let uses_proxy: ProxyMatch = Arc::new(move |url| proxies.apply_to(url)); Self::with_resolution( @@ -119,7 +120,7 @@ impl MediaFetcher { url_policy: UrlPolicy, address_resolver: Arc, uses_proxy: ProxyMatch, - ) -> Result { + ) -> Result { Ok(Self { pinned: pool.client(config, ClientVariant::Media)?, unpinned: pool.client(config, ClientVariant::UnpinnedMedia)?, @@ -164,7 +165,7 @@ impl MediaFetcher { .get(url.clone()) .send() .await - .map_err(crate::custom_httpx::transport::Error::from)?; + .map_err(crate::transport::Error::from)?; if response.status().is_redirection() { if redirects_followed == policy.max_redirects { return Err(Error::TooManyRedirects); @@ -195,7 +196,7 @@ impl MediaFetcher { while let Some(chunk) = response .chunk() .await - .map_err(crate::custom_httpx::transport::Error::from)? + .map_err(crate::transport::Error::from)? { enforce_download_size(bytes.len() as u64 + chunk.len() as u64, policy.max_bytes)?; bytes.extend_from_slice(&chunk); @@ -245,7 +246,7 @@ impl MediaFetcher { .address_resolver .resolve(host, port) .await - .map_err(|error| crate::custom_httpx::transport::Error::Network(error.to_string()))?; + .map_err(|error| crate::transport::Error::Network(error.to_string()))?; validate_addresses(&addresses) } } @@ -346,13 +347,13 @@ impl Resolve for PublicDnsResolver { mod tests { use std::collections::HashSet; - use litellm_http::{HttpSettings, Resolution}; use tokio::{ io::{AsyncReadExt, AsyncWriteExt}, net::TcpListener, }; use super::*; + use crate::{HttpSettings, Resolution}; async fn serve(response: &'static [u8]) -> (Url, tokio::task::JoinHandle<()>) { let listener = TcpListener::bind("127.0.0.1:0") diff --git a/litellm-rust/crates/llms/src/custom_httpx/http_handler.rs b/litellm-rust/crates/http/src/request.rs similarity index 93% rename from litellm-rust/crates/llms/src/custom_httpx/http_handler.rs rename to litellm-rust/crates/http/src/request.rs index e629be37336..874a0f3abf9 100644 --- a/litellm-rust/crates/llms/src/custom_httpx/http_handler.rs +++ b/litellm-rust/crates/http/src/request.rs @@ -13,20 +13,12 @@ use serde_json::{Map, Value}; /// before truncation, so provider bodies are bounded and data-minimized. const UPSTREAM_ERROR_BODY_MAX_CHARS: usize = 256; -#[allow( - dead_code, - reason = "used by the OCR architecture in the next stacked PR" -)] pub enum HeaderPolicy<'a> { All, Only(&'a [&'a str]), Except(&'a [&'a str]), } -#[allow( - dead_code, - reason = "used by the OCR architecture in the next stacked PR" -)] pub fn with_headers( builder: reqwest::RequestBuilder, headers: &[(String, String)], @@ -107,18 +99,6 @@ pub fn has_bearer_auth(headers: &[(String, String)]) -> bool { }) } -#[allow( - dead_code, - reason = "used by the OCR architecture in the next stacked PR" -)] -pub fn deserialize_optional_param<'de, D, T>(deserializer: D) -> Result>, D::Error> -where - D: serde::Deserializer<'de>, - T: serde::Deserialize<'de>, -{ - as serde::Deserialize>::deserialize(deserializer).map(Some) -} - #[cfg(test)] mod tests { use serde_json::json; diff --git a/litellm-rust/crates/llms/src/custom_httpx/transport.rs b/litellm-rust/crates/http/src/transport.rs similarity index 88% rename from litellm-rust/crates/llms/src/custom_httpx/transport.rs rename to litellm-rust/crates/http/src/transport.rs index c42cdf410f6..8814925bbf2 100644 --- a/litellm-rust/crates/llms/src/custom_httpx/transport.rs +++ b/litellm-rust/crates/http/src/transport.rs @@ -46,11 +46,8 @@ mod tests { .send() .await .expect_err("invalid port"); - let error = crate::custom_httpx::transport::Error::from_reqwest_before_dispatch(error); - assert!(matches!( - error, - crate::custom_httpx::transport::Error::Connect(_) - )); + let error = crate::transport::Error::from_reqwest_before_dispatch(error); + assert!(matches!(error, crate::transport::Error::Connect(_))); assert!(!error.to_string().contains("secret")); assert!(!error.to_string().contains("private")); } @@ -76,7 +73,7 @@ mod tests { .await .expect_err("nothing listens on the port"); let root_cause = root_cause(&error).expect("reqwest reports a cause"); - let message = crate::custom_httpx::transport::Error::from(error).to_string(); + let message = crate::transport::Error::from(error).to_string(); assert!(message.contains(&root_cause), "{message}"); assert!(!message.contains("secret")); } @@ -105,8 +102,8 @@ mod tests { let error = response.expect_err("server does not respond"); assert!(error.is_timeout()); assert!(matches!( - crate::custom_httpx::transport::Error::from_reqwest_before_dispatch(error), - crate::custom_httpx::transport::Error::Network(_) + crate::transport::Error::from_reqwest_before_dispatch(error), + crate::transport::Error::Network(_) )); } } diff --git a/litellm-rust/crates/llms/AGENTS.md b/litellm-rust/crates/llms/AGENTS.md index 09fe20cd9d6..bd1c58142fd 100644 --- a/litellm-rust/crates/llms/AGENTS.md +++ b/litellm-rust/crates/llms/AGENTS.md @@ -1,4 +1,4 @@ -litellm-llms mirrors `litellm/llms/`: base config traits, provider transformations, and the `custom_httpx` handlers. See `../core/AGENTS.md` for how the crates layer. +litellm-llms mirrors `litellm/llms/`: base config traits, provider transformations, and the OCR request handler in `base_llm/ocr/handler.rs`. Transport code (clients, media fetching, header helpers, transport errors) lives in `litellm-http`. See `../core/AGENTS.md` for how the crates layer. ## Python/Rust transformation pairs diff --git a/litellm-rust/crates/llms/Cargo.toml b/litellm-rust/crates/llms/Cargo.toml index a81a4427b4d..7afc4171ca8 100644 --- a/litellm-rust/crates/llms/Cargo.toml +++ b/litellm-rust/crates/llms/Cargo.toml @@ -6,7 +6,7 @@ license.workspace = true repository.workspace = true [features] -test-support = [] +test-support = ["litellm-http/test-support"] [dependencies] litellm-types.workspace = true diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/cohere_parse_transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/cohere_parse_transformation.rs index f55f6b067e4..86ee0d96895 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/cohere_parse_transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/cohere_parse_transformation.rs @@ -5,6 +5,7 @@ use crate::{ base_llm::ocr::{ document::{inline_remote_document, validate_inline_document}, error::Error, + handler::OcrClient, transformation::{ BaseOcrConfig, LiteLLMOcrResponse, OcrDocument, OcrRequestContext, OcrResponseFormat, PreparedOcrRequest, @@ -13,7 +14,6 @@ use crate::{ cohere::ocr::transformation::{ CohereOptions, CohereParseConfig, CohereRequest, validate_document, }, - custom_httpx::llm_http_handler::OcrClient, }; #[derive(Default)] @@ -108,7 +108,7 @@ impl BaseOcrConfig for AzureAICohereParseConfig { } fn validate_request_body(&self, body: &Value) -> Result<(), Error> { - let document = crate::custom_httpx::llm_http_handler::body_document(body)?; + let document = crate::base_llm::ocr::handler::body_document(body)?; validate_document(&document)?; validate_inline_document(&document) } diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs index 2e398d0287e..a347375510d 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs @@ -14,19 +14,16 @@ use serde_json::{Map, Value}; use serde_with::serde_as; use tokio::time::Instant; -use crate::{ - base_llm::ocr::{ - document::InlineDocument, - error::Error, - transformation::{ - BaseOcrConfig, DecodedOcrResponse, LiteLLMOcrResponse, OCR_INLINE_MAX_BYTES, - OCR_POLL_RETRY_SECS, OcrConnection, OcrCredentialInputs, OcrDocument, OcrPage, - OcrPageDimensions, OcrResponseContext, OcrResponseFormat, OcrUsageInfo, - PreparedOcrRequest, ResolvedOcrCredentials, credential_env, - decode_and_normalize_response, decode_response, - }, +use crate::base_llm::ocr::{ + document::InlineDocument, + error::Error, + handler::{CallHooks, OcrClient, read_json_response}, + transformation::{ + BaseOcrConfig, DecodedOcrResponse, LiteLLMOcrResponse, OCR_INLINE_MAX_BYTES, + OCR_POLL_RETRY_SECS, OcrConnection, OcrCredentialInputs, OcrDocument, OcrPage, + OcrPageDimensions, OcrResponseContext, OcrResponseFormat, OcrUsageInfo, PreparedOcrRequest, + ResolvedOcrCredentials, credential_env, decode_and_normalize_response, decode_response, }, - custom_httpx::llm_http_handler::{CallHooks, OcrClient, read_json_response}, }; const AZURE_DI_API_VERSION: &str = "2024-11-30"; @@ -440,7 +437,7 @@ async fn read_operation_response( hooks: &dyn CallHooks, ) -> Result, Error> { if response.status() != reqwest::StatusCode::ACCEPTED { - let bytes = crate::custom_httpx::llm_http_handler::read_response_bytes( + let bytes = crate::base_llm::ocr::handler::read_response_bytes( response, connection.max_response_bytes, ) @@ -462,11 +459,9 @@ async fn read_operation_response( { return Err(Error::PollOrigin); } - let bytes = crate::custom_httpx::llm_http_handler::read_response_bytes( - response, - connection.max_response_bytes, - ) - .await?; + let bytes = + crate::base_llm::ocr::handler::read_response_bytes(response, connection.max_response_bytes) + .await?; hooks.response_received(&bytes).await?; poll_operation(http_client, operation, headers, connection, native, hooks).await } @@ -491,21 +486,19 @@ async fn poll_operation( let builder = http_client .get(url.clone()) .timeout(remaining.min(connection.timeout)); - let builder = crate::custom_httpx::http_handler::with_headers( + let builder = litellm_http::request::with_headers( builder, headers, - crate::custom_httpx::http_handler::HeaderPolicy::Only(&[ + litellm_http::request::HeaderPolicy::Only(&[ AZURE_DI_SUBSCRIPTION_HEADER, "authorization", ]), ); - let response = tokio::time::timeout_at( - deadline, - crate::custom_httpx::http_handler::http_request(builder), - ) - .await - .map_err(|_| Error::PollTimeout)? - .map_err(crate::custom_httpx::transport::Error::from)?; + let response = + tokio::time::timeout_at(deadline, litellm_http::request::http_request(builder)) + .await + .map_err(|_| Error::PollTimeout)? + .map_err(litellm_http::transport::Error::from)?; let retry = response .headers() .get(reqwest::header::RETRY_AFTER) @@ -580,8 +573,8 @@ impl AzureDocumentIntelligenceOcrConfig { config: &AzureAuthInputs, env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> Result, Error> { - if crate::custom_httpx::http_handler::has_header(&connection.extra_headers, "authorization") - || crate::custom_httpx::http_handler::has_header( + if litellm_http::request::has_header(&connection.extra_headers, "authorization") + || litellm_http::request::has_header( &connection.extra_headers, AZURE_DI_SUBSCRIPTION_HEADER, ) diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs index 7ef051e8986..4a04910aa9a 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs @@ -7,12 +7,12 @@ use crate::{ base_llm::ocr::{ document::{inline_remote_document, validate_inline_document}, error::Error, + handler::OcrClient, transformation::{ BaseOcrConfig, LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrRequestContext, OcrResponseFormat, PreparedOcrRequest, credential_env, }, }, - custom_httpx::llm_http_handler::OcrClient, mistral::ocr::transformation::{MistralOcrConfig, MistralOcrRequest}, }; @@ -107,7 +107,7 @@ impl BaseOcrConfig for AzureAiOcrConfig { } fn validate_request_body(&self, body: &Value) -> Result<(), Error> { - validate_inline_document(&crate::custom_httpx::llm_http_handler::body_document(body)?) + validate_inline_document(&crate::base_llm::ocr::handler::body_document(body)?) } } @@ -134,8 +134,7 @@ impl AzureAiOcrConfig { env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> Result, Error> { Self::resolve_api_base(connection.api_base.as_deref(), env_lookup)?; - if crate::custom_httpx::http_handler::has_header(&connection.extra_headers, "authorization") - { + if litellm_http::request::has_header(&connection.extra_headers, "authorization") { if config.azure_ad_token_provider.is_some() { super::common_utils::resolve_entra(config, env_lookup).await?; } diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/document.rs b/litellm-rust/crates/llms/src/base_llm/ocr/document.rs index 8737232a075..7ff88c6b843 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/document.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/document.rs @@ -1,18 +1,14 @@ use base64::{Engine, engine::general_purpose::STANDARD}; use data_url::{DataUrl, DataUrlError, forgiving_base64::DecodeError, mime::Mime}; +use litellm_http::{ + media::{DownloadPolicy, Error as MediaError, MediaFetcher}, + transport::Error as TransportError, +}; use reqwest::Url; -use crate::{ - base_llm::ocr::{ - error::Error, - transformation::{ - OCR_INLINE_MAX_BYTES, OCR_MAX_FETCH_REDIRECTS, OcrConnection, OcrDocument, - }, - }, - custom_httpx::{ - media::{DownloadPolicy, Error as MediaError, MediaFetcher}, - transport::Error as TransportError, - }, +use crate::base_llm::ocr::{ + error::Error, + transformation::{OCR_INLINE_MAX_BYTES, OCR_MAX_FETCH_REDIRECTS, OcrConnection, OcrDocument}, }; pub struct InlineDocument<'a>(DataUrl<'a>); @@ -196,10 +192,8 @@ mod tests { .redirect(reqwest::redirect::Policy::none()) .build() .unwrap(); - let client = crate::custom_httpx::llm_http_handler::OcrClient::for_test( - provider_http, - document_http, - ); + let client = + crate::base_llm::ocr::handler::OcrClient::for_test(provider_http, document_http); let converted = inline_remote_document( client.document_fetcher(), OcrDocument::ImageUrl { diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/error.rs b/litellm-rust/crates/llms/src/base_llm/ocr/error.rs index 3061a9fe2b2..9fce387beb5 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/error.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/error.rs @@ -95,11 +95,11 @@ pub enum Error { #[error(transparent)] Auth(#[from] litellm_auth::Error), #[error(transparent)] - Transport(#[from] crate::custom_httpx::transport::Error), + Transport(#[from] litellm_http::transport::Error), #[error(transparent)] Params(#[from] litellm_core_utils::params::Error), #[error(transparent)] - Headers(#[from] crate::custom_httpx::http_handler::HeaderError), + Headers(#[from] litellm_http::request::HeaderError), } impl From for Error { @@ -125,9 +125,7 @@ impl Error { pub fn http_status_code(&self) -> Option { match self { Self::Provider { status, .. } - | Self::Transport(crate::custom_httpx::transport::Error::Http { status, .. }) => { - Some(*status) - } + | Self::Transport(litellm_http::transport::Error::Http { status, .. }) => Some(*status), error if error.is_request() => Some(400), _ => None, } diff --git a/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs b/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs similarity index 94% rename from litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs rename to litellm-rust/crates/llms/src/base_llm/ocr/handler.rs index 58dc03eea2d..b6f266928b1 100644 --- a/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs @@ -2,22 +2,20 @@ use bytes::{Bytes, BytesMut}; use futures_util::future::BoxFuture; use litellm_auth_gcp::VertexAuth; use litellm_host::event::WireRequest; -use litellm_http::{ClientVariant, HttpClientConfig, HttpClientPool}; +use litellm_http::{ + ClientVariant, HttpClientConfig, HttpClientPool, + media::{MediaFetcher, UrlPolicy}, + request::{HeaderPolicy, execute_http_request, with_headers}, + transport, +}; use serde::{Serialize, de::DeserializeOwned}; use serde_json::Value; -use crate::{ - base_llm::ocr::{ - error::Error, - transformation::{ - BaseOcrConfig, DecodedOcrResponse, LiteLLMOcrResponse, OcrDocument, OcrResponseContext, - PreparedOcrRequest, decode_request_value, decode_response, - }, - }, - custom_httpx::{ - http_handler::{HeaderPolicy, execute_http_request, with_headers}, - media::{MediaFetcher, UrlPolicy}, - transport, +use crate::base_llm::ocr::{ + error::Error, + transformation::{ + BaseOcrConfig, DecodedOcrResponse, LiteLLMOcrResponse, OcrDocument, OcrResponseContext, + PreparedOcrRequest, decode_request_value, decode_response, }, }; diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/mod.rs b/litellm-rust/crates/llms/src/base_llm/ocr/mod.rs index 7194efbb203..1231633431e 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/mod.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/mod.rs @@ -1,3 +1,4 @@ pub mod document; pub mod error; +pub mod handler; pub mod transformation; diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs index f215546849d..be4551709a1 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs @@ -12,11 +12,9 @@ use serde::{ use serde_json::{Map, Value}; use serde_with::serde_as; -use crate::{ - base_llm::ocr::error::Error, - custom_httpx::llm_http_handler::{ - CallHooks, OcrClient, read_response_bytes, transform_request_body, - }, +use crate::base_llm::ocr::{ + error::Error, + handler::{CallHooks, OcrClient, read_response_bytes, transform_request_body}, }; pub const OCR_RESPONSE_MAX_BYTES: usize = 64 * 1024 * 1024; diff --git a/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs b/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs index 2528c967f41..da6cf90ffcf 100644 --- a/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs @@ -7,17 +7,15 @@ use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; use serde_with::serde_as; -use crate::{ - base_llm::ocr::{ - document::InlineDocument, - error::Error, - transformation::{ - BaseOcrConfig, LiteLLMOcrResponse, OCR_INLINE_MAX_BYTES, OcrConnection, OcrDocument, - OcrPage, OcrPageImage, OcrResponseFormat, OcrUsageInfo, PreparedOcrRequest, - credential_env, decode_and_normalize_response, decode_response_value, - }, +use crate::base_llm::ocr::{ + document::InlineDocument, + error::Error, + handler::OcrClient, + transformation::{ + BaseOcrConfig, LiteLLMOcrResponse, OCR_INLINE_MAX_BYTES, OcrConnection, OcrDocument, + OcrPage, OcrPageImage, OcrResponseFormat, OcrUsageInfo, PreparedOcrRequest, credential_env, + decode_and_normalize_response, decode_response_value, }, - custom_httpx::llm_http_handler::OcrClient, }; const COHERE_PARSE_API_BASE: &str = "https://api.cohere.com"; @@ -163,7 +161,7 @@ impl BaseOcrConfig for CohereParseConfig { } fn validate_request_body(&self, body: &Value) -> Result<(), Error> { - validate_document(&crate::custom_httpx::llm_http_handler::body_document(body)?) + validate_document(&crate::base_llm::ocr::handler::body_document(body)?) } } @@ -173,8 +171,7 @@ impl CohereParseConfig { connection: &OcrConnection, env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> Result, Error> { - if crate::custom_httpx::http_handler::has_header(&connection.extra_headers, "authorization") - { + if litellm_http::request::has_header(&connection.extra_headers, "authorization") { return Ok(connection.extra_headers.clone()); } let key = connection diff --git a/litellm-rust/crates/llms/src/custom_httpx/mod.rs b/litellm-rust/crates/llms/src/custom_httpx/mod.rs deleted file mode 100644 index 057cb796c09..00000000000 --- a/litellm-rust/crates/llms/src/custom_httpx/mod.rs +++ /dev/null @@ -1,4 +0,0 @@ -pub mod http_handler; -pub mod llm_http_handler; -pub mod media; -pub mod transport; diff --git a/litellm-rust/crates/llms/src/lib.rs b/litellm-rust/crates/llms/src/lib.rs index 884fa739992..8d1bb366ed4 100644 --- a/litellm-rust/crates/llms/src/lib.rs +++ b/litellm-rust/crates/llms/src/lib.rs @@ -3,7 +3,6 @@ pub mod azure_ai; pub mod base_llm; pub mod bedrock; pub mod cohere; -pub mod custom_httpx; pub mod mistral; pub mod openai; pub mod reducto; diff --git a/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs b/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs index 9028f09c5ab..95658837fc3 100644 --- a/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs @@ -2,16 +2,13 @@ use litellm_core_utils::{call_arguments::CallArguments, params::OpaqueParams, ur use serde::{Deserialize, Serialize}; use serde_json::Value; -use crate::{ - base_llm::ocr::{ - error::Error, - transformation::{ - BaseOcrConfig, LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrPage, - OcrResponseFormat, OcrUsageInfo, PreparedOcrRequest, credential_env, - decode_and_normalize_response, - }, +use crate::base_llm::ocr::{ + error::Error, + handler::OcrClient, + transformation::{ + BaseOcrConfig, LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrPage, OcrResponseFormat, + OcrUsageInfo, PreparedOcrRequest, credential_env, decode_and_normalize_response, }, - custom_httpx::llm_http_handler::OcrClient, }; const MISTRAL_OCR_API_BASE: &str = "https://api.mistral.ai/v1"; @@ -129,8 +126,7 @@ impl MistralOcrConfig { connection: &OcrConnection, env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> Result, Error> { - if crate::custom_httpx::http_handler::has_header(&connection.extra_headers, "authorization") - { + if litellm_http::request::has_header(&connection.extra_headers, "authorization") { return Ok(connection.extra_headers.clone()); } let api_key = connection diff --git a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs index ec876fafb8f..740f0ced090 100644 --- a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs @@ -8,18 +8,14 @@ use litellm_core_utils::{ use serde::{Deserialize, Deserializer, Serialize}; use serde_json::{Map, Value, json}; -use crate::{ - base_llm::ocr::{ - document::InlineDocument, - error::Error, - transformation::{ - BaseOcrConfig, LiteLLMOcrResponse, OCR_INLINE_MAX_BYTES, OcrConnection, OcrDocument, - OcrPage, OcrRequestContext, OcrResponseFormat, OcrUsageInfo, PreparedOcrRequest, - credential_env, decode_and_normalize_response, - }, - }, - custom_httpx::llm_http_handler::{ - CallHooks, OcrClient, build_http_request, guardrail_document, +use crate::base_llm::ocr::{ + document::InlineDocument, + error::Error, + handler::{CallHooks, OcrClient, build_http_request, guardrail_document}, + transformation::{ + BaseOcrConfig, LiteLLMOcrResponse, OCR_INLINE_MAX_BYTES, OcrConnection, OcrDocument, + OcrPage, OcrRequestContext, OcrResponseFormat, OcrUsageInfo, PreparedOcrRequest, + credential_env, decode_and_normalize_response, }, }; @@ -437,7 +433,7 @@ fn resolve_headers( connection: &OcrConnection, env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> Result, Error> { - if crate::custom_httpx::http_handler::has_header(&connection.extra_headers, "authorization") { + if litellm_http::request::has_header(&connection.extra_headers, "authorization") { return Ok(connection.extra_headers.clone()); } let api_key = connection @@ -515,25 +511,21 @@ async fn upload_bytes_async( )?) .multipart(reqwest::multipart::Form::new().part("file", part)) .timeout(connection.timeout); - let builder = crate::custom_httpx::http_handler::with_headers( + let builder = litellm_http::request::with_headers( builder, headers, - crate::custom_httpx::http_handler::HeaderPolicy::Except(&[ - "content-type", - "content-length", - ]), + litellm_http::request::HeaderPolicy::Except(&["content-type", "content-length"]), ); - let response = crate::custom_httpx::http_handler::http_request(builder) + let response = litellm_http::request::http_request(builder) .await - .map_err(crate::custom_httpx::transport::Error::from)?; - let uploaded = - crate::custom_httpx::llm_http_handler::read_json_response::( - response, - false, - connection.max_response_bytes, - ) - .await? - .data; + .map_err(litellm_http::transport::Error::from)?; + let uploaded = crate::base_llm::ocr::handler::read_json_response::( + response, + false, + connection.max_response_bytes, + ) + .await? + .data; let file_id = uploaded .file_id .as_deref() diff --git a/litellm-rust/crates/llms/src/vertex_ai/ocr/deepseek_transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/ocr/deepseek_transformation.rs index 588b5243004..8009a65ff77 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/ocr/deepseek_transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/ocr/deepseek_transformation.rs @@ -4,16 +4,14 @@ use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; use super::transformation::VertexAiOcrConfig; -use crate::{ - base_llm::ocr::{ - error::Error, - transformation::{ - BaseOcrConfig, LiteLLMOcrResponse, OcrDocument, OcrPage, OcrPageDimensions, - OcrPageImage, OcrRequestContext, OcrResponseFormat, OcrUsageInfo, PreparedOcrRequest, - credential_env, decode_and_normalize_response, decode_response_value, - }, +use crate::base_llm::ocr::{ + error::Error, + handler::OcrClient, + transformation::{ + BaseOcrConfig, LiteLLMOcrResponse, OcrDocument, OcrPage, OcrPageDimensions, OcrPageImage, + OcrRequestContext, OcrResponseFormat, OcrUsageInfo, PreparedOcrRequest, credential_env, + decode_and_normalize_response, decode_response_value, }, - custom_httpx::llm_http_handler::OcrClient, }; const DEFAULT_API_BASE: &str = "https://aiplatform.googleapis.com"; diff --git a/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs index c2cb23d0010..a50e8261aa3 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs @@ -7,12 +7,12 @@ use crate::{ base_llm::ocr::{ document::{inline_remote_document, validate_inline_document}, error::Error, + handler::OcrClient, transformation::{ BaseOcrConfig, LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrEnvironment, OcrRequestContext, OcrResponseFormat, PreparedOcrRequest, credential_env, }, }, - custom_httpx::llm_http_handler::OcrClient, mistral::ocr::transformation::{MistralOcrConfig, MistralOcrRequest}, }; @@ -112,7 +112,7 @@ impl BaseOcrConfig for VertexAiOcrConfig { } fn validate_request_body(&self, body: &Value) -> Result<(), Error> { - validate_inline_document(&crate::custom_httpx::llm_http_handler::body_document(body)?) + validate_inline_document(&crate::base_llm::ocr::handler::body_document(body)?) } } diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index 61c5947ed9e..19d28f76b6f 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -1,7 +1,6 @@ use litellm_core::{Error, audio_transcription, chat_completions, messages, responses}; -use litellm_llms::{ - base_llm::ocr::error::Error as OcrError, custom_httpx::transport::Error as TransportError, -}; +use litellm_http::transport::Error as TransportError; +use litellm_llms::base_llm::ocr::error::Error as OcrError; use pyo3::{ exceptions::{PyRuntimeError, PyValueError}, prelude::*, diff --git a/litellm-rust/crates/python-bridge/src/http.rs b/litellm-rust/crates/python-bridge/src/http.rs index 1fc3e4a60f1..7e9a5f093b4 100644 --- a/litellm-rust/crates/python-bridge/src/http.rs +++ b/litellm-rust/crates/python-bridge/src/http.rs @@ -8,8 +8,8 @@ use litellm_core_utils::settings::ProcessEnvironment; use litellm_http::{ HttpClientConfig, HttpClientPool, HttpSettings, HttpSettingsLayer, Resolution, SslVerify, Unsupported, + media::{PublicDnsResolver, UrlPolicy}, }; -use litellm_llms::custom_httpx::media::{PublicDnsResolver, UrlPolicy}; use pyo3::{prelude::*, types::PyDict}; use crate::{errors::RustBridgeDeclined, python_settings::PythonSettings}; diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs index c1b3f59df58..1a9b170f661 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -4,7 +4,7 @@ use litellm_core::messages::{ route::{Messages, MessagesCall, MessagesOp, MessagesOpResult, MessagesOutput}, }; use litellm_host_python::{InvokeError, RouteHost, from_py, lookup, to_py}; -use litellm_llms::custom_httpx::transport::Error as TransportError; +use litellm_http::transport::Error as TransportError; use pyo3::{ exceptions::{PyException, PyValueError}, gc::{PyTraverseError, PyVisit}, diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs index 0ae56efbf02..b0a6acdebfd 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs @@ -15,10 +15,9 @@ pub(super) fn to_pyerr(error: Error) -> PyErr { body, headers, } => upstream_error(py, status, body, headers)?, - Error::Transport(litellm_llms::custom_httpx::transport::Error::Http { - status, - body, - }) => upstream_error(py, status, body, Vec::new())?, + Error::Transport(litellm_http::transport::Error::Http { status, body }) => { + upstream_error(py, status, body, Vec::new())? + } Error::RequestFormat => { let error = core_error_to_pyerr(Error::RequestFormat.into()); error diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs index f9d7024c824..190f37d075d 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs @@ -9,7 +9,7 @@ use host::OcrRouteHost; use litellm_auth_gcp::VertexAuth; use litellm_callbacks_legacy::{LegacySurface, PublicCall, run_legacy_call}; use litellm_core::ocr::route::ocr_machine; -use litellm_llms::custom_httpx::llm_http_handler::OcrClient; +use litellm_llms::base_llm::ocr::handler::OcrClient; use pyo3::{ prelude::*, types::{PyDict, PyTuple}, From c0705f31b4b1846647f4305430ab666f33ed1d5a Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 20:51:55 -0700 Subject: [PATCH 04/23] fix(rust): read OCR env-backed constants instead of hardcoding their defaults Native OCR hardcoded the default of five Python constants that come from env vars, so an operator setting them saw no effect: REQUEST_TIMEOUT (Rust used 600s, Python 6000s), MAX_IMAGE_URL_DOWNLOAD_SIZE_MB (0 disables document downloads), AZURE_OPERATION_POLLING_TIMEOUT, AZURE_DOCUMENT_INTELLIGENCE_API_VERSION and AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI. OcrSettings reads them through Lookup with Python's parsing, the bridge builds it per call and OcrClient carries it into the connection. A zero per-call timeout now falls back to REQUEST_TIMEOUT, matching `timeout or request_timeout`. Co-Authored-By: Claude Opus 5 --- litellm-rust/crates/core/src/ocr/handler.rs | 2 +- litellm-rust/crates/core/src/ocr/prepare.rs | 10 +- litellm-rust/crates/core/src/ocr/types.rs | 2 +- .../tests/azure_document_intelligence_ocr.rs | 60 +++++++- litellm-rust/crates/core/tests/ocr.rs | 2 + .../document_intelligence/transformation.rs | 58 +++++--- .../crates/llms/src/base_llm/ocr/document.rs | 2 +- .../crates/llms/src/base_llm/ocr/handler.rs | 14 ++ .../crates/llms/src/base_llm/ocr/mod.rs | 1 + .../crates/llms/src/base_llm/ocr/settings.rs | 137 ++++++++++++++++++ .../llms/src/base_llm/ocr/transformation.rs | 57 ++++++-- .../python-bridge/src/routes/ocr/mod.rs | 4 +- .../python-bridge/src/routes/ocr/project.rs | 2 +- 13 files changed, 297 insertions(+), 54 deletions(-) create mode 100644 litellm-rust/crates/llms/src/base_llm/ocr/settings.rs diff --git a/litellm-rust/crates/core/src/ocr/handler.rs b/litellm-rust/crates/core/src/ocr/handler.rs index f49976de043..126e79e20e7 100644 --- a/litellm-rust/crates/core/src/ocr/handler.rs +++ b/litellm-rust/crates/core/src/ocr/handler.rs @@ -22,7 +22,7 @@ pub(crate) async fn perform_ocr_request( ) -> Result { request.response_format()?; let config = request.config; - let request = prepare_request(request, caller_document); + let request = prepare_request(request, caller_document, client.settings()); let hooks = OcrCallHooks::new(host.clone(), &request, config); config.ocr(client, &request, &hooks).await } diff --git a/litellm-rust/crates/core/src/ocr/prepare.rs b/litellm-rust/crates/core/src/ocr/prepare.rs index 8ac038290b7..72c35469f6d 100644 --- a/litellm-rust/crates/core/src/ocr/prepare.rs +++ b/litellm-rust/crates/core/src/ocr/prepare.rs @@ -1,6 +1,7 @@ use litellm_auth::{InputSource, SecretValue, Sourced}; -use litellm_llms::base_llm::ocr::transformation::{ - OcrConnection, OcrCredentialInputs, PreparedOcrRequest, credential_env, +use litellm_llms::base_llm::ocr::{ + settings::OcrSettings, + transformation::{OcrConnection, OcrCredentialInputs, PreparedOcrRequest, credential_env}, }; use super::provider_config::OcrProvider; @@ -9,6 +10,7 @@ use crate::ocr::types::{LiteLLMOcrRequest, ResolvedOcrRequest}; pub(crate) fn prepare_request( request: ResolvedOcrRequest, caller_document: bool, + settings: &OcrSettings, ) -> PreparedOcrRequest { let credentials = request.credentials.clone(); let api_base_env = match request.config.provider() { @@ -51,7 +53,7 @@ pub(crate) fn prepare_request( PreparedOcrRequest { model, document, - connection: OcrConnection::new(resolved, transport), + connection: OcrConnection::new(resolved, transport, settings.clone()), caller_document, optional_params, input_sources, @@ -61,7 +63,7 @@ pub(crate) fn prepare_request( #[cfg(test)] pub(crate) fn prepare_request_for_test(request: ResolvedOcrRequest) -> PreparedOcrRequest { - prepare_request(request, true) + prepare_request(request, true, &OcrSettings::default()) } #[cfg(test)] diff --git a/litellm-rust/crates/core/src/ocr/types.rs b/litellm-rust/crates/core/src/ocr/types.rs index 6316088dec8..59c9cec8da9 100644 --- a/litellm-rust/crates/core/src/ocr/types.rs +++ b/litellm-rust/crates/core/src/ocr/types.rs @@ -277,7 +277,7 @@ mod tests { vec![("x-a".to_string(), "1".to_string())] ); assert_eq!(request.transport.extra_headers_source, InputSource::Request); - assert_eq!(request.transport.timeout, Duration::from_secs(7)); + assert_eq!(request.transport.timeout, Some(Duration::from_secs(7))); assert_eq!(request.input_sources.len(), 2); let defaulted = LiteLLMOcrRequest::from_inputs( diff --git a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs index 3cbe6fe3159..6dc9bfa5e7e 100644 --- a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs +++ b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs @@ -1,10 +1,12 @@ use litellm_host::event::{CallEvent, MachineEvent}; -use litellm_llms::base_llm::ocr::error::Error; +use litellm_llms::base_llm::ocr::{error::Error, settings::OcrSettings}; use rstest::rstest; use serde_json::{Value, json}; use super::{ - test_support::{MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request}, + test_support::{ + MockResponse, mock_server, ocr_client, perform_ocr, perform_ocr_with, wire_request, + }, wire::{OcrWireRequest, decode_request}, }; use crate::ocr::route::LocalOcrHost; @@ -200,6 +202,42 @@ async fn immediate_response_normalizes_pages_and_preserves_native() { ); } +#[tokio::test] +async fn client_settings_choose_the_api_version_and_the_inch_to_pixel_dpi() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "status":"succeeded", + "analyzeResult":{"pages":[{"pageNumber":1,"width":8.5,"height":11,"unit":"inch"}]} + }))]) + .await; + let client = ocr_client().with_settings(OcrSettings { + document_intelligence_api_version: "2099-01-01".into(), + document_intelligence_dpi: 72, + ..OcrSettings::default() + }); + + let result = crate::ocr::client::perform( + &client, + wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})), + ) + .await + .unwrap(); + server.await.unwrap(); + + let target = seen.lock().unwrap()[0] + .split_whitespace() + .nth(1) + .unwrap() + .to_string(); + assert_eq!( + query_value(&format!("{base}{target}"), "api-version").as_deref(), + Some("2099-01-01") + ); + assert_eq!( + serde_json::to_value(&result.pages[0].dimensions).unwrap(), + json!({"width":612,"height":792,"dpi":72}) + ); +} + #[tokio::test] async fn accepted_response_polls_to_success_with_only_credentials() { let operation = json!({"status":"succeeded","analyzeResult":{"pages":[]}}); @@ -425,13 +463,19 @@ async fn polling_deadline_bounds_retry_delay() { }, ]) .await; - let mut request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})); - request.transport.poll_timeout = std::time::Duration::from_millis(100); + let request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})); + let client = ocr_client().with_settings(OcrSettings { + poll_timeout: std::time::Duration::from_millis(100), + ..OcrSettings::default() + }); - let error = tokio::time::timeout(std::time::Duration::from_secs(1), perform_ocr(request)) - .await - .unwrap() - .unwrap_err(); + let error = tokio::time::timeout( + std::time::Duration::from_secs(1), + crate::ocr::client::perform(&client, request), + ) + .await + .unwrap() + .unwrap_err(); server.await.unwrap(); assert!(error.to_string().contains("timed out")); } diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index b999c43de8b..f87f16cd033 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -13,6 +13,7 @@ use litellm_http::{ use litellm_llms::base_llm::ocr::{ error::Error as OcrError, handler::OcrClient, + settings::OcrSettings, transformation::{LiteLLMOcrResponse, OCR_RESPONSE_MAX_BYTES, OcrTransportConfig}, }; use rstest::rstest; @@ -185,6 +186,7 @@ async fn ocr_client_uses_the_injected_http_pool_configuration() { &Resolution::from(&settings).config, UrlPolicy::default(), VertexAuth::default(), + OcrSettings::default(), ) .unwrap(); crate::ocr::client::perform(&client, wire_request("mistral/model", &base, json!({}))) diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs index a347375510d..5fb20d5900a 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs @@ -18,6 +18,7 @@ use crate::base_llm::ocr::{ document::InlineDocument, error::Error, handler::{CallHooks, OcrClient, read_json_response}, + settings::OcrSettings, transformation::{ BaseOcrConfig, DecodedOcrResponse, LiteLLMOcrResponse, OCR_INLINE_MAX_BYTES, OCR_POLL_RETRY_SECS, OcrConnection, OcrCredentialInputs, OcrDocument, OcrPage, @@ -26,9 +27,7 @@ use crate::base_llm::ocr::{ }, }; -const AZURE_DI_API_VERSION: &str = "2024-11-30"; const AZURE_DI_SUBSCRIPTION_HEADER: &str = "Ocp-Apim-Subscription-Key"; -const AZURE_DI_DEFAULT_DPI: i64 = 96; const AZURE_DI_DEFAULT_WIDTH: f64 = 8.5; const AZURE_DI_DEFAULT_HEIGHT: f64 = 11.0; @@ -195,7 +194,15 @@ impl BaseOcrConfig for AzureDocumentIntelligenceOcrConfig { let endpoint = nonblank(request.connection.api_base.clone()) .or_else(|| nonblank(credential_env(AZURE_DI_ENDPOINT_ENV))) .ok_or_else(|| Error::Auth(litellm_auth::Error::ProviderAuthentication("Missing Azure Document Intelligence API Base - Set AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT or pass api_base".into())))?; - self.build_ocr_url(&endpoint, &request.model, optional_params) + self.build_ocr_url( + &endpoint, + &request.model, + optional_params, + &request + .connection + .settings + .document_intelligence_api_version, + ) } fn transform_ocr_request( @@ -214,12 +221,13 @@ impl BaseOcrConfig for AzureDocumentIntelligenceOcrConfig { raw_response: &[u8], request_format: OcrResponseFormat, ) -> Result { - decode_and_normalize_response( - model, - raw_response, - request_format, - transform_completed_response, - ) + decode_and_normalize_response(model, raw_response, request_format, |model, response| { + transform_completed_response( + model, + response, + OcrSettings::default().document_intelligence_dpi, + ) + }) } async fn async_transform_ocr_response( @@ -240,7 +248,11 @@ impl BaseOcrConfig for AzureDocumentIntelligenceOcrConfig { .await?; Ok(LiteLLMOcrResponse { provider_native_response: decoded.native, - ..transform_completed_response(model, decoded.data)? + ..transform_completed_response( + model, + decoded.data, + context.connection.settings.document_intelligence_dpi, + )? }) } } @@ -353,6 +365,7 @@ fn build_request(document: OcrDocument) -> Result Result { if response.status != Some(OperationStatus::Succeeded) { return Err(Error::OperationStatus( @@ -366,7 +379,7 @@ fn transform_completed_response( let pages = result .pages .into_iter() - .map(transform_azure_page) + .map(|page| transform_azure_page(page, dpi)) .collect::, _>>()?; let pages_processed = i64::try_from(pages.len()).map_err(|_| Error::NumericRange("pages"))?; Ok(LiteLLMOcrResponse { @@ -381,7 +394,7 @@ fn transform_completed_response( }) } -fn transform_azure_page(page: AzureDocumentIntelligencePage) -> Result { +fn transform_azure_page(page: AzureDocumentIntelligencePage, dpi: i64) -> Result { let index = page .page_number .unwrap_or(1) @@ -391,6 +404,7 @@ fn transform_azure_page(page: AzureDocumentIntelligencePage) -> Result Result Result { - let scale = if unit == "inch" { - AZURE_DI_DEFAULT_DPI as f64 - } else { - 1.0 - }; +fn convert_dimensions( + width: f64, + height: f64, + unit: &str, + dpi: i64, +) -> Result { + let scale = if unit == "inch" { dpi as f64 } else { 1.0 }; Ok(OcrPageDimensions { width: Some(pixel_dimension(width, scale, "page.width")?), height: Some(pixel_dimension(height, scale, "page.height")?), - dpi: Some(AZURE_DI_DEFAULT_DPI), + dpi: Some(dpi), }) } @@ -475,7 +490,7 @@ async fn poll_operation( hooks: &dyn CallHooks, ) -> Result, Error> { let deadline = Instant::now() - .checked_add(connection.poll_timeout) + .checked_add(connection.settings.poll_timeout) .ok_or(Error::PollTimeout)?; loop { @@ -544,13 +559,14 @@ impl AzureDocumentIntelligenceOcrConfig { endpoint: &str, model: &str, params: &DocumentIntelligenceParams, + api_version: &str, ) -> Result { let model = format!("{}:analyze", model_id(model)?); ApiUrl::parse(endpoint) .and_then(|url| url.complete_path(&["documentintelligence", "documentModels", &model])) .map(|url| { url.append_query_pairs( - [("api-version", AZURE_DI_API_VERSION)] + [("api-version", api_version)] .into_iter() .chain(params.pages.iter().map(|pages| ("pages", pages.as_str()))) .chain( diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/document.rs b/litellm-rust/crates/llms/src/base_llm/ocr/document.rs index 7ff88c6b843..724625b8208 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/document.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/document.rs @@ -68,7 +68,7 @@ pub async fn inline_remote_document( url, DownloadPolicy { timeout: connection.timeout, - max_bytes: connection.max_download_bytes, + max_bytes: connection.settings.max_download_bytes, max_redirects: OCR_MAX_FETCH_REDIRECTS, }, ) diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs b/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs index b6f266928b1..9410f673d29 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs @@ -13,6 +13,7 @@ use serde_json::Value; use crate::base_llm::ocr::{ error::Error, + settings::OcrSettings, transformation::{ BaseOcrConfig, DecodedOcrResponse, LiteLLMOcrResponse, OcrDocument, OcrResponseContext, PreparedOcrRequest, decode_request_value, decode_response, @@ -33,6 +34,7 @@ pub struct OcrClient { polling_http: reqwest::Client, document_fetcher: MediaFetcher, vertex_auth: VertexAuth, + settings: OcrSettings, } impl OcrClient { @@ -41,12 +43,14 @@ impl OcrClient { config: &HttpClientConfig, url_policy: UrlPolicy, vertex_auth: VertexAuth, + settings: OcrSettings, ) -> Result { Ok(Self { provider_http: pool.client(config, ClientVariant::Provider)?, polling_http: pool.client(config, ClientVariant::NoRedirect)?, document_fetcher: MediaFetcher::new(pool, config, url_policy)?, vertex_auth, + settings, }) } @@ -66,6 +70,10 @@ impl OcrClient { &self.vertex_auth } + pub fn settings(&self) -> &OcrSettings { + &self.settings + } + #[cfg(any(test, feature = "test-support"))] pub fn for_test(provider_http: reqwest::Client, document_http: reqwest::Client) -> Self { Self { @@ -76,8 +84,14 @@ impl OcrClient { .expect("test polling client builds"), document_fetcher: MediaFetcher::for_test(document_http), vertex_auth: VertexAuth::default(), + settings: OcrSettings::default(), } } + + #[cfg(any(test, feature = "test-support"))] + pub fn with_settings(self, settings: OcrSettings) -> Self { + Self { settings, ..self } + } } /// Rust counterpart of `BaseLLMHTTPHandler.async_ocr`: prepare the provider request, diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/mod.rs b/litellm-rust/crates/llms/src/base_llm/ocr/mod.rs index 1231633431e..e81f71b253d 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/mod.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/mod.rs @@ -1,4 +1,5 @@ pub mod document; pub mod error; pub mod handler; +pub mod settings; pub mod transformation; diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/settings.rs b/litellm-rust/crates/llms/src/base_llm/ocr/settings.rs new file mode 100644 index 00000000000..239a5b22000 --- /dev/null +++ b/litellm-rust/crates/llms/src/base_llm/ocr/settings.rs @@ -0,0 +1,137 @@ +use std::time::Duration; + +use litellm_core_utils::settings::Lookup; + +#[derive(Clone, Debug, PartialEq)] +pub struct OcrSettings { + pub request_timeout: Duration, + pub max_download_bytes: u64, + pub poll_timeout: Duration, + pub document_intelligence_api_version: String, + pub document_intelligence_dpi: i64, +} + +impl Default for OcrSettings { + fn default() -> Self { + Self { + request_timeout: Duration::from_secs(6000), + max_download_bytes: megabytes(50.0), + poll_timeout: Duration::from_secs(120), + document_intelligence_api_version: "2024-11-30".into(), + document_intelligence_dpi: 96, + } + } +} + +impl OcrSettings { + pub fn from_environment(env: &impl Lookup) -> Self { + let defaults = Self::default(); + Self { + request_timeout: env + .parsed::("REQUEST_TIMEOUT") + .and_then(|seconds| Duration::try_from_secs_f64(seconds).ok()) + .unwrap_or(defaults.request_timeout), + max_download_bytes: env + .parsed::("MAX_IMAGE_URL_DOWNLOAD_SIZE_MB") + .filter(|size| size.is_finite()) + .map_or(defaults.max_download_bytes, megabytes), + poll_timeout: env + .parsed::("AZURE_OPERATION_POLLING_TIMEOUT") + .map_or(defaults.poll_timeout, |seconds| { + Duration::from_secs(seconds.max(0).unsigned_abs()) + }), + document_intelligence_api_version: env + .get("AZURE_DOCUMENT_INTELLIGENCE_API_VERSION") + .unwrap_or(defaults.document_intelligence_api_version), + document_intelligence_dpi: env + .parsed("AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI") + .unwrap_or(defaults.document_intelligence_dpi), + } + } +} + +fn megabytes(size: f64) -> u64 { + (size * 1024.0 * 1024.0) as u64 +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + fn env_of(values: &'static [(&'static str, &'static str)]) -> impl Fn(&str) -> Option { + move |name| { + values + .iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| value.to_string()) + } + } + + #[test] + fn an_empty_environment_keeps_the_python_defaults() { + assert_eq!( + OcrSettings::from_environment(&env_of(&[])), + OcrSettings::default() + ); + } + + #[test] + fn every_setting_follows_its_environment_variable() { + let settings = OcrSettings::from_environment(&env_of(&[ + ("REQUEST_TIMEOUT", "30.5"), + ("MAX_IMAGE_URL_DOWNLOAD_SIZE_MB", "0.5"), + ("AZURE_OPERATION_POLLING_TIMEOUT", " 600 "), + ("AZURE_DOCUMENT_INTELLIGENCE_API_VERSION", "2025-01-01"), + ("AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI", "72"), + ])); + assert_eq!( + settings, + OcrSettings { + request_timeout: Duration::from_millis(30_500), + max_download_bytes: 512 * 1024, + poll_timeout: Duration::from_secs(600), + document_intelligence_api_version: "2025-01-01".into(), + document_intelligence_dpi: 72, + } + ); + } + + #[rstest] + #[case::zero_disables_downloads("0", 0)] + #[case::negative_rejects_every_download("-1", 0)] + #[case::fraction_truncates_like_int("0.0000001", 0)] + #[case::unparsable_keeps_the_default("big", 50 * 1024 * 1024)] + fn download_size_converts_megabytes_like_python( + #[case] value: &'static str, + #[case] bytes: u64, + ) { + let env = + move |name: &str| (name == "MAX_IMAGE_URL_DOWNLOAD_SIZE_MB").then(|| value.to_string()); + assert_eq!( + OcrSettings::from_environment(&env).max_download_bytes, + bytes + ); + } + + #[test] + fn a_negative_polling_timeout_expires_immediately() { + let env = + |name: &str| (name == "AZURE_OPERATION_POLLING_TIMEOUT").then(|| "-5".to_string()); + assert_eq!( + OcrSettings::from_environment(&env).poll_timeout, + Duration::ZERO + ); + } + + #[test] + fn an_empty_api_version_is_sent_as_is_like_python_str_of_getenv() { + let env = + |name: &str| (name == "AZURE_DOCUMENT_INTELLIGENCE_API_VERSION").then(String::new); + assert_eq!( + OcrSettings::from_environment(&env).document_intelligence_api_version, + "" + ); + } +} diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs index be4551709a1..5d1a0c8e0ed 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs @@ -15,14 +15,12 @@ use serde_with::serde_as; use crate::base_llm::ocr::{ error::Error, handler::{CallHooks, OcrClient, read_response_bytes, transform_request_body}, + settings::OcrSettings, }; pub const OCR_RESPONSE_MAX_BYTES: usize = 64 * 1024 * 1024; -pub const OCR_HTTP_TIMEOUT_SECS: u64 = 600; pub const OCR_INLINE_MAX_BYTES: usize = 50 * 1024 * 1024; -pub const OCR_DOWNLOAD_MAX_BYTES: u64 = 50 * 1024 * 1024; pub const OCR_MAX_FETCH_REDIRECTS: usize = 10; -pub const OCR_POLL_TIMEOUT_SECS: u64 = 120; pub const OCR_POLL_RETRY_SECS: u64 = 2; #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] @@ -114,10 +112,8 @@ impl OcrCredentialInputs { pub struct OcrTransportConfig { pub extra_headers: Vec<(String, String)>, pub extra_headers_source: InputSource, - pub timeout: Duration, - pub max_download_bytes: u64, + pub timeout: Option, pub max_response_bytes: usize, - pub poll_timeout: Duration, } impl Default for OcrTransportConfig { @@ -125,10 +121,8 @@ impl Default for OcrTransportConfig { Self { extra_headers: Vec::new(), extra_headers_source: InputSource::Deployment, - timeout: Duration::from_secs(OCR_HTTP_TIMEOUT_SECS), - max_download_bytes: OCR_DOWNLOAD_MAX_BYTES, + timeout: None, max_response_bytes: OCR_RESPONSE_MAX_BYTES, - poll_timeout: Duration::from_secs(OCR_POLL_TIMEOUT_SECS), } } } @@ -143,7 +137,7 @@ impl OcrTransportConfig { Self { extra_headers, extra_headers_source, - timeout: timeout.unwrap_or(self.timeout), + timeout: timeout.or(self.timeout), ..self } } @@ -164,13 +158,16 @@ pub struct OcrConnection { pub extra_headers: Vec<(String, String)>, pub extra_headers_source: InputSource, pub timeout: Duration, - pub max_download_bytes: u64, pub max_response_bytes: usize, - pub poll_timeout: Duration, + pub settings: OcrSettings, } impl OcrConnection { - pub fn new(credentials: ResolvedOcrCredentials, transport: OcrTransportConfig) -> Self { + pub fn new( + credentials: ResolvedOcrCredentials, + transport: OcrTransportConfig, + settings: OcrSettings, + ) -> Self { let api_key_source = credentials .api_key .as_ref() @@ -188,10 +185,12 @@ impl OcrConnection { api_base_source, extra_headers: transport.extra_headers, extra_headers_source: transport.extra_headers_source, - timeout: transport.timeout, - max_download_bytes: transport.max_download_bytes, + timeout: transport + .timeout + .filter(|timeout| !timeout.is_zero()) + .unwrap_or(settings.request_timeout), max_response_bytes: transport.max_response_bytes, - poll_timeout: transport.poll_timeout, + settings, } } } @@ -201,6 +200,7 @@ impl Default for OcrConnection { Self::new( ResolvedOcrCredentials::default(), OcrTransportConfig::default(), + OcrSettings::default(), ) } } @@ -573,6 +573,31 @@ mod tests { use super::*; + #[test] + fn connection_timeout_falls_back_to_the_request_timeout_setting_like_a_python_or() { + let settings = OcrSettings { + request_timeout: Duration::from_secs(42), + ..OcrSettings::default() + }; + let timeout = |call: Option| { + OcrConnection::new( + ResolvedOcrCredentials::default(), + OcrTransportConfig { + timeout: call, + ..OcrTransportConfig::default() + }, + settings.clone(), + ) + .timeout + }; + assert_eq!(timeout(None), Duration::from_secs(42)); + assert_eq!(timeout(Some(Duration::ZERO)), Duration::from_secs(42)); + assert_eq!( + timeout(Some(Duration::from_secs(5))), + Duration::from_secs(5) + ); + } + #[test] fn normalized_response_rejects_invalid_shared_fields() { for fields in [ diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs index 190f37d075d..bb845f48783 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs @@ -9,7 +9,8 @@ use host::OcrRouteHost; use litellm_auth_gcp::VertexAuth; use litellm_callbacks_legacy::{LegacySurface, PublicCall, run_legacy_call}; use litellm_core::ocr::route::ocr_machine; -use litellm_llms::base_llm::ocr::handler::OcrClient; +use litellm_core_utils::settings::ProcessEnvironment; +use litellm_llms::base_llm::ocr::{handler::OcrClient, settings::OcrSettings}; use pyo3::{ prelude::*, types::{PyDict, PyTuple}, @@ -43,6 +44,7 @@ fn run_ocr( &config, http::url_policy(py)?, VERTEX_AUTH.clone(), + OcrSettings::from_environment(&ProcessEnvironment), ) .map_err(|error| RustBridgeDeclined::new_err(error.to_string()))?; run_legacy_call( diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs index 5dd2aa804b8..697b935a1d4 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs @@ -592,7 +592,7 @@ kwargs = { ); assert_eq!( projected.transport.timeout, - std::time::Duration::from_secs(5) + Some(std::time::Duration::from_secs(5)) ); }); } From 0d76359dc9a4e1dba45020626f143e1f1f294bff Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 20:57:24 -0700 Subject: [PATCH 05/23] fix(rust): resolve OCR provider env fallbacks through the secret manager Python reads every provider credential fallback (MISTRAL_API_KEY, AZURE_AI_API_KEY, AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT, Azure AD and Vertex env, ...) through get_secret_str, which consults the configured key_management_system before os.environ. Native OCR read std::env directly, so a key held only in the vault went missing and a stale env copy silently won. OcrClient now carries an injected secret Lookup that the connection exposes to providers and auth crates; the bridge backs it with settings.secret -> get_secret_str, pure Rust keeps the process env. Co-Authored-By: Claude Opus 5 --- litellm-rust/crates/core/src/ocr/handler.rs | 2 +- litellm-rust/crates/core/src/ocr/prepare.rs | 23 +++++-- litellm-rust/crates/core/tests/ocr.rs | 25 +++++++ .../ocr/cohere_parse_transformation.rs | 2 +- .../document_intelligence/transformation.rs | 10 +-- .../llms/src/azure_ai/ocr/transformation.rs | 12 ++-- .../crates/llms/src/base_llm/ocr/handler.rs | 15 ++++- .../crates/llms/src/base_llm/ocr/settings.rs | 4 +- .../llms/src/base_llm/ocr/transformation.rs | 18 +++-- .../llms/src/cohere/ocr/transformation.rs | 6 +- .../llms/src/mistral/ocr/transformation.rs | 6 +- .../llms/src/reducto/ocr/transformation.rs | 6 +- .../vertex_ai/ocr/deepseek_transformation.rs | 7 +- .../llms/src/vertex_ai/ocr/transformation.rs | 9 +-- .../python-bridge/src/python_settings.rs | 65 ++++++++++++++++++- .../python-bridge/src/routes/ocr/mod.rs | 5 +- litellm/rust_bridge/settings.py | 6 ++ .../test_litellm/rust_bridge/test_settings.py | 46 +++++++++++++ 18 files changed, 226 insertions(+), 41 deletions(-) diff --git a/litellm-rust/crates/core/src/ocr/handler.rs b/litellm-rust/crates/core/src/ocr/handler.rs index 126e79e20e7..19037e49033 100644 --- a/litellm-rust/crates/core/src/ocr/handler.rs +++ b/litellm-rust/crates/core/src/ocr/handler.rs @@ -22,7 +22,7 @@ pub(crate) async fn perform_ocr_request( ) -> Result { request.response_format()?; let config = request.config; - let request = prepare_request(request, caller_document, client.settings()); + let request = prepare_request(request, caller_document, client); let hooks = OcrCallHooks::new(host.clone(), &request, config); config.ocr(client, &request, &hooks).await } diff --git a/litellm-rust/crates/core/src/ocr/prepare.rs b/litellm-rust/crates/core/src/ocr/prepare.rs index 72c35469f6d..ed8c7fba503 100644 --- a/litellm-rust/crates/core/src/ocr/prepare.rs +++ b/litellm-rust/crates/core/src/ocr/prepare.rs @@ -1,7 +1,7 @@ use litellm_auth::{InputSource, SecretValue, Sourced}; use litellm_llms::base_llm::ocr::{ - settings::OcrSettings, - transformation::{OcrConnection, OcrCredentialInputs, PreparedOcrRequest, credential_env}, + handler::OcrClient, + transformation::{OcrConnection, OcrCredentialInputs, PreparedOcrRequest}, }; use super::provider_config::OcrProvider; @@ -10,7 +10,7 @@ use crate::ocr::types::{LiteLLMOcrRequest, ResolvedOcrRequest}; pub(crate) fn prepare_request( request: ResolvedOcrRequest, caller_document: bool, - settings: &OcrSettings, + client: &OcrClient, ) -> PreparedOcrRequest { let credentials = request.credentials.clone(); let api_base_env = match request.config.provider() { @@ -23,14 +23,14 @@ pub(crate) fn prepare_request( request .config .get_api_key_env_var() - .and_then(credential_env) + .and_then(|name| client.secrets().get(name)) .map(|value| Sourced::new(SecretValue::new(value), InputSource::Environment)) }) }); let dynamic_api_base = credentials.dynamic_api_base.or_else(|| { credentials.api_base.clone().or_else(|| { api_base_env - .and_then(credential_env) + .and_then(|name| client.secrets().get(name)) .map(|value| Sourced::new(value, InputSource::Environment)) }) }); @@ -53,7 +53,12 @@ pub(crate) fn prepare_request( PreparedOcrRequest { model, document, - connection: OcrConnection::new(resolved, transport, settings.clone()), + connection: OcrConnection::new( + resolved, + transport, + client.settings().clone(), + client.secrets().clone(), + ), caller_document, optional_params, input_sources, @@ -63,7 +68,11 @@ pub(crate) fn prepare_request( #[cfg(test)] pub(crate) fn prepare_request_for_test(request: ResolvedOcrRequest) -> PreparedOcrRequest { - prepare_request(request, true, &OcrSettings::default()) + prepare_request( + request, + true, + &OcrClient::for_test(reqwest::Client::new(), reqwest::Client::new()), + ) } #[cfg(test)] diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index f87f16cd033..61d59a38065 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -174,6 +174,30 @@ async fn facade_retains_native_response_when_requested() { ); } +#[tokio::test] +async fn provider_key_fallback_reads_the_injected_secret_source() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let request = decode_request(OcrWireRequest { + model: "mistral/model".into(), + document: json!({"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}), + api_key: None, + api_base: Some(base.clone()), + custom_llm_provider: None, + extra_headers: None, + optional_params: Default::default(), + input_sources: Default::default(), + timeout_seconds: Some(2.0), + }) + .unwrap(); + let client = ocr_client().with_secrets(Arc::new(|name: &str| { + (name == "MISTRAL_API_KEY").then(|| "from-secret-manager".to_string()) + })); + + crate::ocr::client::perform(&client, request).await.unwrap(); + server.await.unwrap(); + assert!(seen.lock().unwrap()[0].contains("authorization: Bearer from-secret-manager")); +} + #[tokio::test] async fn ocr_client_uses_the_injected_http_pool_configuration() { let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; @@ -187,6 +211,7 @@ async fn ocr_client_uses_the_injected_http_pool_configuration() { UrlPolicy::default(), VertexAuth::default(), OcrSettings::default(), + Arc::new(litellm_core_utils::settings::ProcessEnvironment), ) .unwrap(); crate::ocr::client::perform(&client, wire_request("mistral/model", &base, json!({}))) diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/cohere_parse_transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/cohere_parse_transformation.rs index 86ee0d96895..045d8744bc9 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/cohere_parse_transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/cohere_parse_transformation.rs @@ -53,7 +53,7 @@ impl BaseOcrConfig for AzureAICohereParseConfig { ) -> Result { let base = super::transformation::AzureAiOcrConfig::resolve_api_base( request.connection.api_base.as_deref(), - &crate::base_llm::ocr::transformation::credential_env, + &|name: &str| request.connection.secret(name), )?; self.get_complete_url(&base) } diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs index 5fb20d5900a..8e6182f454f 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs @@ -23,7 +23,7 @@ use crate::base_llm::ocr::{ BaseOcrConfig, DecodedOcrResponse, LiteLLMOcrResponse, OCR_INLINE_MAX_BYTES, OCR_POLL_RETRY_SECS, OcrConnection, OcrCredentialInputs, OcrDocument, OcrPage, OcrPageDimensions, OcrResponseContext, OcrResponseFormat, OcrUsageInfo, PreparedOcrRequest, - ResolvedOcrCredentials, credential_env, decode_and_normalize_response, decode_response, + ResolvedOcrCredentials, decode_and_normalize_response, decode_response, }, }; @@ -181,8 +181,10 @@ impl BaseOcrConfig for AzureDocumentIntelligenceOcrConfig { &request.input_sources, )? }; - self.resolve_headers(&request.connection, &config, &credential_env) - .await + self.resolve_headers(&request.connection, &config, &|name: &str| { + request.connection.secret(name) + }) + .await } fn get_complete_url( @@ -192,7 +194,7 @@ impl BaseOcrConfig for AzureDocumentIntelligenceOcrConfig { _environment: &Self::Environment, ) -> Result { let endpoint = nonblank(request.connection.api_base.clone()) - .or_else(|| nonblank(credential_env(AZURE_DI_ENDPOINT_ENV))) + .or_else(|| nonblank(request.connection.secret(AZURE_DI_ENDPOINT_ENV))) .ok_or_else(|| Error::Auth(litellm_auth::Error::ProviderAuthentication("Missing Azure Document Intelligence API Base - Set AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT or pass api_base".into())))?; self.build_ocr_url( &endpoint, diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs index 4a04910aa9a..cd20e75df85 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs @@ -10,7 +10,7 @@ use crate::{ handler::OcrClient, transformation::{ BaseOcrConfig, LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrRequestContext, - OcrResponseFormat, PreparedOcrRequest, credential_env, + OcrResponseFormat, PreparedOcrRequest, }, }, mistral::ocr::transformation::{MistralOcrConfig, MistralOcrRequest}, @@ -57,8 +57,10 @@ impl BaseOcrConfig for AzureAiOcrConfig { &request.input_sources, )? }; - self.resolve_headers(&request.connection, &config, &credential_env) - .await + self.resolve_headers(&request.connection, &config, &|name: &str| { + request.connection.secret(name) + }) + .await } fn get_complete_url( @@ -67,7 +69,9 @@ impl BaseOcrConfig for AzureAiOcrConfig { _optional_params: &Self::OcrParams, _environment: &Self::Environment, ) -> Result { - self.build_ocr_url(request.connection.api_base.as_deref(), &credential_env) + self.build_ocr_url(request.connection.api_base.as_deref(), &|name: &str| { + request.connection.secret(name) + }) } fn transform_ocr_request( diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs b/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs index 9410f673d29..91fb6461770 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/handler.rs @@ -13,7 +13,7 @@ use serde_json::Value; use crate::base_llm::ocr::{ error::Error, - settings::OcrSettings, + settings::{OcrSettings, Secrets}, transformation::{ BaseOcrConfig, DecodedOcrResponse, LiteLLMOcrResponse, OcrDocument, OcrResponseContext, PreparedOcrRequest, decode_request_value, decode_response, @@ -35,6 +35,7 @@ pub struct OcrClient { document_fetcher: MediaFetcher, vertex_auth: VertexAuth, settings: OcrSettings, + secrets: Secrets, } impl OcrClient { @@ -44,6 +45,7 @@ impl OcrClient { url_policy: UrlPolicy, vertex_auth: VertexAuth, settings: OcrSettings, + secrets: Secrets, ) -> Result { Ok(Self { provider_http: pool.client(config, ClientVariant::Provider)?, @@ -51,6 +53,7 @@ impl OcrClient { document_fetcher: MediaFetcher::new(pool, config, url_policy)?, vertex_auth, settings, + secrets, }) } @@ -74,6 +77,10 @@ impl OcrClient { &self.settings } + pub fn secrets(&self) -> &Secrets { + &self.secrets + } + #[cfg(any(test, feature = "test-support"))] pub fn for_test(provider_http: reqwest::Client, document_http: reqwest::Client) -> Self { Self { @@ -85,6 +92,7 @@ impl OcrClient { document_fetcher: MediaFetcher::for_test(document_http), vertex_auth: VertexAuth::default(), settings: OcrSettings::default(), + secrets: std::sync::Arc::new(litellm_core_utils::settings::ProcessEnvironment), } } @@ -92,6 +100,11 @@ impl OcrClient { pub fn with_settings(self, settings: OcrSettings) -> Self { Self { settings, ..self } } + + #[cfg(any(test, feature = "test-support"))] + pub fn with_secrets(self, secrets: Secrets) -> Self { + Self { secrets, ..self } + } } /// Rust counterpart of `BaseLLMHTTPHandler.async_ocr`: prepare the provider request, diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/settings.rs b/litellm-rust/crates/llms/src/base_llm/ocr/settings.rs index 239a5b22000..276b2ca1311 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/settings.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/settings.rs @@ -1,7 +1,9 @@ -use std::time::Duration; +use std::{sync::Arc, time::Duration}; use litellm_core_utils::settings::Lookup; +pub type Secrets = Arc; + #[derive(Clone, Debug, PartialEq)] pub struct OcrSettings { pub request_timeout: Duration, diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs index 5d1a0c8e0ed..3960282b580 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs @@ -1,9 +1,10 @@ -use std::{collections::BTreeMap, future::Future, time::Duration}; +use std::{collections::BTreeMap, future::Future, sync::Arc, time::Duration}; use litellm_auth::{InputSource, SecretValue, Sourced, TokenProviderHandle}; use litellm_core_utils::{ call_arguments::CallArguments, serde_compat::{FiniteF64, LaxI64}, + settings::ProcessEnvironment, }; use serde::{ Deserialize, Serialize, @@ -15,7 +16,7 @@ use serde_with::serde_as; use crate::base_llm::ocr::{ error::Error, handler::{CallHooks, OcrClient, read_response_bytes, transform_request_body}, - settings::OcrSettings, + settings::{OcrSettings, Secrets}, }; pub const OCR_RESPONSE_MAX_BYTES: usize = 64 * 1024 * 1024; @@ -160,6 +161,7 @@ pub struct OcrConnection { pub timeout: Duration, pub max_response_bytes: usize, pub settings: OcrSettings, + pub secrets: Secrets, } impl OcrConnection { @@ -167,6 +169,7 @@ impl OcrConnection { credentials: ResolvedOcrCredentials, transport: OcrTransportConfig, settings: OcrSettings, + secrets: Secrets, ) -> Self { let api_key_source = credentials .api_key @@ -191,8 +194,13 @@ impl OcrConnection { .unwrap_or(settings.request_timeout), max_response_bytes: transport.max_response_bytes, settings, + secrets, } } + + pub fn secret(&self, name: &str) -> Option { + self.secrets.get(name) + } } impl Default for OcrConnection { @@ -201,6 +209,7 @@ impl Default for OcrConnection { ResolvedOcrCredentials::default(), OcrTransportConfig::default(), OcrSettings::default(), + Arc::new(ProcessEnvironment), ) } } @@ -563,10 +572,6 @@ pub fn decode_and_normalize_response( }) } -pub fn credential_env(name: &str) -> Option { - std::env::var(name).ok() -} - #[cfg(test)] mod tests { use serde_json::json; @@ -587,6 +592,7 @@ mod tests { ..OcrTransportConfig::default() }, settings.clone(), + Arc::new(ProcessEnvironment), ) .timeout }; diff --git a/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs b/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs index da6cf90ffcf..d141c68db38 100644 --- a/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs @@ -13,7 +13,7 @@ use crate::base_llm::ocr::{ handler::OcrClient, transformation::{ BaseOcrConfig, LiteLLMOcrResponse, OCR_INLINE_MAX_BYTES, OcrConnection, OcrDocument, - OcrPage, OcrPageImage, OcrResponseFormat, OcrUsageInfo, PreparedOcrRequest, credential_env, + OcrPage, OcrPageImage, OcrResponseFormat, OcrUsageInfo, PreparedOcrRequest, decode_and_normalize_response, decode_response_value, }, }; @@ -122,7 +122,9 @@ impl BaseOcrConfig for CohereParseConfig { request: &PreparedOcrRequest, _client: &OcrClient, ) -> Result { - self.resolve_headers(&request.connection, &credential_env) + self.resolve_headers(&request.connection, &|name: &str| { + request.connection.secret(name) + }) } fn get_complete_url( diff --git a/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs b/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs index 95658837fc3..2b14372fbec 100644 --- a/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs @@ -7,7 +7,7 @@ use crate::base_llm::ocr::{ handler::OcrClient, transformation::{ BaseOcrConfig, LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrPage, OcrResponseFormat, - OcrUsageInfo, PreparedOcrRequest, credential_env, decode_and_normalize_response, + OcrUsageInfo, PreparedOcrRequest, decode_and_normalize_response, }, }; @@ -84,7 +84,9 @@ impl BaseOcrConfig for MistralOcrConfig { request: &PreparedOcrRequest, _client: &OcrClient, ) -> Result { - self.resolve_headers(&request.connection, &credential_env) + self.resolve_headers(&request.connection, &|name: &str| { + request.connection.secret(name) + }) } fn get_complete_url( diff --git a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs index 740f0ced090..307ba697316 100644 --- a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs @@ -15,7 +15,7 @@ use crate::base_llm::ocr::{ transformation::{ BaseOcrConfig, LiteLLMOcrResponse, OCR_INLINE_MAX_BYTES, OcrConnection, OcrDocument, OcrPage, OcrRequestContext, OcrResponseFormat, OcrUsageInfo, PreparedOcrRequest, - credential_env, decode_and_normalize_response, + decode_and_normalize_response, }, }; @@ -110,7 +110,9 @@ impl BaseOcrConfig for ReductoParseV3Config { request: &PreparedOcrRequest, _client: &OcrClient, ) -> Result { - resolve_headers(&request.connection, &credential_env) + resolve_headers(&request.connection, &|name: &str| { + request.connection.secret(name) + }) } fn get_complete_url( diff --git a/litellm-rust/crates/llms/src/vertex_ai/ocr/deepseek_transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/ocr/deepseek_transformation.rs index 8009a65ff77..6fa0b9c5977 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/ocr/deepseek_transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/ocr/deepseek_transformation.rs @@ -9,7 +9,7 @@ use crate::base_llm::ocr::{ handler::OcrClient, transformation::{ BaseOcrConfig, LiteLLMOcrResponse, OcrDocument, OcrPage, OcrPageDimensions, OcrPageImage, - OcrRequestContext, OcrResponseFormat, OcrUsageInfo, PreparedOcrRequest, credential_env, + OcrRequestContext, OcrResponseFormat, OcrUsageInfo, PreparedOcrRequest, decode_and_normalize_response, decode_response_value, }, }; @@ -126,8 +126,9 @@ impl BaseOcrConfig for VertexAIDeepSeekOCRConfig { &request.optional_params, &request.input_sources, )?; - let location = vertex::get_vertex_ai_location(&config, &credential_env) - .unwrap_or_else(|| DEFAULT_LOCATION.to_string()); + let location = + vertex::get_vertex_ai_location(&config, &|name: &str| request.connection.secret(name)) + .unwrap_or_else(|| DEFAULT_LOCATION.to_string()); self.get_complete_url( request.connection.api_base.as_deref(), &environment.project_id, diff --git a/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs index a50e8261aa3..f7941db9364 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs @@ -10,7 +10,7 @@ use crate::{ handler::OcrClient, transformation::{ BaseOcrConfig, LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrEnvironment, - OcrRequestContext, OcrResponseFormat, PreparedOcrRequest, credential_env, + OcrRequestContext, OcrResponseFormat, PreparedOcrRequest, }, }, mistral::ocr::transformation::{MistralOcrConfig, MistralOcrRequest}, @@ -65,8 +65,9 @@ impl BaseOcrConfig for VertexAiOcrConfig { &request.optional_params, &request.input_sources, )?; - let location = vertex::get_vertex_ai_location(&config, &credential_env) - .unwrap_or_else(|| DEFAULT_LOCATION.to_string()); + let location = + vertex::get_vertex_ai_location(&config, &|name: &str| request.connection.secret(name)) + .unwrap_or_else(|| DEFAULT_LOCATION.to_string()); self.build_ocr_url( request.connection.api_base.as_deref(), &environment.project_id, @@ -139,7 +140,7 @@ impl VertexAiOcrConfig { .as_ref() .map(litellm_auth::SecretValue::expose), config, - &credential_env, + &|name: &str| connection.secret(name), ) .await .map_err(Error::from) diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs index 79921d67452..272e711ada5 100644 --- a/litellm-rust/crates/python-bridge/src/python_settings.rs +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -1,3 +1,4 @@ +use litellm_core_utils::settings::Lookup; use pyo3::prelude::*; const MODULE: &str = "litellm.rust_bridge.settings"; @@ -29,6 +30,23 @@ impl PythonSettings { } } +pub(crate) struct PythonSecrets; + +impl Lookup for PythonSecrets { + fn get(&self, name: &str) -> Option { + Python::attach(|py| { + py.import(MODULE) + .and_then(|module| module.getattr("secret")?.call1((name,))) + .and_then(|value| value.extract::>()) + .unwrap_or_else(|error| { + let _ = + PythonSettings::warn(py, &format!("reading secret {name} failed: {error}")); + None + }) + }) + } +} + #[cfg(test)] pub(crate) const CONTRACT: &str = include_str!("../python_settings.json"); @@ -36,9 +54,10 @@ pub(crate) const CONTRACT: &str = include_str!("../python_settings.json"); mod tests { use std::{collections::BTreeSet, ffi::CString}; + use litellm_core_utils::settings::Lookup; use pyo3::{prelude::*, types::PyDict}; - use super::{CONTRACT, PythonSettings}; + use super::{CONTRACT, PythonSecrets, PythonSettings}; #[test] fn every_settings_group_is_in_the_python_contract() { @@ -62,4 +81,48 @@ mod tests { assert_eq!(read, declared); }); } + + #[test] + fn secrets_come_from_the_python_secret_reader_and_a_failed_read_is_unset() { + Python::initialize(); + Python::attach(|py| { + py.run( + c" +import sys +import types +settings = types.ModuleType('litellm.rust_bridge.settings') +settings.warnings = [] +def secret(name): + if name == 'BROKEN': + raise RuntimeError('vault down') + return {'MISTRAL_API_KEY': 'from-vault'}.get(name) +settings.secret = secret +settings.warn = settings.warnings.append +sys.modules.setdefault('litellm', types.ModuleType('litellm')) +sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bridge')) +sys.modules['litellm.rust_bridge.settings'] = settings +", + None, + None, + ) + .unwrap(); + }); + assert_eq!( + PythonSecrets.get("MISTRAL_API_KEY").as_deref(), + Some("from-vault") + ); + assert_eq!(PythonSecrets.get("ABSENT"), None); + assert_eq!(PythonSecrets.get("BROKEN"), None); + Python::attach(|py| { + let warnings: Vec = py + .import("litellm.rust_bridge.settings") + .unwrap() + .getattr("warnings") + .unwrap() + .extract() + .unwrap(); + assert_eq!(warnings.len(), 1); + assert!(warnings[0].contains("BROKEN") && warnings[0].contains("vault down")); + }); + } } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs index bb845f48783..966b24a82e7 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs @@ -3,7 +3,7 @@ mod errors; mod host; mod project; -use std::sync::LazyLock; +use std::sync::{Arc, LazyLock}; use host::OcrRouteHost; use litellm_auth_gcp::VertexAuth; @@ -16,7 +16,7 @@ use pyo3::{ types::{PyDict, PyTuple}, }; -use crate::{errors::RustBridgeDeclined, http}; +use crate::{errors::RustBridgeDeclined, http, python_settings::PythonSecrets}; const SURFACE: LegacySurface = LegacySurface { call_type: "ocr", @@ -45,6 +45,7 @@ fn run_ocr( http::url_policy(py)?, VERTEX_AUTH.clone(), OcrSettings::from_environment(&ProcessEnvironment), + Arc::new(PythonSecrets), ) .map_err(|error| RustBridgeDeclined::new_err(error.to_string()))?; run_legacy_call( diff --git a/litellm/rust_bridge/settings.py b/litellm/rust_bridge/settings.py index e170f93b198..210ef7ac6b4 100644 --- a/litellm/rust_bridge/settings.py +++ b/litellm/rust_bridge/settings.py @@ -30,6 +30,12 @@ def warn(message: str) -> None: verbose_logger.warning("%s", message) +def secret(name: str) -> str | None: + from litellm.secret_managers.main import get_secret_str + + return get_secret_str(name) + + def url_policy() -> UrlPolicy: import litellm diff --git a/tests/test_litellm/rust_bridge/test_settings.py b/tests/test_litellm/rust_bridge/test_settings.py index f75145c2b2c..f3baf463b87 100644 --- a/tests/test_litellm/rust_bridge/test_settings.py +++ b/tests/test_litellm/rust_bridge/test_settings.py @@ -3,12 +3,15 @@ import logging from pathlib import Path from typing import Final +import httpx import pytest from pydantic import TypeAdapter import litellm +from litellm.integrations.custom_secret_manager import CustomSecretManager from litellm.llms.custom_httpx.http_handler import default_user_agent from litellm.rust_bridge import settings +from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem CONTRACT_PATH: Final = Path(__file__).parents[3] / "litellm-rust/crates/python-bridge/python_settings.json" @@ -73,3 +76,46 @@ def test_warn_reaches_the_litellm_logger(caplog: pytest.LogCaptureFixture) -> No settings.warn("ssl_ecdh_curve 'secp521r1' is not supported") assert [record.getMessage() for record in caplog.records] == ["ssl_ecdh_curve 'secp521r1' is not supported"] + + +class _VaultSecrets(CustomSecretManager): + def __init__(self, secrets: dict[str, str]) -> None: + super().__init__(secret_manager_name="rust_bridge_settings_test") + self.secrets = secrets + + async def async_read_secret( + self, + secret_name: str, + optional_params: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + ) -> str | None: + return self.secrets.get(secret_name) + + def sync_read_secret( + self, + secret_name: str, + optional_params: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + ) -> str | None: + return self.secrets.get(secret_name) + + +def test_secret_prefers_the_secret_manager_and_falls_back_to_the_environment_on_a_miss( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("MISTRAL_API_KEY", "stale-env-key") + monkeypatch.setenv("REDUCTO_API_KEY", "env-only-key") + monkeypatch.setattr(litellm, "secret_manager_client", _VaultSecrets({"MISTRAL_API_KEY": "vault-key"})) + monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.CUSTOM) + monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode="read_only")) + + assert settings.secret("MISTRAL_API_KEY") == "vault-key" + assert settings.secret("REDUCTO_API_KEY") == "env-only-key" + assert settings.secret("ABSENT_KEY") is None + + +def test_secret_reads_the_environment_without_a_secret_manager(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("MISTRAL_API_KEY", "env-key") + monkeypatch.setattr(litellm, "secret_manager_client", None) + + assert settings.secret("MISTRAL_API_KEY") == "env-key" From 1ee4b62e9c2536fbf973830f324e62d1c02e57f1 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 21:01:37 -0700 Subject: [PATCH 06/23] fix(rust): honor vertex_project, vertex_location and enable_azure_ad_token_refresh globals Python resolves the Vertex project and location as call params, then the litellm.vertex_project / litellm.vertex_location globals, then env, and Azure AD token refresh from litellm.enable_azure_ad_token_refresh alone. Native OCR skipped the globals, so a config.yaml litellm_settings value silently fell through to the credential's project and us-central1, and a managed identity setup without an API key failed. The bridge now reads them through a provider_defaults settings group into OcrSettings, and VertexConfig / AzureAuthInputs slot them in at Python's precedence. Co-Authored-By: Claude Opus 5 --- litellm-rust/Cargo.lock | 1 + litellm-rust/crates/auth-azure/Cargo.toml | 1 + litellm-rust/crates/auth-azure/src/types.rs | 48 ++++++++++++++++--- litellm-rust/crates/auth-gcp/src/lib.rs | 47 ++++++++++++++---- .../crates/core/tests/vertex_ai_ocr.rs | 25 +++++++++- .../llms/src/azure_ai/ocr/common_utils.rs | 16 ++++++- .../document_intelligence/transformation.rs | 8 +--- .../llms/src/azure_ai/ocr/transformation.rs | 8 +--- .../crates/llms/src/base_llm/ocr/settings.rs | 8 ++++ .../llms/src/vertex_ai/ocr/common_utils.rs | 18 ++++++- .../vertex_ai/ocr/deepseek_transformation.rs | 9 ++-- .../llms/src/vertex_ai/ocr/transformation.rs | 12 ++--- .../crates/python-bridge/python_settings.json | 5 ++ .../python-bridge/src/python_settings.rs | 4 +- .../python-bridge/src/routes/ocr/mod.rs | 32 ++++++++++++- litellm/rust_bridge/settings.py | 17 +++++++ .../test_litellm/rust_bridge/test_settings.py | 13 +++++ 17 files changed, 221 insertions(+), 51 deletions(-) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 5fbddcaffcf..0cbca96ad57 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1976,6 +1976,7 @@ dependencies = [ "azure_identity", "litellm-auth", "moka", + "rstest", "serde_json", "sha2 0.10.9", "strum", diff --git a/litellm-rust/crates/auth-azure/Cargo.toml b/litellm-rust/crates/auth-azure/Cargo.toml index 9f8260c7b3f..8099506d2e5 100644 --- a/litellm-rust/crates/auth-azure/Cargo.toml +++ b/litellm-rust/crates/auth-azure/Cargo.toml @@ -18,4 +18,5 @@ azure_core = "1.0.0" azure_identity = { version = "1.0.0", features = ["tokio"] } [dev-dependencies] +rstest.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/auth-azure/src/types.rs b/litellm-rust/crates/auth-azure/src/types.rs index 2a510de1f43..87e883a6a54 100644 --- a/litellm-rust/crates/auth-azure/src/types.rs +++ b/litellm-rust/crates/auth-azure/src/types.rs @@ -1,11 +1,10 @@ -use serde_json::{Map, Value}; use std::collections::BTreeMap; -use strum::EnumString; -use litellm_auth::Error; use litellm_auth::{ - CredentialResolverHandle, InputSource, SecretValue, Sourced, TokenProviderHandle, + CredentialResolverHandle, Error, InputSource, SecretValue, Sourced, TokenProviderHandle, }; +use serde_json::{Map, Value}; +use strum::EnumString; pub const DEFAULT_AZURE_SCOPE: &str = "https://cognitiveservices.azure.com/.default"; @@ -52,6 +51,16 @@ pub struct AzureAuthInputs { } impl AzureAuthInputs { + pub fn or_configured_token_refresh(self, enabled: bool) -> Self { + if *self.enable_azure_ad_token_refresh.value() || !enabled { + return self; + } + Self { + enable_azure_ad_token_refresh: Sourced::new(true, InputSource::Deployment), + ..self + } + } + #[cfg(test)] pub fn from_optional_params(params: &Map) -> Result { Self::from_sourced_optional_params(params, &BTreeMap::new()) @@ -115,12 +124,12 @@ fn source_for(sources: &BTreeMap, name: &str) -> InputSourc #[cfg(test)] mod tests { - use serde_json::json; - use std::collections::BTreeMap; - use super::{AzureAuthInputs, AzureCredentialType, ConfigValue}; use litellm_auth::{InputSource, Sourced}; + use serde_json::json; + + use super::{AzureAuthInputs, AzureCredentialType, ConfigValue}; #[test] fn selector_parsing_is_exact() { @@ -189,4 +198,29 @@ mod tests { assert!(!debug.contains("token-value")); assert!(!debug.contains("secret-value")); } + + #[rstest::rstest] + #[case::global_turns_refresh_on(json!({}), true, true, InputSource::Deployment)] + #[case::global_overrides_a_call_false_like_python(json!({"enable_azure_ad_token_refresh": false}), true, true, InputSource::Deployment)] + #[case::call_true_survives_a_global_false(json!({"enable_azure_ad_token_refresh": true}), false, true, InputSource::Request)] + #[case::both_off(json!({}), false, false, InputSource::Request)] + fn token_refresh_follows_the_configured_global( + #[case] params: serde_json::Value, + #[case] global: bool, + #[case] enabled: bool, + #[case] source: InputSource, + ) { + let sources = BTreeMap::from([( + "enable_azure_ad_token_refresh".to_string(), + InputSource::Request, + )]); + let inputs = + AzureAuthInputs::from_sourced_optional_params(params.as_object().unwrap(), &sources) + .unwrap() + .or_configured_token_refresh(global); + assert_eq!( + inputs.enable_azure_ad_token_refresh, + Sourced::new(enabled, source) + ); + } } diff --git a/litellm-rust/crates/auth-gcp/src/lib.rs b/litellm-rust/crates/auth-gcp/src/lib.rs index f8402624edc..bf619fee144 100644 --- a/litellm-rust/crates/auth-gcp/src/lib.rs +++ b/litellm-rust/crates/auth-gcp/src/lib.rs @@ -1,17 +1,13 @@ -use std::collections::BTreeMap; -use std::future::Future; -use std::path::Path; -use std::pin::Pin; -use std::sync::Arc; +use std::{collections::BTreeMap, future::Future, path::Path, pin::Pin, sync::Arc}; use gcp_auth::{CustomServiceAccount, TokenProvider}; +use litellm_auth::{ + CredentialPlacement, Error, InputSource, SecretValue, Sourced, http::apply_credential, +}; use moka::future::Cache; use serde_json::{Map, Value}; use sha2::{Digest, Sha256}; -use litellm_auth::http::apply_credential; -use litellm_auth::{CredentialPlacement, Error, InputSource, SecretValue, Sourced}; - const CLOUD_PLATFORM_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform"; const GOOGLE_OAUTH_TOKEN_ENDPOINT: &str = "https://oauth2.googleapis.com/token"; const GOOGLE_APPLICATION_CREDENTIALS_ENV: &str = "GOOGLE_APPLICATION_CREDENTIALS"; @@ -45,6 +41,16 @@ impl VertexConfig { }) } + pub fn or_configured(self, project_id: Option<&str>, location: Option<&str>) -> Self { + let configured = + |value: Option<&str>| value.filter(|value| !value.is_empty()).map(str::to_string); + Self { + project_id: self.project_id.or_else(|| configured(project_id)), + location: self.location.or_else(|| configured(location)), + ..self + } + } + pub fn project_id(&self) -> Option<&str> { self.project_id.as_deref() } @@ -571,4 +577,29 @@ mod tests { assert_eq!(loads.load(Ordering::SeqCst), 1); assert_eq!(calls.load(Ordering::SeqCst), 4); } + + #[test] + fn configured_defaults_sit_between_call_params_and_the_environment() { + let env = |name: &str| Some(format!("env-{name}")); + let from_config = + VertexConfig::default().or_configured(Some("global-project"), Some("global-location")); + assert_eq!( + get_vertex_ai_project(&from_config, &env).as_deref(), + Some("global-project") + ); + assert_eq!( + get_vertex_ai_location(&from_config, &env).as_deref(), + Some("global-location") + ); + let from_call = + config(json!({"vertex_project":"call-project","vertex_location":"call-location"})) + .or_configured(Some("global-project"), Some("global-location")); + assert_eq!(from_call.project_id(), Some("call-project")); + assert_eq!(from_call.location(), Some("call-location")); + let empty_global = VertexConfig::default().or_configured(Some(""), None); + assert_eq!( + get_vertex_ai_project(&empty_global, &env).as_deref(), + Some("env-VERTEXAI_PROJECT") + ); + } } diff --git a/litellm-rust/crates/core/tests/vertex_ai_ocr.rs b/litellm-rust/crates/core/tests/vertex_ai_ocr.rs index 1f1186c7827..399b7cac39a 100644 --- a/litellm-rust/crates/core/tests/vertex_ai_ocr.rs +++ b/litellm-rust/crates/core/tests/vertex_ai_ocr.rs @@ -1,8 +1,8 @@ use litellm_auth::InputSource; -use litellm_llms::base_llm::ocr::transformation::OcrResponseFormat; +use litellm_llms::base_llm::ocr::{settings::OcrSettings, transformation::OcrResponseFormat}; use serde_json::{Value, json}; -use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; +use super::test_support::{MockResponse, mock_server, ocr_client, perform_ocr, wire_request}; fn request_body(request: &str) -> Value { serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() @@ -48,6 +48,27 @@ async fn facade_executes_vertex_mistral_with_resolved_project_and_location() { ); } +#[tokio::test] +async fn configured_project_and_location_apply_when_the_call_sets_neither() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let client = ocr_client().with_settings(OcrSettings { + vertex_project: Some("configured-project".into()), + vertex_location: Some("europe-west4".into()), + ..OcrSettings::default() + }); + + crate::ocr::client::perform( + &client, + wire_request("vertex_ai/mistral-ocr-maas", &base, json!({})), + ) + .await + .unwrap(); + server.await.unwrap(); + assert!(seen.lock().unwrap()[0].starts_with( + "POST /v1/projects/configured-project/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict " + )); +} + #[tokio::test] async fn supplied_authorization_is_forwarded_without_a_static_token() { let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/common_utils.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/common_utils.rs index 26eeeb6635c..9c2f3f70b91 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/common_utils.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/common_utils.rs @@ -3,7 +3,21 @@ use std::sync::OnceLock; use litellm_auth::{InputSource, Sourced}; use litellm_auth_azure::{AzureAuthInputs, AzureAuthService}; -use crate::base_llm::ocr::{error::Error, transformation::OcrConnection}; +use crate::base_llm::ocr::{ + error::Error, + transformation::{OcrConnection, PreparedOcrRequest}, +}; + +pub(crate) fn azure_auth_inputs(request: &PreparedOcrRequest) -> Result { + Ok(AzureAuthInputs { + azure_ad_token_provider: request.azure_ad_token_provider.clone(), + ..AzureAuthInputs::from_sourced_optional_params( + &request.optional_params, + &request.input_sources, + )? + } + .or_configured_token_refresh(request.connection.settings.enable_azure_ad_token_refresh)) +} pub(super) async fn resolve_entra( config: &AzureAuthInputs, diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs index 8e6182f454f..9b27fdbb568 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs @@ -174,13 +174,7 @@ impl BaseOcrConfig for AzureDocumentIntelligenceOcrConfig { request: &PreparedOcrRequest, _client: &OcrClient, ) -> Result { - let config = AzureAuthInputs { - azure_ad_token_provider: request.azure_ad_token_provider.clone(), - ..AzureAuthInputs::from_sourced_optional_params( - &request.optional_params, - &request.input_sources, - )? - }; + let config = crate::azure_ai::ocr::common_utils::azure_auth_inputs(request)?; self.resolve_headers(&request.connection, &config, &|name: &str| { request.connection.secret(name) }) diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs index cd20e75df85..6df83e57eab 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs @@ -50,13 +50,7 @@ impl BaseOcrConfig for AzureAiOcrConfig { request: &PreparedOcrRequest, _client: &OcrClient, ) -> Result { - let config = AzureAuthInputs { - azure_ad_token_provider: request.azure_ad_token_provider.clone(), - ..AzureAuthInputs::from_sourced_optional_params( - &request.optional_params, - &request.input_sources, - )? - }; + let config = crate::azure_ai::ocr::common_utils::azure_auth_inputs(request)?; self.resolve_headers(&request.connection, &config, &|name: &str| { request.connection.secret(name) }) diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/settings.rs b/litellm-rust/crates/llms/src/base_llm/ocr/settings.rs index 276b2ca1311..f5954599b43 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/settings.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/settings.rs @@ -11,6 +11,9 @@ pub struct OcrSettings { pub poll_timeout: Duration, pub document_intelligence_api_version: String, pub document_intelligence_dpi: i64, + pub vertex_project: Option, + pub vertex_location: Option, + pub enable_azure_ad_token_refresh: bool, } impl Default for OcrSettings { @@ -21,6 +24,9 @@ impl Default for OcrSettings { poll_timeout: Duration::from_secs(120), document_intelligence_api_version: "2024-11-30".into(), document_intelligence_dpi: 96, + vertex_project: None, + vertex_location: None, + enable_azure_ad_token_refresh: false, } } } @@ -48,6 +54,7 @@ impl OcrSettings { document_intelligence_dpi: env .parsed("AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI") .unwrap_or(defaults.document_intelligence_dpi), + ..defaults } } } @@ -96,6 +103,7 @@ mod tests { poll_timeout: Duration::from_secs(600), document_intelligence_api_version: "2025-01-01".into(), document_intelligence_dpi: 72, + ..OcrSettings::default() } ); } diff --git a/litellm-rust/crates/llms/src/vertex_ai/ocr/common_utils.rs b/litellm-rust/crates/llms/src/vertex_ai/ocr/common_utils.rs index 979c9526f96..46285874d9f 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/ocr/common_utils.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/ocr/common_utils.rs @@ -1,6 +1,22 @@ use litellm_auth::InputSource; +use litellm_auth_gcp::VertexConfig; -use crate::base_llm::ocr::{error::Error, transformation::OcrConnection}; +use crate::base_llm::ocr::{ + error::Error, + transformation::{OcrConnection, PreparedOcrRequest}, +}; + +pub(super) fn vertex_config(request: &PreparedOcrRequest) -> Result { + let settings = &request.connection.settings; + Ok(VertexConfig::from_sourced_optional_params( + &request.optional_params, + &request.input_sources, + )? + .or_configured( + settings.vertex_project.as_deref(), + settings.vertex_location.as_deref(), + )) +} pub(super) fn validate_destination(connection: &OcrConnection) -> Result<(), Error> { if connection.api_base.is_some() && connection.api_base_source == InputSource::Request { diff --git a/litellm-rust/crates/llms/src/vertex_ai/ocr/deepseek_transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/ocr/deepseek_transformation.rs index 6fa0b9c5977..f0b035621fa 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/ocr/deepseek_transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/ocr/deepseek_transformation.rs @@ -1,9 +1,9 @@ -use litellm_auth_gcp::{self as vertex, VertexConfig}; +use litellm_auth_gcp as vertex; use litellm_core_utils::{call_arguments::CallArguments, params::OpaqueParams, url_utils::ApiUrl}; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; -use super::transformation::VertexAiOcrConfig; +use super::{common_utils::vertex_config, transformation::VertexAiOcrConfig}; use crate::base_llm::ocr::{ error::Error, handler::OcrClient, @@ -122,10 +122,7 @@ impl BaseOcrConfig for VertexAIDeepSeekOCRConfig { _params: &Self::OcrParams, environment: &Self::Environment, ) -> Result { - let config = VertexConfig::from_sourced_optional_params( - &request.optional_params, - &request.input_sources, - )?; + let config = vertex_config(request)?; let location = vertex::get_vertex_ai_location(&config, &|name: &str| request.connection.secret(name)) .unwrap_or_else(|| DEFAULT_LOCATION.to_string()); diff --git a/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs index f7941db9364..2d505ba4342 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs @@ -2,7 +2,7 @@ use litellm_auth_gcp::{self as vertex, VertexConfig}; use litellm_core_utils::{call_arguments::CallArguments, params::OpaqueParams, url_utils::ApiUrl}; use serde_json::Value; -use super::common_utils::validate_destination; +use super::common_utils::{validate_destination, vertex_config}; use crate::{ base_llm::ocr::{ document::{inline_remote_document, validate_inline_document}, @@ -47,10 +47,7 @@ impl BaseOcrConfig for VertexAiOcrConfig { request: &PreparedOcrRequest, client: &OcrClient, ) -> Result { - let config = VertexConfig::from_sourced_optional_params( - &request.optional_params, - &request.input_sources, - )?; + let config = vertex_config(request)?; self.resolve_environment(&request.connection, &config, client) .await } @@ -61,10 +58,7 @@ impl BaseOcrConfig for VertexAiOcrConfig { _optional_params: &Self::OcrParams, environment: &Self::Environment, ) -> Result { - let config = VertexConfig::from_sourced_optional_params( - &request.optional_params, - &request.input_sources, - )?; + let config = vertex_config(request)?; let location = vertex::get_vertex_ai_location(&config, &|name: &str| request.connection.secret(name)) .unwrap_or_else(|| DEFAULT_LOCATION.to_string()); diff --git a/litellm-rust/crates/python-bridge/python_settings.json b/litellm-rust/crates/python-bridge/python_settings.json index a6f5ee9c6f4..4ad3edf682d 100644 --- a/litellm-rust/crates/python-bridge/python_settings.json +++ b/litellm-rust/crates/python-bridge/python_settings.json @@ -14,5 +14,10 @@ "url_policy": [ "user_url_validation", "user_url_allowed_hosts" + ], + "provider_defaults": [ + "vertex_project", + "vertex_location", + "enable_azure_ad_token_refresh" ] } diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs index 272e711ada5..83db4f02500 100644 --- a/litellm-rust/crates/python-bridge/src/python_settings.rs +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -7,16 +7,18 @@ const MODULE: &str = "litellm.rust_bridge.settings"; pub(crate) enum PythonSettings { Http, UrlPolicy, + ProviderDefaults, } impl PythonSettings { #[cfg(test)] - pub(crate) const ALL: [Self; 2] = [Self::Http, Self::UrlPolicy]; + pub(crate) const ALL: [Self; 3] = [Self::Http, Self::UrlPolicy, Self::ProviderDefaults]; pub(crate) fn name(self) -> &'static str { match self { Self::Http => "http_settings", Self::UrlPolicy => "url_policy", + Self::ProviderDefaults => "provider_defaults", } } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs index 966b24a82e7..785f6e48e13 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs @@ -16,7 +16,11 @@ use pyo3::{ types::{PyDict, PyTuple}, }; -use crate::{errors::RustBridgeDeclined, http, python_settings::PythonSecrets}; +use crate::{ + errors::RustBridgeDeclined, + http, + python_settings::{PythonSecrets, PythonSettings}, +}; const SURFACE: LegacySurface = LegacySurface { call_type: "ocr", @@ -44,7 +48,7 @@ fn run_ocr( &config, http::url_policy(py)?, VERTEX_AUTH.clone(), - OcrSettings::from_environment(&ProcessEnvironment), + ocr_settings(py)?, Arc::new(PythonSecrets), ) .map_err(|error| RustBridgeDeclined::new_err(error.to_string()))?; @@ -58,6 +62,30 @@ fn run_ocr( ) } +#[derive(FromPyObject)] +struct PythonProviderDefaults { + vertex_project: Option, + vertex_location: Option, + enable_azure_ad_token_refresh: Option, +} + +fn ocr_settings(py: Python<'_>) -> PyResult { + let defaults: PythonProviderDefaults = PythonSettings::ProviderDefaults + .read(py)? + .extract() + .map_err(|error: PyErr| { + RustBridgeDeclined::new_err(format!( + "litellm provider defaults cannot be used by the Rust route: {error}" + )) + })?; + Ok(OcrSettings { + vertex_project: defaults.vertex_project, + vertex_location: defaults.vertex_location, + enable_azure_ad_token_refresh: defaults.enable_azure_ad_token_refresh == Some(true), + ..OcrSettings::from_environment(&ProcessEnvironment) + }) +} + #[pyfunction] pub(crate) fn ocr( py: Python<'_>, diff --git a/litellm/rust_bridge/settings.py b/litellm/rust_bridge/settings.py index 210ef7ac6b4..037d6d9bd27 100644 --- a/litellm/rust_bridge/settings.py +++ b/litellm/rust_bridge/settings.py @@ -24,6 +24,13 @@ class UrlPolicy: user_url_allowed_hosts: Sequence[str] +@dataclass(frozen=True, slots=True) +class ProviderDefaults: + vertex_project: str | None + vertex_location: str | None + enable_azure_ad_token_refresh: bool | None + + def warn(message: str) -> None: from litellm._logging import verbose_logger @@ -36,6 +43,16 @@ def secret(name: str) -> str | None: return get_secret_str(name) +def provider_defaults() -> ProviderDefaults: + import litellm + + return ProviderDefaults( + vertex_project=litellm.vertex_project, + vertex_location=litellm.vertex_location, + enable_azure_ad_token_refresh=litellm.enable_azure_ad_token_refresh, + ) + + def url_policy() -> UrlPolicy: import litellm diff --git a/tests/test_litellm/rust_bridge/test_settings.py b/tests/test_litellm/rust_bridge/test_settings.py index f3baf463b87..44c5ec42b36 100644 --- a/tests/test_litellm/rust_bridge/test_settings.py +++ b/tests/test_litellm/rust_bridge/test_settings.py @@ -22,6 +22,7 @@ def test_the_rust_contract_matches_the_returned_fields() -> None: assert contract == { "http_settings": [field.name for field in dataclasses.fields(settings.http_settings())], "url_policy": [field.name for field in dataclasses.fields(settings.url_policy())], + "provider_defaults": [field.name for field in dataclasses.fields(settings.provider_defaults())], } @@ -119,3 +120,15 @@ def test_secret_reads_the_environment_without_a_secret_manager(monkeypatch: pyte monkeypatch.setattr(litellm, "secret_manager_client", None) assert settings.secret("MISTRAL_API_KEY") == "env-key" + + +def test_provider_defaults_read_the_litellm_globals(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "vertex_project", "configured-project") + monkeypatch.setattr(litellm, "vertex_location", "europe-west4") + monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", True) + + assert settings.provider_defaults() == settings.ProviderDefaults( + vertex_project="configured-project", + vertex_location="europe-west4", + enable_azure_ad_token_refresh=True, + ) From a987efca2cd40ff6866b0f161b921a646e36f8a0 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 18 Sep 2026 22:39:27 -0700 Subject: [PATCH 07/23] fix(proxy): refuse runtime writes to config-owned settings A write into a settings store for a key the config file declares used to land in the runtime layer and then lose to the config on every read, so the caller saw success while nothing changed. It now raises ConfigOwnedKeyError, and the allowed-IP routes turn that into a 400 naming the key instead of reporting success on a list they never changed. Both allowed-IP routes now build a new list rather than mutating the one the config layer holds, and the os.environ resolver rebuilds the config it is given instead of writing back into it, so a reader can no longer corrupt the raw values the store keeps for provenance. The database reload leaves a config-owned key alone rather than writing a normalized copy back over it, which would now raise and abort the rest of the reconcile pass. --- .../proxy/config_resolvers/settings_store.py | 14 +++- litellm/proxy/proxy_server.py | 55 +++++++++------ .../proxy_setting_endpoints.py | 36 +++++++--- .../config_resolvers/test_settings_store.py | 29 +++++++- tests/test_litellm/proxy/test_proxy_server.py | 70 +++++++++++++++++++ .../test_proxy_setting_endpoints.py | 49 +++++++++++++ 6 files changed, 215 insertions(+), 38 deletions(-) diff --git a/litellm/proxy/config_resolvers/settings_store.py b/litellm/proxy/config_resolvers/settings_store.py index d4ca0e87d2b..98ebfdabd0f 100644 --- a/litellm/proxy/config_resolvers/settings_store.py +++ b/litellm/proxy/config_resolvers/settings_store.py @@ -17,6 +17,14 @@ from litellm.proxy.config_resolvers.settings_rules import ( rule_for, ) + +class ConfigOwnedKeyError(RuntimeError): + def __init__(self, section: Section, key: str) -> None: + super().__init__(f"{section}.{key} is set in the config file and cannot be changed at runtime") + self.section: Final = section + self.key: Final = key + + _EMPTY_VALUES: Final[Mapping[str, JsonValue]] = MappingProxyType({}) _EMPTY_ROWS: Final[Mapping[DbRow, Mapping[str, JsonValue]]] = MappingProxyType({}) @@ -72,8 +80,8 @@ class SettingsStore(MutableMapping[str, JsonValue]): return resolved.value def __setitem__(self, key: str, value: JsonValue) -> None: - if self.owned_by_config(key): - return + if self.owned_by_config(key) and value != self.get(key): + raise ConfigOwnedKeyError(self._section, key) self._runtime_values = MappingProxyType({**self._runtime_values, key: value}) self._deleted_runtime_keys = self._deleted_runtime_keys - frozenset((key,)) @@ -81,7 +89,7 @@ class SettingsStore(MutableMapping[str, JsonValue]): if key not in self: raise KeyError(key) if self.owned_by_config(key): - return + raise ConfigOwnedKeyError(self._section, key) self._runtime_values = MappingProxyType( {key_: value for key_, value in self._runtime_values.items() if key_ != key} ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index daf94b79849..eb9889f1877 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -5221,20 +5221,27 @@ class ProxyConfig: verbose_proxy_logger.warning("Maximum recursion depth (%s) reached while processing config.", max_depth) return config - for key, value in config.items(): - if isinstance(value, dict): - config[key] = self._check_for_os_environ_vars(config=value, depth=depth + 1, max_depth=max_depth) - elif isinstance(value, list): - for item in value: - if isinstance(item, dict): - item = self._check_for_os_environ_vars(config=item, depth=depth + 1, max_depth=max_depth) - # if the value is a string and starts with "os.environ/" - then it's an environment variable - elif isinstance(value, str) and value.startswith("os.environ/"): - resolved = get_secret(value) - if resolved is None and secret_manager_would_be_consulted(value): - verbose_proxy_logger.warning("%s is absent from the configured secret manager", value) - config[key] = resolved - return config + return { + key: self._resolved_config_value(value=value, depth=depth, max_depth=max_depth) + for key, value in config.items() + } + + def _resolved_config_value(self, value: object, depth: int, max_depth: int) -> object: + if isinstance(value, dict): + return self._check_for_os_environ_vars(config=value, depth=depth + 1, max_depth=max_depth) + if isinstance(value, list): + return [ + self._check_for_os_environ_vars(config=item, depth=depth + 1, max_depth=max_depth) + if isinstance(item, dict) + else item + for item in value + ] + if isinstance(value, str) and value.startswith("os.environ/"): + resolved: Final = get_secret(value) + if resolved is None and secret_manager_would_be_consulted(value): + verbose_proxy_logger.warning("%s is absent from the configured secret manager", value) + return resolved + return value def _initialize_secret_manager_from_raw_config( self, config: Mapping[str, object], config_file_path: str | None @@ -7321,7 +7328,9 @@ class ProxyConfig: "disable_auto_add_proxy_admin_to_teams", "apply_user_budget_to_team_keys", ): - if key in db_values and (value := self.settings.get(key)) is not None: + if key not in db_values or self.settings.owned_by_config(key): + continue + if (value := self.settings.get(key)) is not None: self.settings[key] = coerce_bool(value) async def _apply_cache_size_setting( @@ -7331,21 +7340,24 @@ class ProxyConfig: ) -> None: if "user_api_key_cache_max_size" not in db_values and not cache_size_was_db: return + writable: Final = not self.settings.owned_by_config("user_api_key_cache_max_size") cache_value: Final = self.settings.get("user_api_key_cache_max_size") try: cache_max_size: Final = ConfigGeneralSettings.model_validate( MappingProxyType({"user_api_key_cache_max_size": cache_value}) ).user_api_key_cache_max_size except ValidationError: - self.settings.pop("user_api_key_cache_max_size", None) + if writable: + self.settings.pop("user_api_key_cache_max_size", None) verbose_proxy_logger.warning( "Ignoring invalid general_settings.user_api_key_cache_max_size=%r from the DB", cache_value ) return - if cache_max_size is None: - self.settings.pop("user_api_key_cache_max_size", None) - else: - self.settings["user_api_key_cache_max_size"] = cache_max_size + if writable: + if cache_max_size is None: + self.settings.pop("user_api_key_cache_max_size", None) + else: + self.settings["user_api_key_cache_max_size"] = cache_max_size user_api_key_cache.update_in_memory_max_size(cache_max_size) async def _apply_store_model_in_db_setting(self, db_values: Mapping[str, SettingsJsonValue]) -> None: @@ -7357,7 +7369,8 @@ class ProxyConfig: return normalized: Final = coerce_bool(value) store_model_in_db = normalized if isinstance(normalized, bool) else bool(normalized) - self.settings["store_model_in_db"] = store_model_in_db + if not self.settings.owned_by_config("store_model_in_db"): + self.settings["store_model_in_db"] = store_model_in_db async def _apply_retention_settings( self, diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 75431383fbd..b1f0ac7b35b 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -3,7 +3,7 @@ import asyncio import json import os from collections import Counter -from collections.abc import Mapping, Sequence +from collections.abc import Mapping, MutableMapping, Sequence from types import MappingProxyType from typing import ( Final, @@ -24,6 +24,7 @@ from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_keys from litellm.proxy._experimental.mcp_server.tool_search import MCP_TOOL_SEARCH_SETTINGS_KEY from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.config_resolvers.settings_store import ConfigOwnedKeyError from litellm.proxy.config_resolvers.sso import ( SSO_FIELD_ENV_VARS, SSO_SECRET_FIELDS, @@ -489,6 +490,23 @@ async def get_allowed_ips(): return {"data": _allowed_ip} +def _store_allowed_ips(general_settings: MutableMapping[str, object], allowed_ips: Sequence[str]) -> None: + try: + general_settings["allowed_ips"] = list(allowed_ips) + except ConfigOwnedKeyError as owned: + raise HTTPException( + status_code=400, + detail={ + "error": f"{owned.section} key '{owned.key}' is set in the config file and cannot be changed here", + "keys": [owned.key], + "section": owned.section, + "resolution": ( + "edit the config file to change it, or remove it from the file to let the database own it" + ), + }, + ) from owned + + @router.post( "/add/allowed_ip", tags=["Budget & Spend Tracking"], @@ -509,12 +527,10 @@ async def add_allowed_ip( if prisma_client is None: raise Exception("No DB Connected") - _allowed_ips: Final[list] = general_settings.get("allowed_ips", []) - if ip_address.ip not in _allowed_ips: - _allowed_ips.append(ip_address.ip) - general_settings["allowed_ips"] = _allowed_ips - else: + _allowed_ips: Final[Sequence[str]] = general_settings.get("allowed_ips") or [] + if ip_address.ip in _allowed_ips: raise HTTPException(status_code=400, detail="IP address already exists") + _store_allowed_ips(general_settings, (*_allowed_ips, ip_address.ip)) if store_model_in_db is not True: raise HTTPException( @@ -568,12 +584,10 @@ async def delete_allowed_ip( proxy_config, ) - _allowed_ips: Final[list] = general_settings.get("allowed_ips", []) - if ip_address.ip in _allowed_ips: - _allowed_ips.remove(ip_address.ip) - general_settings["allowed_ips"] = _allowed_ips - else: + _allowed_ips: Final[Sequence[str]] = general_settings.get("allowed_ips") or [] + if ip_address.ip not in _allowed_ips: raise HTTPException(status_code=404, detail="IP address not found") + _store_allowed_ips(general_settings, tuple(ip for ip in _allowed_ips if ip != ip_address.ip)) # Load existing config config: Final = await proxy_config.get_config() diff --git a/tests/test_litellm/proxy/config_resolvers/test_settings_store.py b/tests/test_litellm/proxy/config_resolvers/test_settings_store.py index 88ec382b013..18cbc3b0d6f 100644 --- a/tests/test_litellm/proxy/config_resolvers/test_settings_store.py +++ b/tests/test_litellm/proxy/config_resolvers/test_settings_store.py @@ -6,7 +6,7 @@ from unittest.mock import patch import pytest from litellm.proxy.config_resolvers.settings_rules import JsonValue -from litellm.proxy.config_resolvers.settings_store import SettingsStore +from litellm.proxy.config_resolvers.settings_store import ConfigOwnedKeyError, SettingsStore def test_settings_store_matches_plain_dict_mapping_operations() -> None: @@ -162,13 +162,27 @@ def test_settings_store_refuses_a_runtime_write_to_a_config_owned_key() -> None: store: Final = SettingsStore("general_settings") store.load_yaml({"max_parallel_requests": 3}) - store["max_parallel_requests"] = 11 - del store["max_parallel_requests"] + with pytest.raises(ConfigOwnedKeyError) as write: + store["max_parallel_requests"] = 11 + with pytest.raises(ConfigOwnedKeyError): + del store["max_parallel_requests"] + assert "max_parallel_requests" in str(write.value) assert store["max_parallel_requests"] == 3 assert store.source("max_parallel_requests") == "config" +def test_settings_store_accepts_a_write_that_does_not_change_a_config_owned_value() -> None: + store: Final = SettingsStore("general_settings") + store.load_yaml({"master_key": "os.environ/MASTER_KEY"}) + store.apply_runtime_values({"master_key": "sk-resolved"}) + + store["master_key"] = "sk-resolved" + + assert store["master_key"] == "sk-resolved" + assert store.source("master_key") == "config" + + @pytest.mark.timeout(10) def test_settings_store_clear_removes_every_key_the_config_file_does_not_own() -> None: store: Final = SettingsStore("general_settings") @@ -282,3 +296,12 @@ def test_settings_store_starts_with_an_unset_source() -> None: store: Final = SettingsStore("general_settings") assert store.source("unknown") == "unset" + + +def test_settings_store_still_accepts_a_write_to_a_key_the_config_does_not_own() -> None: + store: Final = SettingsStore("general_settings") + store.load_yaml({"allowed_ips": ["1.2.3.4"]}) + + store["max_parallel_requests"] = 7 + + assert store["max_parallel_requests"] == 7 diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 263300d12b1..0d59b022e38 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -3311,6 +3311,76 @@ async def test_load_config_rejects_malformed_role_permissions(tmp_path): await ProxyConfig().load_config(router=MagicMock(), config_file_path=str(config_file)) +def test_os_environ_resolution_leaves_the_config_layer_holding_the_reference(monkeypatch): + from litellm.proxy.proxy_server import ProxyConfig + + monkeypatch.setenv("PROOF_NESTED_SECRET", "sk-nested-value") + proxy_config: Final = ProxyConfig() + config: Final = { + "general_settings": { + "master_key": "os.environ/PROOF_NESTED_SECRET", + "coordination_redis": {"password": "os.environ/PROOF_NESTED_SECRET"}, + } + } + + proxy_config._load_yaml_settings_stores(config) + resolved: Final = proxy_config._check_for_os_environ_vars( + config=proxy_config._config_with_resolved_settings(config) + ) + + assert resolved["general_settings"]["coordination_redis"]["password"] == "sk-nested-value" + assert resolved["general_settings"]["master_key"] == "sk-nested-value" + assert proxy_config.settings.config_value("master_key") == "os.environ/PROOF_NESTED_SECRET" + assert proxy_config.settings.config_value("coordination_redis") == { + "password": "os.environ/PROOF_NESTED_SECRET" + } + + +def test_os_environ_resolution_reaches_dicts_nested_in_a_list(monkeypatch): + from litellm.proxy.proxy_server import ProxyConfig + + monkeypatch.setenv("PROOF_LIST_SECRET", "sk-list-value") + config: Final = {"model_list": [{"litellm_params": {"api_key": "os.environ/PROOF_LIST_SECRET"}}]} + + resolved: Final = ProxyConfig()._check_for_os_environ_vars(config=config) + + assert resolved["model_list"][0]["litellm_params"]["api_key"] == "sk-list-value" + + +@pytest.mark.parametrize("config_cache_size", ("not-a-number", "7")) +@pytest.mark.asyncio +async def test_db_reload_finishes_when_the_config_owns_a_setting_the_db_also_sets(monkeypatch, config_cache_size): + import litellm + from litellm.proxy import proxy_server as proxy_server_module + from litellm.proxy.proxy_server import ProxyConfig + + monkeypatch.setattr(litellm, "user_url_allowed_hosts", [], raising=False) + monkeypatch.setattr(proxy_server_module, "user_api_key_cache", MagicMock(), raising=False) + proxy_config: Final = ProxyConfig() + proxy_config.settings.load_yaml( + { + "store_prompts_in_spend_logs": "os.environ/PROOF_FLAG", + "store_model_in_db": "os.environ/PROOF_FLAG", + "user_api_key_cache_max_size": config_cache_size, + } + ) + monkeypatch.setattr(proxy_server_module, "general_settings", proxy_config.settings, raising=False) + + await proxy_config._update_general_settings( + { + "store_prompts_in_spend_logs": False, + "store_model_in_db": False, + "user_api_key_cache_max_size": 5, + "user_url_allowed_hosts": ["proof.example.com"], + } + ) + + assert litellm.user_url_allowed_hosts == ["proof.example.com"] + assert proxy_config.settings["store_prompts_in_spend_logs"] == "os.environ/PROOF_FLAG" + assert proxy_config.settings["store_model_in_db"] == "os.environ/PROOF_FLAG" + assert proxy_config.settings["user_api_key_cache_max_size"] == config_cache_size + + def test_max_ui_session_budget_default_is_one_dollar(): """LIT-4662: the dashboard session budget default is a product decision; the old 0.25 default locked admins out of auto router Test Connection and the diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index 9860d1bf94a..58201bd14ce 100644 --- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -2662,6 +2662,55 @@ def test_delete_allowed_ip_writes_deleted_audit_log(monkeypatch): app.dependency_overrides.pop(user_api_key_auth, None) +@pytest.mark.parametrize("route", ["/add/allowed_ip", "/delete/allowed_ip"]) +def test_allowed_ip_routes_refuse_a_config_owned_list_with_a_clear_400(route, monkeypatch): + from unittest.mock import AsyncMock, MagicMock + + import litellm.proxy.proxy_server as proxy_server_module + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.config_resolvers.settings_store import SettingsStore + + store = SettingsStore("general_settings") + store.load_yaml({"allowed_ips": ["203.0.113.77"]}) + saved = [] + + fake_prisma = MagicMock() + fake_prisma.db.litellm_auditlog.create = AsyncMock() + + async def _get_config(): + return {"general_settings": {"allowed_ips": ["203.0.113.77"]}} + + async def _save_config(new_config=None): + saved.append(new_config) + return new_config + + monkeypatch.setattr(proxy_server_module, "prisma_client", fake_prisma) + monkeypatch.setattr(proxy_server_module, "store_model_in_db", True) + monkeypatch.setattr(proxy_server_module, "general_settings", store) + monkeypatch.setattr(proxy_server_module.proxy_config, "get_config", _get_config) + monkeypatch.setattr(proxy_server_module.proxy_config, "save_config", _save_config) + + async def _admin_auth(): + return UserAPIKeyAuth( + user_id="config-admin", + api_key="hashed-admin-key", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + app.dependency_overrides[user_api_key_auth] = _admin_auth + try: + ip = "198.51.100.9" if route == "/add/allowed_ip" else "203.0.113.77" + resp = client.post(route, json={"ip": ip}) + + assert resp.status_code == 400, resp.text + assert "allowed_ips" in resp.text + assert list(store["allowed_ips"]) == ["203.0.113.77"] + assert saved == [] + finally: + app.dependency_overrides.pop(user_api_key_auth, None) + + def test_update_ui_theme_settings_writes_audit_log(mock_proxy_config, monkeypatch): """Updating the UI theme must be audited under ui_theme_config.""" from unittest.mock import AsyncMock, MagicMock From 3d2ec852155e01e13e80ce9330de106d23f1430c Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 18 Sep 2026 22:45:27 -0700 Subject: [PATCH 08/23] chore(proxy): keep the new config-owned refusals inside the LIT002 ceiling --- litellm/proxy/proxy_server.py | 4 ++-- .../proxy/ui_crud_endpoints/proxy_setting_endpoints.py | 10 +++++----- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index eb9889f1877..02c021333c3 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -5221,7 +5221,7 @@ class ProxyConfig: verbose_proxy_logger.warning("Maximum recursion depth (%s) reached while processing config.", max_depth) return config - return { + return { # mutable-ok: callers deep-copy and mutate this, and a mappingproxy cannot be deep-copied key: self._resolved_config_value(value=value, depth=depth, max_depth=max_depth) for key, value in config.items() } @@ -5230,7 +5230,7 @@ class ProxyConfig: if isinstance(value, dict): return self._check_for_os_environ_vars(config=value, depth=depth + 1, max_depth=max_depth) if isinstance(value, list): - return [ + return [ # mutable-ok: config values round-trip through json, where a tuple is not a list self._check_for_os_environ_vars(config=item, depth=depth + 1, max_depth=max_depth) if isinstance(item, dict) else item diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index b1f0ac7b35b..7c2abce60e2 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -492,13 +492,13 @@ async def get_allowed_ips(): def _store_allowed_ips(general_settings: MutableMapping[str, object], allowed_ips: Sequence[str]) -> None: try: - general_settings["allowed_ips"] = list(allowed_ips) + general_settings["allowed_ips"] = list(allowed_ips) # mutable-ok: compared against the file's own list except ConfigOwnedKeyError as owned: raise HTTPException( status_code=400, - detail={ + detail={ # mutable-ok: HTTPException serializes its detail as json "error": f"{owned.section} key '{owned.key}' is set in the config file and cannot be changed here", - "keys": [owned.key], + "keys": (owned.key,), "section": owned.section, "resolution": ( "edit the config file to change it, or remove it from the file to let the database own it" @@ -527,7 +527,7 @@ async def add_allowed_ip( if prisma_client is None: raise Exception("No DB Connected") - _allowed_ips: Final[Sequence[str]] = general_settings.get("allowed_ips") or [] + _allowed_ips: Final[Sequence[str]] = general_settings.get("allowed_ips") or () if ip_address.ip in _allowed_ips: raise HTTPException(status_code=400, detail="IP address already exists") _store_allowed_ips(general_settings, (*_allowed_ips, ip_address.ip)) @@ -584,7 +584,7 @@ async def delete_allowed_ip( proxy_config, ) - _allowed_ips: Final[Sequence[str]] = general_settings.get("allowed_ips") or [] + _allowed_ips: Final[Sequence[str]] = general_settings.get("allowed_ips") or () if ip_address.ip not in _allowed_ips: raise HTTPException(status_code=404, detail="IP address not found") _store_allowed_ips(general_settings, tuple(ip for ip in _allowed_ips if ip != ip_address.ip)) From c9158fcc12819fbe3b358b32cf693a506618ada9 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 18 Sep 2026 22:58:08 -0700 Subject: [PATCH 09/23] fix(proxy): keep a config-owned key's resolved value across a database reload Applying a database row dropped the runtime layer for every key the row carried, including keys the config file owns. Those runtime entries hold the env-resolved config values, so after a reload a key written as os.environ/ read back as that literal string. The store now keeps the runtime entry for a key the config owns and clears only the rest. Visible as store_model_in_db silently turning itself off: the reload read the raw reference, coerced it to False, and overwrote the resolved global. --- .../proxy/config_resolvers/settings_store.py | 7 ++++--- .../config_resolvers/test_settings_store.py | 12 +++++++++++- tests/test_litellm/proxy/test_proxy_server.py | 18 ++++++++++++++++++ 3 files changed, 33 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/config_resolvers/settings_store.py b/litellm/proxy/config_resolvers/settings_store.py index 98ebfdabd0f..345f00c35a5 100644 --- a/litellm/proxy/config_resolvers/settings_store.py +++ b/litellm/proxy/config_resolvers/settings_store.py @@ -117,12 +117,13 @@ class SettingsStore(MutableMapping[str, JsonValue]): self._deleted_runtime_keys = frozenset() def _clear_runtime_keys(self, keys: frozenset[str]) -> None: - if not keys: + stale: Final = frozenset(key for key in keys if not self.owned_by_config(key)) + if not stale: return self._runtime_values = MappingProxyType( - {key: value for key, value in self._runtime_values.items() if key not in keys} + {key: value for key, value in self._runtime_values.items() if key not in stale} ) - self._deleted_runtime_keys = self._deleted_runtime_keys - keys + self._deleted_runtime_keys = self._deleted_runtime_keys - stale def _keys(self) -> tuple[str, ...]: return tuple( diff --git a/tests/test_litellm/proxy/config_resolvers/test_settings_store.py b/tests/test_litellm/proxy/config_resolvers/test_settings_store.py index 18cbc3b0d6f..1182bcdce3c 100644 --- a/tests/test_litellm/proxy/config_resolvers/test_settings_store.py +++ b/tests/test_litellm/proxy/config_resolvers/test_settings_store.py @@ -86,7 +86,6 @@ def test_settings_store_keeps_unaffected_runtime_values_on_a_db_row_refresh() -> def test_settings_store_keeps_a_config_owned_key_when_a_db_row_disagrees() -> None: store: Final = SettingsStore("general_settings") store.load_yaml({"changed": "config"}) - store.apply_runtime_values({"changed": "resolved-config"}) store.apply_db_row("general_settings", {"changed": "database"}) @@ -94,6 +93,17 @@ def test_settings_store_keeps_a_config_owned_key_when_a_db_row_disagrees() -> No assert store.source("changed") == "config" +def test_settings_store_keeps_the_resolved_value_of_a_config_owned_key_across_a_db_row() -> None: + store: Final = SettingsStore("general_settings") + store.load_yaml({"changed": "os.environ/SETTING"}) + store.apply_runtime_values({"changed": "resolved-config"}) + + store.apply_db_row("general_settings", {"changed": "database"}) + + assert store["changed"] == "resolved-config" + assert store.source("changed") == "config" + + def test_settings_store_removes_only_runtime_values_affected_by_a_cleared_db_row() -> None: store: Final = SettingsStore("general_settings") store.load_yaml({"template": "os.environ/SETTING"}) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 0d59b022e38..d5719a8382c 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -3381,6 +3381,24 @@ async def test_db_reload_finishes_when_the_config_owns_a_setting_the_db_also_set assert proxy_config.settings["user_api_key_cache_max_size"] == config_cache_size +@pytest.mark.asyncio +async def test_db_reload_keeps_the_resolved_value_of_a_config_owned_env_reference(monkeypatch): + from litellm.proxy import proxy_server as proxy_server_module + from litellm.proxy.proxy_server import ProxyConfig + + monkeypatch.setattr(proxy_server_module, "user_api_key_cache", MagicMock(), raising=False) + monkeypatch.setattr(proxy_server_module, "store_model_in_db", True, raising=False) + proxy_config: Final = ProxyConfig() + proxy_config.settings.load_yaml({"store_model_in_db": "os.environ/PROOF_STORE_FLAG"}) + proxy_config.settings.apply_runtime_values({"store_model_in_db": True}) + monkeypatch.setattr(proxy_server_module, "general_settings", proxy_config.settings, raising=False) + + await proxy_config._update_general_settings({"store_model_in_db": True}) + + assert proxy_config.settings["store_model_in_db"] is True + assert proxy_server_module.store_model_in_db is True + + def test_max_ui_session_budget_default_is_one_dollar(): """LIT-4662: the dashboard session budget default is a product decision; the old 0.25 default locked admins out of auto router Test Connection and the From bdbe265c7020a6533b8dd729708ff01725bac4ad Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 03:14:39 -0700 Subject: [PATCH 10/23] fix(proxy): /key/bulk_update writes only the fields each item carries A bulk item that carried only tags reached the DB with max_budget, team_id, and budget_id as explicit nulls, wiping the key's budget and detaching it from its team. The per-key update is now built from the fields the item actually set, so a field left out keeps its value and an explicit null still clears it, the same as /key/update. Items carrying a field the bulk path cannot apply (object_permission and the like) are rejected with 422 instead of being silently dropped. --- litellm/proxy/_lazy_openapi_snapshot.json | 2 +- .../key_management_endpoints.py | 14 ++-- .../key_management_endpoints.py | 4 +- .../test_key_management_endpoints.py | 83 +++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 +- 5 files changed, 96 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 73ea8cf1991..18849ef5b64 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -19622,7 +19622,7 @@ } } }, - "description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n " + "description": "\nUnified rate-limit error.\n\nEvery rate-limit condition surfaced by litellm \u2014 whether it originated from\nan upstream LLM provider, a vendor batch endpoint, or one of litellm's own\nproxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\nmax-iterations, etc.) \u2014 is raised as an instance of this class.\n\nThe :attr:`category` attribute lets callers distinguish the source. See\n:class:`RateLimitErrorCategory` for the available values.\n" }, "500": { "content": { diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 6c195d713c8..f28101f7072 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3517,7 +3517,10 @@ async def bulk_update_keys( - max_budget: Optional[float] - Max budget for key - team_id: Optional[str] - Team ID associated with key - tags: Optional[List[str]] - Tags for organizing keys - + + Only the fields an item carries are written: a field left out keeps its current value and an + explicit null clears it, the same as /key/update. An item carrying any other field is rejected with 422. + Returns: - total_requested: int - Total number of keys requested for update - successful_updates: List[SuccessfulKeyUpdate] - List of successfully updated keys with their updated info @@ -3586,15 +3589,8 @@ async def bulk_update_keys( for key_update_item in data.keys: try: - update_key_request = UpdateKeyRequest( - key=key_update_item.key, - budget_id=key_update_item.budget_id, - max_budget=key_update_item.max_budget, - team_id=key_update_item.team_id, - tags=key_update_item.tags, - ) updated_key_info = await _process_single_key_update( - update_key_request=update_key_request, + update_key_request=UpdateKeyRequest.model_validate(key_update_item.model_dump(exclude_unset=True)), user_api_key_dict=user_api_key_dict, litellm_changed_by=litellm_changed_by, prisma_client=prisma_client, diff --git a/litellm/types/proxy/management_endpoints/key_management_endpoints.py b/litellm/types/proxy/management_endpoints/key_management_endpoints.py index 63bbaa5ba4e..3e193956d30 100644 --- a/litellm/types/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/types/proxy/management_endpoints/key_management_endpoints.py @@ -25,7 +25,9 @@ class KeySearchWhere(TypedDict): class BulkUpdateKeyRequestItem(BaseModel): - """Individual key update request item""" + """One /key/bulk_update item; only the fields it carries are written, and unknown fields are rejected.""" + + model_config = ConfigDict(extra="forbid") key: str # Key identifier (token) budget_id: str | None = None # Budget ID associated with the key diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index c852307b051..2f0f605b839 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -7097,6 +7097,89 @@ async def test_list_key_helper_applies_search_to_prisma_where(): assert _search_clause("key-id-123", "key-id-123") in where["AND"], f"search not in Prisma where: {where}" +_BULK_UPDATE_TOKEN: Final = "1f2e3d4c5b6a79880123456789abcdef0123456789abcdef0123456789abcdef" + + +async def _bulk_update_one_key(monkeypatch, item_payload: Mapping[str, object]) -> Mapping[str, object]: + """Runs /key/bulk_update with one item against a budgeted team key and returns the row written to the DB.""" + from litellm.proxy.management_endpoints.key_management_endpoints import bulk_update_keys + from litellm.types.proxy.management_endpoints.key_management_endpoints import BulkUpdateKeyRequest + + key_in_db = LiteLLM_VerificationToken( + token=_BULK_UPDATE_TOKEN, user_id="test-user", team_id="team-1", max_budget=100.0, budget_id="budget-1" + ) + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=key_in_db) + mock_prisma_client.update_data = AsyncMock(return_value={"data": {"token": _BULK_UPDATE_TOKEN}}) + _setup_update_key_mocks(monkeypatch, mock_prisma_client) + + with ( + patch( # test-quality-ok: the handler reads the cache and hook singletons from module globals, no injection seam + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: the permission check is a classmethod the handler calls directly, no injection seam + "litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: the audit hook is a classmethod the handler calls directly, no injection seam + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook", + new_callable=AsyncMock, + ), + ): + response = await bulk_update_keys( + data=BulkUpdateKeyRequest.model_validate({"keys": [{"key": _BULK_UPDATE_TOKEN, **item_payload}]}), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + litellm_changed_by=None, + ) + + assert response.failed_updates == [] + return mock_prisma_client.update_data.call_args.kwargs["data"] + + +@pytest.mark.asyncio +async def test_bulk_update_keys_item_without_a_field_leaves_that_column_alone(monkeypatch): + """A tags-only item used to reach the DB with max_budget, team_id, and budget_id as explicit + nulls, so tagging a key wiped its budget and detached it from its team.""" + written = await _bulk_update_one_key(monkeypatch, {"tags": ["team-a"]}) + + assert written["metadata"]["tags"] == ["team-a"] + assert not {"max_budget", "team_id", "budget_id"} & written.keys() + + +@pytest.mark.asyncio +async def test_bulk_update_keys_explicit_null_still_clears_the_field(monkeypatch): + """Sending `"max_budget": null` on an item is a request to remove the budget, as on /key/update.""" + written = await _bulk_update_one_key(monkeypatch, {"max_budget": None}) + + assert written["max_budget"] is None + assert not {"team_id", "budget_id"} & written.keys() + + +def test_bulk_update_keys_rejects_a_field_the_bulk_path_cannot_apply(): + """`object_permission` used to be accepted with 200 and dropped, leaving an item that carried + nothing but the key, so the call wiped the key's budget instead of granting the permission.""" + from fastapi import FastAPI + + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.management_endpoints.key_management_endpoints import router + + test_app = FastAPI() + test_app.include_router(router) + test_app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ) + response = TestClient(test_app).post( + "/key/bulk_update", + json={"keys": [{"key": _BULK_UPDATE_TOKEN, "object_permission": {"vector_stores": ["vs-1"]}}]}, + ) + + assert response.status_code == 422, response.text + assert "object_permission" in response.text + + @pytest.mark.asyncio async def test_generate_key_negative_max_budget(): """ diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index a5ede63bf7f..a9174603290 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -7694,6 +7694,9 @@ export interface paths { * - team_id: Optional[str] - Team ID associated with key * - tags: Optional[List[str]] - Tags for organizing keys * + * Only the fields an item carries are written: a field left out keeps its current value and an + * explicit null clears it, the same as /key/update. An item carrying any other field is rejected with 422. + * * Returns: * - total_requested: int - Total number of keys requested for update * - successful_updates: List[SuccessfulKeyUpdate] - List of successfully updated keys with their updated info @@ -25237,7 +25240,7 @@ export interface components { }; /** * BulkUpdateKeyRequestItem - * @description Individual key update request item + * @description One /key/bulk_update item; only the fields it carries are written, and unknown fields are rejected. */ BulkUpdateKeyRequestItem: { /** Budget Id */ From ad4da0f8e6b1e22ec7ecd8fb0b0e3ad138b807c2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 03:28:42 -0700 Subject: [PATCH 11/23] chore(proxy): regenerate the lazy OpenAPI snapshot on Python 3.12 and drop a test helper docstring --- litellm/proxy/_lazy_openapi_snapshot.json | 2 +- .../proxy/management_endpoints/test_key_management_endpoints.py | 1 - 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 18849ef5b64..73ea8cf1991 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -19622,7 +19622,7 @@ } } }, - "description": "\nUnified rate-limit error.\n\nEvery rate-limit condition surfaced by litellm \u2014 whether it originated from\nan upstream LLM provider, a vendor batch endpoint, or one of litellm's own\nproxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\nmax-iterations, etc.) \u2014 is raised as an instance of this class.\n\nThe :attr:`category` attribute lets callers distinguish the source. See\n:class:`RateLimitErrorCategory` for the available values.\n" + "description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n " }, "500": { "content": { diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 2f0f605b839..9d023e129aa 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -7101,7 +7101,6 @@ _BULK_UPDATE_TOKEN: Final = "1f2e3d4c5b6a79880123456789abcdef0123456789abcdef012 async def _bulk_update_one_key(monkeypatch, item_payload: Mapping[str, object]) -> Mapping[str, object]: - """Runs /key/bulk_update with one item against a budgeted team key and returns the row written to the DB.""" from litellm.proxy.management_endpoints.key_management_endpoints import bulk_update_keys from litellm.types.proxy.management_endpoints.key_management_endpoints import BulkUpdateKeyRequest From b746ac44563589a0e2b407b064470c2ee18b4b27 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 03:44:07 -0700 Subject: [PATCH 12/23] fix(proxy): accept object_permission on /key/bulk_update items instead of 422 --- .../key_management_endpoints.py | 3 +- .../key_management_endpoints.py | 12 ++++-- .../test_key_management_endpoints.py | 41 +++++++++---------- ui/litellm-dashboard/src/lib/http/schema.d.ts | 6 ++- 4 files changed, 34 insertions(+), 28 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index f28101f7072..959fee7b010 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3517,9 +3517,10 @@ async def bulk_update_keys( - max_budget: Optional[float] - Max budget for key - team_id: Optional[str] - Team ID associated with key - tags: Optional[List[str]] - Tags for organizing keys + - object_permission: Optional[LiteLLM_ObjectPermissionBase] - key-specific object permission, as on /key/update Only the fields an item carries are written: a field left out keeps its current value and an - explicit null clears it, the same as /key/update. An item carrying any other field is rejected with 422. + explicit null clears it, the same as /key/update. Returns: - total_requested: int - Total number of keys requested for update diff --git a/litellm/types/proxy/management_endpoints/key_management_endpoints.py b/litellm/types/proxy/management_endpoints/key_management_endpoints.py index 3e193956d30..001bc3c0d51 100644 --- a/litellm/types/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/types/proxy/management_endpoints/key_management_endpoints.py @@ -5,7 +5,12 @@ from pydantic import BaseModel, ConfigDict, model_validator from typing_extensions import ReadOnly, TypedDict from litellm.models.verification_token import LiteLLM_VerificationToken -from litellm.proxy._types import GenerateKeyRequest, RegenerateKeyRequest, UpdateKeyRequest +from litellm.proxy._types import ( + GenerateKeyRequest, + LiteLLM_ObjectPermissionBase, + RegenerateKeyRequest, + UpdateKeyRequest, +) from litellm.types.llms.base import LiteLLMPydanticObjectBase from litellm.types.proxy.management_endpoints.internal_user_endpoints import InsensitiveContains @@ -25,15 +30,14 @@ class KeySearchWhere(TypedDict): class BulkUpdateKeyRequestItem(BaseModel): - """One /key/bulk_update item; only the fields it carries are written, and unknown fields are rejected.""" - - model_config = ConfigDict(extra="forbid") + """One /key/bulk_update item; only the fields it carries are written.""" key: str # Key identifier (token) budget_id: str | None = None # Budget ID associated with the key max_budget: float | None = None # Max budget for key team_id: str | None = None # Team ID associated with key tags: list[str] | None = None # Tags for organizing keys + object_permission: LiteLLM_ObjectPermissionBase | None = None class BulkUpdateKeyRequest(BaseModel): diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 9d023e129aa..1bf5018900a 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -7100,7 +7100,7 @@ async def test_list_key_helper_applies_search_to_prisma_where(): _BULK_UPDATE_TOKEN: Final = "1f2e3d4c5b6a79880123456789abcdef0123456789abcdef0123456789abcdef" -async def _bulk_update_one_key(monkeypatch, item_payload: Mapping[str, object]) -> Mapping[str, object]: +async def _bulk_update_one_key(monkeypatch, item_payload: Mapping[str, object]) -> AsyncMock: from litellm.proxy.management_endpoints.key_management_endpoints import bulk_update_keys from litellm.types.proxy.management_endpoints.key_management_endpoints import BulkUpdateKeyRequest @@ -7109,6 +7109,10 @@ async def _bulk_update_one_key(monkeypatch, item_payload: Mapping[str, object]) ) mock_prisma_client = AsyncMock() mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=key_in_db) + mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock( + return_value=MagicMock(object_permission_id="objperm-bulk") + ) mock_prisma_client.update_data = AsyncMock(return_value={"data": {"token": _BULK_UPDATE_TOKEN}}) _setup_update_key_mocks(monkeypatch, mock_prisma_client) @@ -7135,14 +7139,18 @@ async def _bulk_update_one_key(monkeypatch, item_payload: Mapping[str, object]) ) assert response.failed_updates == [] - return mock_prisma_client.update_data.call_args.kwargs["data"] + return mock_prisma_client + + +def _written_key_row(prisma: AsyncMock) -> Mapping[str, object]: + return prisma.update_data.call_args.kwargs["data"] @pytest.mark.asyncio async def test_bulk_update_keys_item_without_a_field_leaves_that_column_alone(monkeypatch): """A tags-only item used to reach the DB with max_budget, team_id, and budget_id as explicit nulls, so tagging a key wiped its budget and detached it from its team.""" - written = await _bulk_update_one_key(monkeypatch, {"tags": ["team-a"]}) + written = _written_key_row(await _bulk_update_one_key(monkeypatch, {"tags": ["team-a"]})) assert written["metadata"]["tags"] == ["team-a"] assert not {"max_budget", "team_id", "budget_id"} & written.keys() @@ -7151,32 +7159,23 @@ async def test_bulk_update_keys_item_without_a_field_leaves_that_column_alone(mo @pytest.mark.asyncio async def test_bulk_update_keys_explicit_null_still_clears_the_field(monkeypatch): """Sending `"max_budget": null` on an item is a request to remove the budget, as on /key/update.""" - written = await _bulk_update_one_key(monkeypatch, {"max_budget": None}) + written = _written_key_row(await _bulk_update_one_key(monkeypatch, {"max_budget": None})) assert written["max_budget"] is None assert not {"team_id", "budget_id"} & written.keys() -def test_bulk_update_keys_rejects_a_field_the_bulk_path_cannot_apply(): +@pytest.mark.asyncio +async def test_bulk_update_keys_object_permission_is_granted_not_dropped(monkeypatch): """`object_permission` used to be accepted with 200 and dropped, leaving an item that carried nothing but the key, so the call wiped the key's budget instead of granting the permission.""" - from fastapi import FastAPI + prisma = await _bulk_update_one_key(monkeypatch, {"object_permission": {"vector_stores": ["vs-1"]}}) - from litellm.proxy.auth.user_api_key_auth import user_api_key_auth - from litellm.proxy.management_endpoints.key_management_endpoints import router - - test_app = FastAPI() - test_app.include_router(router) - test_app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" - ) - response = TestClient(test_app).post( - "/key/bulk_update", - json={"keys": [{"key": _BULK_UPDATE_TOKEN, "object_permission": {"vector_stores": ["vs-1"]}}]}, - ) - - assert response.status_code == 422, response.text - assert "object_permission" in response.text + upserted = prisma.db.litellm_objectpermissiontable.upsert.call_args.kwargs["data"]["create"] + assert upserted["vector_stores"] == ["vs-1"] + written = _written_key_row(prisma) + assert written["object_permission_id"] == "objperm-bulk" + assert not {"max_budget", "team_id", "budget_id"} & written.keys() @pytest.mark.asyncio diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index a9174603290..4ada2c3372b 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -7693,9 +7693,10 @@ export interface paths { * - max_budget: Optional[float] - Max budget for key * - team_id: Optional[str] - Team ID associated with key * - tags: Optional[List[str]] - Tags for organizing keys + * - object_permission: Optional[LiteLLM_ObjectPermissionBase] - key-specific object permission, as on /key/update * * Only the fields an item carries are written: a field left out keeps its current value and an - * explicit null clears it, the same as /key/update. An item carrying any other field is rejected with 422. + * explicit null clears it, the same as /key/update. * * Returns: * - total_requested: int - Total number of keys requested for update @@ -25240,7 +25241,7 @@ export interface components { }; /** * BulkUpdateKeyRequestItem - * @description One /key/bulk_update item; only the fields it carries are written, and unknown fields are rejected. + * @description One /key/bulk_update item; only the fields it carries are written. */ BulkUpdateKeyRequestItem: { /** Budget Id */ @@ -25249,6 +25250,7 @@ export interface components { key: string; /** Max Budget */ max_budget?: number | null; + object_permission?: components["schemas"]["LiteLLM_ObjectPermissionBase"] | null; /** Tags */ tags?: string[] | null; /** Team Id */ From f982d3e0469590fe8a1d05843fd6d9a04dfbc56a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 03:55:48 -0700 Subject: [PATCH 13/23] docs(proxy): state /key/bulk_update null handling as /key/update parity --- .../proxy/management_endpoints/key_management_endpoints.py | 4 ++-- ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 959fee7b010..033ada2c50d 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3519,8 +3519,8 @@ async def bulk_update_keys( - tags: Optional[List[str]] - Tags for organizing keys - object_permission: Optional[LiteLLM_ObjectPermissionBase] - key-specific object permission, as on /key/update - Only the fields an item carries are written: a field left out keeps its current value and an - explicit null clears it, the same as /key/update. + Only the fields an item carries are written: a field left out keeps its current value, and a field + sent explicitly, null included, is applied exactly as /key/update applies it. Returns: - total_requested: int - Total number of keys requested for update diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 4ada2c3372b..e1a4f4a743d 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -7695,8 +7695,8 @@ export interface paths { * - tags: Optional[List[str]] - Tags for organizing keys * - object_permission: Optional[LiteLLM_ObjectPermissionBase] - key-specific object permission, as on /key/update * - * Only the fields an item carries are written: a field left out keeps its current value and an - * explicit null clears it, the same as /key/update. + * Only the fields an item carries are written: a field left out keeps its current value, and a field + * sent explicitly, null included, is applied exactly as /key/update applies it. * * Returns: * - total_requested: int - Total number of keys requested for update From e0db862781378ff27468845afbb44f3df6ebaada Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 04:31:55 -0700 Subject: [PATCH 14/23] fix(cost): bill DeepSeek V4.1 Flash and V4 Pro at their off-peak rates outside peak hours DeepSeek charges half the listed rate outside 01:00-04:00 and 06:00-10:00 UTC Monday to Friday, so every deepseek-flash, deepseek-v4-flash, deepseek-v4-flash-vision-exp, and deepseek-v4-pro entry now carries an off_peak_pricing block with those windows and the halved input, output, and cache-hit rates. The generated cost map schema picks up the block, and the regression tests pin the peak and off-peak cost of one call at fixed moments. --- ...odel_prices_and_context_window_backup.json | 224 ++++++++++++++++++ model_prices_and_context_window.json | 224 ++++++++++++++++++ model_prices_and_context_window.schema.json | 108 +++++++++ .../deepseek/test_deepseek_cost_calculator.py | 70 ++++++ .../test_litellm/test_model_prices_schema.py | 41 ++++ tests/test_litellm/test_utils.py | 32 +++ 6 files changed, 699 insertions(+) create mode 100644 tests/test_litellm/llms/deepseek/test_deepseek_cost_calculator.py diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 32619a86247..4dbf0337894 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -59567,6 +59567,34 @@ "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", + "off_peak_pricing": { + "cache_read_input_token_cost": 3e-09, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "windows": [ + { + "hours_utc": [ + "00:00-01:00", + "04:00-06:00", + "10:00-00:00" + ], + "weekdays": [ + 1, + 2, + 3, + 4, + 5 + ] + }, + { + "hours_utc": "00:00-00:00", + "weekdays": [ + 6, + 7 + ] + } + ] + }, "output_cost_per_token": 1.2e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ @@ -59593,6 +59621,34 @@ "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", + "off_peak_pricing": { + "cache_read_input_token_cost": 3e-09, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "windows": [ + { + "hours_utc": [ + "00:00-01:00", + "04:00-06:00", + "10:00-00:00" + ], + "weekdays": [ + 1, + 2, + 3, + 4, + 5 + ] + }, + { + "hours_utc": "00:00-00:00", + "weekdays": [ + 6, + 7 + ] + } + ] + }, "output_cost_per_token": 1.2e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ @@ -59619,6 +59675,34 @@ "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", + "off_peak_pricing": { + "cache_read_input_token_cost": 3e-09, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "windows": [ + { + "hours_utc": [ + "00:00-01:00", + "04:00-06:00", + "10:00-00:00" + ], + "weekdays": [ + 1, + 2, + 3, + 4, + 5 + ] + }, + { + "hours_utc": "00:00-00:00", + "weekdays": [ + 6, + 7 + ] + } + ] + }, "output_cost_per_token": 1.2e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ @@ -59645,6 +59729,34 @@ "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", + "off_peak_pricing": { + "cache_read_input_token_cost": 2.2e-08, + "input_cost_per_token": 6.6e-07, + "output_cost_per_token": 1.98e-06, + "windows": [ + { + "hours_utc": [ + "00:00-01:00", + "04:00-06:00", + "10:00-00:00" + ], + "weekdays": [ + 1, + 2, + 3, + 4, + 5 + ] + }, + { + "hours_utc": "00:00-00:00", + "weekdays": [ + 6, + 7 + ] + } + ] + }, "output_cost_per_token": 3.96e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ @@ -59671,6 +59783,34 @@ "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", + "off_peak_pricing": { + "cache_read_input_token_cost": 3e-09, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "windows": [ + { + "hours_utc": [ + "00:00-01:00", + "04:00-06:00", + "10:00-00:00" + ], + "weekdays": [ + 1, + 2, + 3, + 4, + 5 + ] + }, + { + "hours_utc": "00:00-00:00", + "weekdays": [ + 6, + 7 + ] + } + ] + }, "output_cost_per_token": 1.2e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ @@ -59697,6 +59837,34 @@ "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", + "off_peak_pricing": { + "cache_read_input_token_cost": 3e-09, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "windows": [ + { + "hours_utc": [ + "00:00-01:00", + "04:00-06:00", + "10:00-00:00" + ], + "weekdays": [ + 1, + 2, + 3, + 4, + 5 + ] + }, + { + "hours_utc": "00:00-00:00", + "weekdays": [ + 6, + 7 + ] + } + ] + }, "output_cost_per_token": 1.2e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ @@ -59723,6 +59891,34 @@ "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", + "off_peak_pricing": { + "cache_read_input_token_cost": 3e-09, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "windows": [ + { + "hours_utc": [ + "00:00-01:00", + "04:00-06:00", + "10:00-00:00" + ], + "weekdays": [ + 1, + 2, + 3, + 4, + 5 + ] + }, + { + "hours_utc": "00:00-00:00", + "weekdays": [ + 6, + 7 + ] + } + ] + }, "output_cost_per_token": 1.2e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ @@ -59749,6 +59945,34 @@ "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", + "off_peak_pricing": { + "cache_read_input_token_cost": 2.2e-08, + "input_cost_per_token": 6.6e-07, + "output_cost_per_token": 1.98e-06, + "windows": [ + { + "hours_utc": [ + "00:00-01:00", + "04:00-06:00", + "10:00-00:00" + ], + "weekdays": [ + 1, + 2, + 3, + 4, + 5 + ] + }, + { + "hours_utc": "00:00-00:00", + "weekdays": [ + 6, + 7 + ] + } + ] + }, "output_cost_per_token": 3.96e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 32619a86247..4dbf0337894 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -59567,6 +59567,34 @@ "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", + "off_peak_pricing": { + "cache_read_input_token_cost": 3e-09, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "windows": [ + { + "hours_utc": [ + "00:00-01:00", + "04:00-06:00", + "10:00-00:00" + ], + "weekdays": [ + 1, + 2, + 3, + 4, + 5 + ] + }, + { + "hours_utc": "00:00-00:00", + "weekdays": [ + 6, + 7 + ] + } + ] + }, "output_cost_per_token": 1.2e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ @@ -59593,6 +59621,34 @@ "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", + "off_peak_pricing": { + "cache_read_input_token_cost": 3e-09, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "windows": [ + { + "hours_utc": [ + "00:00-01:00", + "04:00-06:00", + "10:00-00:00" + ], + "weekdays": [ + 1, + 2, + 3, + 4, + 5 + ] + }, + { + "hours_utc": "00:00-00:00", + "weekdays": [ + 6, + 7 + ] + } + ] + }, "output_cost_per_token": 1.2e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ @@ -59619,6 +59675,34 @@ "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", + "off_peak_pricing": { + "cache_read_input_token_cost": 3e-09, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "windows": [ + { + "hours_utc": [ + "00:00-01:00", + "04:00-06:00", + "10:00-00:00" + ], + "weekdays": [ + 1, + 2, + 3, + 4, + 5 + ] + }, + { + "hours_utc": "00:00-00:00", + "weekdays": [ + 6, + 7 + ] + } + ] + }, "output_cost_per_token": 1.2e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ @@ -59645,6 +59729,34 @@ "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", + "off_peak_pricing": { + "cache_read_input_token_cost": 2.2e-08, + "input_cost_per_token": 6.6e-07, + "output_cost_per_token": 1.98e-06, + "windows": [ + { + "hours_utc": [ + "00:00-01:00", + "04:00-06:00", + "10:00-00:00" + ], + "weekdays": [ + 1, + 2, + 3, + 4, + 5 + ] + }, + { + "hours_utc": "00:00-00:00", + "weekdays": [ + 6, + 7 + ] + } + ] + }, "output_cost_per_token": 3.96e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ @@ -59671,6 +59783,34 @@ "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", + "off_peak_pricing": { + "cache_read_input_token_cost": 3e-09, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "windows": [ + { + "hours_utc": [ + "00:00-01:00", + "04:00-06:00", + "10:00-00:00" + ], + "weekdays": [ + 1, + 2, + 3, + 4, + 5 + ] + }, + { + "hours_utc": "00:00-00:00", + "weekdays": [ + 6, + 7 + ] + } + ] + }, "output_cost_per_token": 1.2e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ @@ -59697,6 +59837,34 @@ "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", + "off_peak_pricing": { + "cache_read_input_token_cost": 3e-09, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "windows": [ + { + "hours_utc": [ + "00:00-01:00", + "04:00-06:00", + "10:00-00:00" + ], + "weekdays": [ + 1, + 2, + 3, + 4, + 5 + ] + }, + { + "hours_utc": "00:00-00:00", + "weekdays": [ + 6, + 7 + ] + } + ] + }, "output_cost_per_token": 1.2e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ @@ -59723,6 +59891,34 @@ "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", + "off_peak_pricing": { + "cache_read_input_token_cost": 3e-09, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 6e-07, + "windows": [ + { + "hours_utc": [ + "00:00-01:00", + "04:00-06:00", + "10:00-00:00" + ], + "weekdays": [ + 1, + 2, + 3, + 4, + 5 + ] + }, + { + "hours_utc": "00:00-00:00", + "weekdays": [ + 6, + 7 + ] + } + ] + }, "output_cost_per_token": 1.2e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ @@ -59749,6 +59945,34 @@ "max_output_tokens": 393216, "max_tokens": 393216, "mode": "chat", + "off_peak_pricing": { + "cache_read_input_token_cost": 2.2e-08, + "input_cost_per_token": 6.6e-07, + "output_cost_per_token": 1.98e-06, + "windows": [ + { + "hours_utc": [ + "00:00-01:00", + "04:00-06:00", + "10:00-00:00" + ], + "weekdays": [ + 1, + 2, + 3, + 4, + 5 + ] + }, + { + "hours_utc": "00:00-00:00", + "weekdays": [ + 6, + 7 + ] + } + ] + }, "output_cost_per_token": 3.96e-06, "source": "https://api-docs.deepseek.com/quick_start/pricing", "supported_endpoints": [ diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index 8d79c560175..44b2569defd 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -457,6 +457,114 @@ "type": "number", "minimum": 0 }, + "off_peak_pricing": { + "type": "object", + "description": "Rates that replace the same-named base fields while the request falls inside the stated UTC windows.", + "properties": { + "hours_utc": { + "description": "UTC \"HH:MM-HH:MM\" window, or a list of them; a window may wrap past midnight.", + "oneOf": [ + { + "type": "string", + "pattern": "^([01]\\d|2[0-3]):[0-5]\\d-([01]\\d|2[0-3]):[0-5]\\d$" + }, + { + "type": "array", + "items": { + "type": "string", + "pattern": "^([01]\\d|2[0-3]):[0-5]\\d-([01]\\d|2[0-3]):[0-5]\\d$" + }, + "minItems": 1 + } + ] + }, + "windows": { + "type": "array", + "items": { + "type": "object", + "properties": { + "hours_utc": { + "description": "UTC \"HH:MM-HH:MM\" window, or a list of them; a window may wrap past midnight.", + "oneOf": [ + { + "type": "string", + "pattern": "^([01]\\d|2[0-3]):[0-5]\\d-([01]\\d|2[0-3]):[0-5]\\d$" + }, + { + "type": "array", + "items": { + "type": "string", + "pattern": "^([01]\\d|2[0-3]):[0-5]\\d-([01]\\d|2[0-3]):[0-5]\\d$" + }, + "minItems": 1 + } + ] + }, + "weekdays": { + "type": "array", + "description": "ISO-8601 weekday numbers (1 = Monday .. 7 = Sunday) or English day names the window applies on.", + "items": { + "oneOf": [ + { + "type": "integer", + "minimum": 1, + "maximum": 7 + }, + { + "type": "string", + "pattern": "(?i)^(mon|monday|tue|tues|tuesday|wed|wednesday|thu|thur|thurs|thursday|fri|friday|sat|saturday|sun|sunday)$" + } + ] + }, + "minItems": 1 + } + }, + "required": [ + "hours_utc" + ], + "additionalProperties": false + }, + "minItems": 1 + }, + "weekday_timezone": { + "type": "string", + "description": "IANA zone the weekdays of each window are read on; defaults to UTC." + }, + "input_cost_per_token": { + "type": "number", + "minimum": 0 + }, + "output_cost_per_token": { + "type": "number", + "minimum": 0 + }, + "output_cost_per_reasoning_token": { + "type": "number", + "minimum": 0 + }, + "cache_read_input_token_cost": { + "type": "number", + "minimum": 0 + }, + "cache_creation_input_token_cost": { + "type": "number", + "minimum": 0 + } + }, + "anyOf": [ + { + "required": [ + "hours_utc" + ] + }, + { + "required": [ + "windows" + ] + } + ], + "additionalProperties": false + }, "output_cost_per_audio_token": { "type": "number", "minimum": 0 diff --git a/tests/test_litellm/llms/deepseek/test_deepseek_cost_calculator.py b/tests/test_litellm/llms/deepseek/test_deepseek_cost_calculator.py new file mode 100644 index 00000000000..c3a4cdad0ac --- /dev/null +++ b/tests/test_litellm/llms/deepseek/test_deepseek_cost_calculator.py @@ -0,0 +1,70 @@ +from datetime import datetime, timezone +from typing import Final + +import pytest + +import litellm +from litellm._internal_context import pinned_billing_time +from litellm.types.utils import ModelResponse, PromptTokensDetailsWrapper, Usage + +PEAK_MOMENTS: Final = ( + pytest.param(datetime(2026, 9, 22, 8, 0, tzinfo=timezone.utc), id="tuesday-08:00"), + pytest.param(datetime(2026, 9, 25, 9, 59, tzinfo=timezone.utc), id="friday-09:59"), + pytest.param(datetime(2026, 9, 21, 1, 0, tzinfo=timezone.utc), id="monday-01:00"), +) +OFF_PEAK_MOMENTS: Final = ( + pytest.param(datetime(2026, 9, 26, 2, 0, tzinfo=timezone.utc), id="saturday-02:00"), + pytest.param(datetime(2026, 9, 27, 8, 0, tzinfo=timezone.utc), id="sunday-08:00"), + pytest.param(datetime(2026, 9, 21, 0, 30, tzinfo=timezone.utc), id="monday-00:30"), + pytest.param(datetime(2026, 9, 23, 5, 0, tzinfo=timezone.utc), id="wednesday-05:00"), + pytest.param(datetime(2026, 9, 24, 10, 0, tzinfo=timezone.utc), id="thursday-10:00"), + pytest.param(datetime(2026, 9, 22, 12, 0, tzinfo=timezone.utc), id="tuesday-12:00"), +) +PEAK_COST_PER_MILLION_IN_AND_OUT_WITH_400K_CACHE_HITS: Final = { + "deepseek-flash": 1.3824, + "deepseek-v4-pro": 4.7696, +} + + +def one_million_in_and_out_with_400k_cache_hits(model: str) -> ModelResponse: + return ModelResponse( + model=model, + usage=Usage( + prompt_tokens=1_000_000, + completion_tokens=1_000_000, + total_tokens=2_000_000, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=400_000), + ), + ) + + +def deepseek_cost_at(model: str, moment: datetime) -> float: + with pinned_billing_time(moment): + return litellm.completion_cost( + completion_response=one_million_in_and_out_with_400k_cache_hits(model), + model=model, + custom_llm_provider="deepseek", + ) + + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize(("model", "peak_cost"), PEAK_COST_PER_MILLION_IN_AND_OUT_WITH_400K_CACHE_HITS.items()) +@pytest.mark.parametrize("moment", PEAK_MOMENTS) +def test_deepseek_bills_the_listed_rate_during_weekday_peak_hours(model: str, peak_cost: float, moment: datetime): + assert deepseek_cost_at(model, moment) == pytest.approx(peak_cost) + + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize(("model", "peak_cost"), PEAK_COST_PER_MILLION_IN_AND_OUT_WITH_400K_CACHE_HITS.items()) +@pytest.mark.parametrize("moment", OFF_PEAK_MOMENTS) +def test_deepseek_bills_half_the_listed_rate_off_peak(model: str, peak_cost: float, moment: datetime): + assert deepseek_cost_at(model, moment) == pytest.approx(peak_cost / 2) + + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize("alias", ("deepseek-v4-flash", "deepseek-v4-flash-vision-exp", "deepseek/deepseek-flash")) +def test_deepseek_flash_aliases_follow_the_same_off_peak_schedule(alias: str): + saturday: Final = datetime(2026, 9, 26, 2, 0, tzinfo=timezone.utc) + assert deepseek_cost_at(alias, saturday) == pytest.approx( + PEAK_COST_PER_MILLION_IN_AND_OUT_WITH_400K_CACHE_HITS["deepseek-flash"] / 2 + ) diff --git a/tests/test_litellm/test_model_prices_schema.py b/tests/test_litellm/test_model_prices_schema.py index 918aff806c1..052278631e2 100644 --- a/tests/test_litellm/test_model_prices_schema.py +++ b/tests/test_litellm/test_model_prices_schema.py @@ -367,6 +367,47 @@ def test_active_mistral_chat_rows_price_cache_reads_below_input(path: Path): assert drifted == [] +DEEPSEEK_PRICED_ROWS: Final = tuple( + f"{prefix}{name}" + for name in ("deepseek-flash", "deepseek-v4-flash", "deepseek-v4-flash-vision-exp", "deepseek-v4-pro") + for prefix in ("", "deepseek/") +) +DEEPSEEK_OFF_PEAK_WINDOWS: Final = ( + {"hours_utc": ["00:00-01:00", "04:00-06:00", "10:00-00:00"], "weekdays": [1, 2, 3, 4, 5]}, + {"hours_utc": "00:00-00:00", "weekdays": [6, 7]}, +) +DEEPSEEK_HALVED_RATES: Final = ("input_cost_per_token", "output_cost_per_token", "cache_read_input_token_cost") + + +def deepseek_off_peak_drift(entry: Mapping[str, object]) -> str | None: + block: Final = entry.get("off_peak_pricing") + if not isinstance(block, dict): + return "no off_peak_pricing block" + if tuple(block.get("windows", ())) != DEEPSEEK_OFF_PEAK_WINDOWS: + return f"windows={block.get('windows')}" + halved: Final = {rate: block.get(rate) for rate in DEEPSEEK_HALVED_RATES} + expected: Final = {rate: float(str(entry[rate])) / 2 for rate in DEEPSEEK_HALVED_RATES} + mismatched: Final = { + rate for rate in DEEPSEEK_HALVED_RATES if halved[rate] != pytest.approx(expected[rate], rel=1e-9) + } + return f"off-peak rates {halved} are not half of the listed rates" if mismatched else None + + +@pytest.mark.parametrize("path", (PRICES_PATH, BACKUP_PRICES_PATH), ids=("main", "backup")) +def test_deepseek_rows_bill_half_rate_outside_weekday_peak_hours(path: Path): + """DeepSeek charges half its listed rate outside 01:00-04:00 and 06:00-10:00 UTC Monday to + Friday (api-docs.deepseek.com/quick_start/pricing, read 2026-09-19), so every row on that + pricing page carries an off_peak_pricing block with those windows and the halved rates.""" + rows: Mapping[str, object] = json.loads(path.read_text()) + drifted: Final = { + name: deepseek_off_peak_drift(entry) + for name in DEEPSEEK_PRICED_ROWS + if isinstance(entry := rows.get(name), dict) and deepseek_off_peak_drift(entry) is not None + } + assert drifted == {} + assert all(name in rows for name in DEEPSEEK_PRICED_ROWS) + + PROVIDER_LABELS_WITHOUT_A_MODEL_SET: Final = frozenset({"sagemaker", "bedrock_converse"}) MODES_SERVED_OUTSIDE_THE_LLM_PROVIDER_REGISTRY: Final = frozenset({"search", "evaluation"}) VERTEX_FAMILIES_A_VERTEX_WILDCARD_GRANT_DOES_NOT_LIST: Final = frozenset( diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 8bb9a7e86a3..537c5ed45a4 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -952,6 +952,38 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "supports_image_size": {"type": "boolean"}, "supports_native_structured_output": {"type": "boolean"}, "use_openai_responses_path": {"type": "boolean"}, + "off_peak_pricing": { + "type": "object", + "properties": { + "hours_utc": { + "oneOf": [{"type": "string"}, {"type": "array", "items": {"type": "string"}}], + }, + "windows": { + "type": "array", + "items": { + "type": "object", + "properties": { + "hours_utc": { + "oneOf": [{"type": "string"}, {"type": "array", "items": {"type": "string"}}], + }, + "weekdays": { + "type": "array", + "items": {"oneOf": [{"type": "integer"}, {"type": "string"}]}, + }, + }, + "required": ["hours_utc"], + "additionalProperties": False, + }, + }, + "weekday_timezone": {"type": "string"}, + "input_cost_per_token": {"type": "number"}, + "output_cost_per_token": {"type": "number"}, + "output_cost_per_reasoning_token": {"type": "number"}, + "cache_read_input_token_cost": {"type": "number"}, + "cache_creation_input_token_cost": {"type": "number"}, + }, + "additionalProperties": False, + }, "tiered_pricing": { "type": "array", "items": { From df6a222cb88cb9e44b1b8649d11c17466afc0d3c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 04:59:26 -0700 Subject: [PATCH 15/23] fix(proxy): validate bulk object_permission against the key's team as /key/update does --- .../key_management_endpoints.py | 52 +++++++++++++++---- .../test_key_management_endpoints.py | 40 ++++++++++++-- 2 files changed, 78 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 033ada2c50d..4554a85b225 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2798,9 +2798,18 @@ async def _process_single_key_update( llm_router=llm_router, ) + key_request: Final = await _with_validated_object_permission( + update_key_request=update_key_request, + team_obj=team_obj, + existing_key_row=existing_key_row, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_api_key_dict=user_api_key_dict, + ) + # Prepare update data non_default_values = await prepare_key_update_data( - data=update_key_request, existing_key_row=existing_key_row, prisma_client=prisma_client, llm_router=llm_router + data=key_request, existing_key_row=existing_key_row, prisma_client=prisma_client, llm_router=llm_router ) await _enforce_custom_key_policy( @@ -2809,7 +2818,7 @@ async def _process_single_key_update( operation="update", existing_key_row=existing_key_row, non_default_values=non_default_values, - request=update_key_request, + request=key_request, ), ) @@ -2825,15 +2834,15 @@ async def _process_single_key_update( existing_key_row=existing_key_row, prisma_client=prisma_client, ) - _data: Final = {**update_values, "token": update_key_request.key} + _data: Final = {**update_values, "token": key_request.key} response: Final[Mapping[str, object] | None] = cast( # cast-ok: every update_data branch returns a str-keyed dict "Mapping[str, object] | None", - await prisma_client.update_data(token=update_key_request.key, data=_data), + await prisma_client.update_data(token=key_request.key, data=_data), ) # Delete cache await _delete_cache_key_object( - hashed_token=_hash_token_if_needed(update_key_request.key), + hashed_token=_hash_token_if_needed(key_request.key), user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) @@ -2842,17 +2851,15 @@ async def _process_single_key_update( # authenticating against the access groups it just lost. await sync_key_update_access_group_membership( prisma_client=prisma_client, - key_token=_hash_token_if_needed( - _resolve_token_to_update(data=update_key_request, existing_key_row=existing_key_row) - ), - data=update_key_request, + key_token=_hash_token_if_needed(_resolve_token_to_update(data=key_request, existing_key_row=existing_key_row)), + data=key_request, existing_key_row=existing_key_row, ) # Trigger async hook asyncio.create_task( KeyManagementEventHooks.async_key_updated_hook( - data=update_key_request, + data=key_request, existing_key_row=existing_key_row, response=response, user_api_key_dict=user_api_key_dict, @@ -2875,6 +2882,31 @@ async def _process_single_key_update( return updated_key_info +async def _with_validated_object_permission( + update_key_request: UpdateKeyRequest, + team_obj: LiteLLM_TeamTableCachedObj | None, + existing_key_row: LiteLLM_VerificationToken, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + user_api_key_dict: UserAPIKeyAuth, +) -> UpdateKeyRequest: + if update_key_request.object_permission is None: + return update_key_request + normalized_object_permission: Final = await _validate_mcp_servers_for_key_update( + data=update_key_request, + team_obj=team_obj, + existing_key_row=existing_key_row, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + is_proxy_admin=user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value, + ) + if normalized_object_permission is None: + return update_key_request + return update_key_request.model_copy( + update=MappingProxyType({"object_permission": LiteLLM_ObjectPermissionBase(**normalized_object_permission)}) + ) + + async def _validate_mcp_servers_for_key_update( data: "UpdateKeyRequest", team_obj: Optional["LiteLLM_TeamTableCachedObj"], diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 1bf5018900a..93c5a8c3ded 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -80,7 +80,11 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( validate_key_team_change, ) from litellm.proxy.proxy_server import app -from litellm.types.proxy.management_endpoints.key_management_endpoints import CustomKeyPolicyRequest +from litellm.types.proxy.management_endpoints.key_management_endpoints import ( + BulkUpdateKeyRequest, + BulkUpdateKeyResponse, + CustomKeyPolicyRequest, +) client = TestClient(app) @@ -7098,23 +7102,29 @@ async def test_list_key_helper_applies_search_to_prisma_where(): _BULK_UPDATE_TOKEN: Final = "1f2e3d4c5b6a79880123456789abcdef0123456789abcdef0123456789abcdef" +_BULK_UPDATE_TEAM: Final = LiteLLM_TeamTableCachedObj(team_id="team-1") -async def _bulk_update_one_key(monkeypatch, item_payload: Mapping[str, object]) -> AsyncMock: +async def _run_bulk_update_on_one_key( + monkeypatch, item_payload: Mapping[str, object], team: LiteLLM_TeamTableCachedObj = _BULK_UPDATE_TEAM +) -> tuple[BulkUpdateKeyResponse, AsyncMock]: from litellm.proxy.management_endpoints.key_management_endpoints import bulk_update_keys - from litellm.types.proxy.management_endpoints.key_management_endpoints import BulkUpdateKeyRequest key_in_db = LiteLLM_VerificationToken( token=_BULK_UPDATE_TOKEN, user_id="test-user", team_id="team-1", max_budget=100.0, budget_id="budget-1" ) mock_prisma_client = AsyncMock() mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=key_in_db) + mock_prisma_client.get_data = AsyncMock(return_value=key_in_db) mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=None) mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock( return_value=MagicMock(object_permission_id="objperm-bulk") ) mock_prisma_client.update_data = AsyncMock(return_value={"data": {"token": _BULK_UPDATE_TOKEN}}) _setup_update_key_mocks(monkeypatch, mock_prisma_client) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", AsyncMock(return_value=team) + ) with ( patch( # test-quality-ok: the handler reads the cache and hook singletons from module globals, no injection seam @@ -7138,8 +7148,13 @@ async def _bulk_update_one_key(monkeypatch, item_payload: Mapping[str, object]) litellm_changed_by=None, ) + return response, mock_prisma_client + + +async def _bulk_update_one_key(monkeypatch, item_payload: Mapping[str, object]) -> AsyncMock: + response, prisma = await _run_bulk_update_on_one_key(monkeypatch, item_payload) assert response.failed_updates == [] - return mock_prisma_client + return prisma def _written_key_row(prisma: AsyncMock) -> Mapping[str, object]: @@ -7178,6 +7193,23 @@ async def test_bulk_update_keys_object_permission_is_granted_not_dropped(monkeyp assert not {"max_budget", "team_id", "budget_id"} & written.keys() +@pytest.mark.asyncio +async def test_bulk_update_keys_object_permission_outside_the_team_allowlist_is_refused(monkeypatch): + """A bulk item's object_permission is checked against the key's team exactly as /key/update + checks it, so a team key cannot be granted a search tool its team does not allow.""" + team = LiteLLM_TeamTableCachedObj( + team_id="team-1", + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="op-team-1", search_tools=["team-search"]), + ) + response, prisma = await _run_bulk_update_on_one_key( + monkeypatch, {"object_permission": {"search_tools": ["other-search"]}}, team=team + ) + + assert response.successful_updates == [] + assert "not allowed by team 'team-1'" in response.failed_updates[0].failed_reason + prisma.update_data.assert_not_called() + + @pytest.mark.asyncio async def test_generate_key_negative_max_budget(): """ From dc02e5f5fb2b9f856a7fa80f33ef33588a6529a7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 05:12:12 -0700 Subject: [PATCH 16/23] test(proxy): stub the existing key's team in the bulk item policy tests --- .../management_endpoints/test_key_management_endpoints.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 93c5a8c3ded..e2a68988ee2 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -13189,6 +13189,11 @@ async def _process_single_key_update_under_policy(prisma_client: AsyncMock, data "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook", new_callable=AsyncMock, ), + patch( # test-quality-ok: the existing key's team is outside the policy path, as in the /key/update tests + "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", + new_callable=AsyncMock, + return_value=None, + ), ): return await _process_single_key_update( update_key_request=data, From 0a000217229880d612a63e483fd6eefb71370ee6 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 19 Sep 2026 07:33:02 -0700 Subject: [PATCH 17/23] fix(rust): resolve Mistral OCR credentials in Python's env order Python resolves the Mistral key as api_key, MISTRAL_AZURE_API_KEY, then MISTRAL_API_KEY, and the base as api_base, MISTRAL_AZURE_API_BASE, then the public endpoint, never reading MISTRAL_API_BASE. Native OCR read MISTRAL_API_KEY and MISTRAL_API_BASE instead, so with the Azure pair set it sent the call to a different endpoint with a different key. Empty env values now fall through like Python's `or` chain. Co-Authored-By: Claude Opus 5 --- litellm-rust/crates/core/src/ocr/prepare.rs | 22 ++++++++++-------- litellm-rust/crates/core/tests/ocr.rs | 25 ++++++++++++++++----- 2 files changed, 32 insertions(+), 15 deletions(-) diff --git a/litellm-rust/crates/core/src/ocr/prepare.rs b/litellm-rust/crates/core/src/ocr/prepare.rs index ed8c7fba503..f1e1dcaaa6d 100644 --- a/litellm-rust/crates/core/src/ocr/prepare.rs +++ b/litellm-rust/crates/core/src/ocr/prepare.rs @@ -13,24 +13,28 @@ pub(crate) fn prepare_request( client: &OcrClient, ) -> PreparedOcrRequest { let credentials = request.credentials.clone(); - let api_base_env = match request.config.provider() { - OcrProvider::Mistral => Some("MISTRAL_API_BASE"), - OcrProvider::AzureAi => Some("AZURE_AI_API_BASE"), - OcrProvider::Cohere | OcrProvider::Reducto | OcrProvider::VertexAi => None, + let (preferred_api_key_env, api_base_env) = match request.config.provider() { + OcrProvider::Mistral => ( + Some("MISTRAL_AZURE_API_KEY"), + Some("MISTRAL_AZURE_API_BASE"), + ), + OcrProvider::AzureAi => (None, Some("AZURE_AI_API_BASE")), + OcrProvider::Cohere | OcrProvider::Reducto | OcrProvider::VertexAi => (None, None), }; + let secret = |name: &str| client.secrets().truthy(name); let dynamic_api_key = credentials.dynamic_api_key.or_else(|| { credentials.api_key.clone().or_else(|| { - request - .config - .get_api_key_env_var() - .and_then(|name| client.secrets().get(name)) + preferred_api_key_env + .into_iter() + .chain(request.config.get_api_key_env_var()) + .find_map(secret) .map(|value| Sourced::new(SecretValue::new(value), InputSource::Environment)) }) }); let dynamic_api_base = credentials.dynamic_api_base.or_else(|| { credentials.api_base.clone().or_else(|| { api_base_env - .and_then(|name| client.secrets().get(name)) + .and_then(secret) .map(|value| Sourced::new(value, InputSource::Environment)) }) }); diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index 61d59a38065..3aedc7b9023 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -174,14 +174,30 @@ async fn facade_retains_native_response_when_requested() { ); } +#[rstest] +#[case::plain_key(&[("MISTRAL_API_KEY", "plain")], "plain")] +#[case::azure_key_wins(&[("MISTRAL_AZURE_API_KEY", "azure"), ("MISTRAL_API_KEY", "plain")], "azure")] +#[case::empty_azure_key_falls_through(&[("MISTRAL_AZURE_API_KEY", ""), ("MISTRAL_API_KEY", "plain")], "plain")] #[tokio::test] -async fn provider_key_fallback_reads_the_injected_secret_source() { +async fn mistral_env_fallbacks_follow_python_through_the_injected_secret_source( + #[case] secrets: &'static [(&'static str, &'static str)], + #[case] expected_key: &str, +) { let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let secret_base = base.clone(); + let client = ocr_client().with_secrets(Arc::new(move |name: &str| match name { + "MISTRAL_AZURE_API_BASE" => Some(secret_base.clone()), + "MISTRAL_API_BASE" => Some("http://127.0.0.1:9/never-read".into()), + _ => secrets + .iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| value.to_string()), + })); let request = decode_request(OcrWireRequest { model: "mistral/model".into(), document: json!({"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}), api_key: None, - api_base: Some(base.clone()), + api_base: None, custom_llm_provider: None, extra_headers: None, optional_params: Default::default(), @@ -189,13 +205,10 @@ async fn provider_key_fallback_reads_the_injected_secret_source() { timeout_seconds: Some(2.0), }) .unwrap(); - let client = ocr_client().with_secrets(Arc::new(|name: &str| { - (name == "MISTRAL_API_KEY").then(|| "from-secret-manager".to_string()) - })); crate::ocr::client::perform(&client, request).await.unwrap(); server.await.unwrap(); - assert!(seen.lock().unwrap()[0].contains("authorization: Bearer from-secret-manager")); + assert!(seen.lock().unwrap()[0].contains(&format!("authorization: Bearer {expected_key}"))); } #[tokio::test] From c735cc3db14e357be69c8e9be51455f212320920 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 08:00:12 -0700 Subject: [PATCH 18/23] test(cost): point dated snapshot tests at a date the cost map cannot carry The azure row of test_get_model_info_falls_back_from_dated_snapshot_to_undated_entry used gpt-5.6-luna-2026-07-09, which main's cost map carries as an exact azure key, so the lookup returned the dated key and the required misc test job failed on main. All three dated snapshot tests now use a 2099-01-01 snapshot date, so they keep exercising the strip path whatever real snapshots the map gains --- tests/test_litellm/test_cost_calculator.py | 2 +- tests/test_litellm/test_utils.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 50eec369c07..aef17f3d5d0 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -113,7 +113,7 @@ def test_completion_cost_uses_response_model_for_dynamic_routing(_local_model_co def test_completion_cost_strips_dated_azure_snapshot_model(_local_model_cost_map: None) -> None: dated_response = ModelResponse( - model="gpt-5.6-luna-2026-07-09", + model="gpt-5.6-luna-2099-01-01", choices=[Choices(index=0, message=Message(role="assistant", content="hi"), finish_reason="stop")], usage=Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150), ) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 71af3cf76a6..d2eb40bedca 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -186,8 +186,8 @@ def test_get_model_info_strips_openai_finetune_ids_without_a_custom_suffix(local @pytest.mark.parametrize( ("model", "custom_llm_provider", "expected_key"), [ - ("gpt-5.6-luna-2026-07-09", "openai", "gpt-5.6-luna"), - ("gpt-5.6-luna-2026-07-09", "azure", "azure/gpt-5.6-luna"), + ("gpt-5.6-luna-2099-01-01", "openai", "gpt-5.6-luna"), + ("gpt-5.6-luna-2099-01-01", "azure", "azure/gpt-5.6-luna"), ], ) def test_get_model_info_falls_back_from_dated_snapshot_to_undated_entry( From 0074b943a65087b4c7fde9897703ab5109e35d05 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 19 Sep 2026 08:30:11 -0700 Subject: [PATCH 19/23] fix(rust): read proxy env vars in urllib's order Python resolves proxies through urllib.request.getproxies_environment: the lowercase variable wins, an empty value is unset, an empty lowercase value clears the uppercase one, and under CGI only the uppercase HTTP_PROXY is forgotten because a client can set it with a Proxy header. The Rust route took the uppercase variable even when empty and dropped every proxy under CGI, so provider calls could skip a required egress proxy --- litellm-rust/crates/http/src/proxy.rs | 48 ++++++++++++++++++++------- 1 file changed, 36 insertions(+), 12 deletions(-) diff --git a/litellm-rust/crates/http/src/proxy.rs b/litellm-rust/crates/http/src/proxy.rs index e51ce3141e5..e771631435d 100644 --- a/litellm-rust/crates/http/src/proxy.rs +++ b/litellm-rust/crates/http/src/proxy.rs @@ -11,19 +11,17 @@ pub struct EnvironmentProxies { impl EnvironmentProxies { pub fn from_environment(env: &impl Lookup) -> Self { - if env.get("REQUEST_METHOD").is_some() { - return Self::default(); - } - let first = |upper: &str, lower: &str| { - env.get(upper) - .or_else(|| env.get(lower)) + let lowercase_first = |upper: Option<&str>, lower: &str| { + env.get(lower) + .or_else(|| upper.and_then(|name| env.truthy(name))) .unwrap_or_default() }; + let is_cgi = env.get("REQUEST_METHOD").is_some(); Self { - all: first("ALL_PROXY", "all_proxy"), - http: first("HTTP_PROXY", "http_proxy"), - https: first("HTTPS_PROXY", "https_proxy"), - no: first("NO_PROXY", "no_proxy"), + all: lowercase_first(Some("ALL_PROXY"), "all_proxy"), + http: lowercase_first((!is_cgi).then_some("HTTP_PROXY"), "http_proxy"), + https: lowercase_first(Some("HTTPS_PROXY"), "https_proxy"), + no: lowercase_first(Some("NO_PROXY"), "no_proxy"), } } @@ -78,8 +76,6 @@ mod tests { #[case::all_covers_https(&[("ALL_PROXY", "http://proxy:3128")], "https://api.test/", true)] #[case::lowercase(&[("https_proxy", "http://proxy:3128")], "https://api.test/", true)] #[case::no_proxy_bypass(&[("HTTPS_PROXY", "http://proxy:3128"), ("NO_PROXY", "api.test")], "https://api.test/", false)] - #[case::cgi_ignores_everything(&[("HTTPS_PROXY", "http://proxy:3128"), ("REQUEST_METHOD", "GET")], "https://api.test/", false)] - #[case::uppercase_wins_even_when_empty(&[("HTTPS_PROXY", ""), ("https_proxy", "http://proxy:3128")], "https://api.test/", false)] fn proxies_follow_the_injected_environment( #[case] env: &'static [(&'static str, &'static str)], #[case] target: &str, @@ -89,6 +85,34 @@ mod tests { assert_eq!(proxies.apply_to(&url(target)), expected); } + #[rstest] + #[case::lowercase_wins(&[("HTTPS_PROXY", "http://upper:3128"), ("https_proxy", "http://lower:3128")], &[("https_proxy", "http://lower:3128")])] + #[case::empty_uppercase_falls_through_to_lowercase(&[("HTTPS_PROXY", ""), ("https_proxy", "http://lower:3128")], &[("https_proxy", "http://lower:3128")])] + #[case::empty_uppercase_alone_is_unset(&[("HTTPS_PROXY", "")], &[])] + #[case::empty_lowercase_clears_the_uppercase_value(&[("HTTPS_PROXY", "http://upper:3128"), ("https_proxy", "")], &[])] + #[case::lowercase_no_proxy_wins(&[("NO_PROXY", "upper.test"), ("no_proxy", "lower.test")], &[("no_proxy", "lower.test")])] + #[case::cgi_forgets_the_client_settable_http_proxy(&[("REQUEST_METHOD", "GET"), ("HTTP_PROXY", "http://attacker:3128")], &[])] + #[case::cgi_keeps_lowercase_http_proxy(&[("REQUEST_METHOD", "GET"), ("HTTP_PROXY", "http://attacker:3128"), ("http_proxy", "http://lower:3128")], &[("http_proxy", "http://lower:3128")])] + #[case::cgi_keeps_every_other_variable(&[("REQUEST_METHOD", "GET"), ("HTTPS_PROXY", "http://proxy:3128"), ("ALL_PROXY", "http://all:3128"), ("NO_PROXY", "internal.test")], &[("HTTPS_PROXY", "http://proxy:3128"), ("ALL_PROXY", "http://all:3128"), ("NO_PROXY", "internal.test")])] + fn variables_resolve_like_urllib_getproxies_environment( + #[case] env: &'static [(&'static str, &'static str)], + #[case] equivalent: &'static [(&'static str, &'static str)], + ) { + assert_eq!( + EnvironmentProxies::from_environment(&env_of(env)), + EnvironmentProxies::from_environment(&env_of(equivalent)) + ); + } + + #[test] + fn a_cgi_request_still_proxies_https_through_the_configured_proxy() { + let proxies = EnvironmentProxies::from_environment(&env_of(&[ + ("REQUEST_METHOD", "GET"), + ("HTTPS_PROXY", "http://proxy:3128"), + ])); + assert!(proxies.apply_to(&url("https://api.test/"))); + } + #[test] fn an_empty_environment_proxies_nothing() { let proxies = EnvironmentProxies::from_environment(&env_of(&[])); From b341d21a7657ae002baecd3e82df071436464963 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 19 Sep 2026 08:30:35 -0700 Subject: [PATCH 20/23] fix(rust): redact proxy credentials in Debug and build the proxy matcher once EnvironmentProxies holds raw proxy URLs, which can carry user:password, and it sits inside HttpSettings and HttpClientConfig, so any {:?} of those would print the password. Derive veil's Redact like the auth crate does. NO_PROXY stays readable because it holds no credentials. The media fetcher also rebuilt the hyper-util matcher for every URL and redirect hop. Build it once when the fetcher is created --- litellm-rust/Cargo.lock | 1 + litellm-rust/crates/http/Cargo.toml | 1 + litellm-rust/crates/http/src/media.rs | 3 +-- litellm-rust/crates/http/src/proxy.rs | 32 +++++++++++++++++++++------ 4 files changed, 28 insertions(+), 9 deletions(-) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 0cbca96ad57..726d2f484da 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2145,6 +2145,7 @@ dependencies = [ "serde_json", "thiserror 2.0.19", "tokio", + "veil", "webpki-roots", ] diff --git a/litellm-rust/crates/http/Cargo.toml b/litellm-rust/crates/http/Cargo.toml index 4f94f37a8d5..d4457f5685c 100644 --- a/litellm-rust/crates/http/Cargo.toml +++ b/litellm-rust/crates/http/Cargo.toml @@ -17,6 +17,7 @@ rustls.workspace = true serde_json.workspace = true thiserror.workspace = true tokio.workspace = true +veil.workspace = true webpki-roots.workspace = true [dev-dependencies] diff --git a/litellm-rust/crates/http/src/media.rs b/litellm-rust/crates/http/src/media.rs index ae3f55b476a..3b29c9e28a7 100644 --- a/litellm-rust/crates/http/src/media.rs +++ b/litellm-rust/crates/http/src/media.rs @@ -103,8 +103,7 @@ impl MediaFetcher { config: &HttpClientConfig, url_policy: UrlPolicy, ) -> Result { - let proxies = config.proxies.clone(); - let uses_proxy: ProxyMatch = Arc::new(move |url| proxies.apply_to(url)); + let uses_proxy: ProxyMatch = Arc::new(config.proxies.matcher()); Self::with_resolution( pool, config, diff --git a/litellm-rust/crates/http/src/proxy.rs b/litellm-rust/crates/http/src/proxy.rs index e771631435d..7fedaff4418 100644 --- a/litellm-rust/crates/http/src/proxy.rs +++ b/litellm-rust/crates/http/src/proxy.rs @@ -1,10 +1,14 @@ use hyper_util::client::proxy::matcher::Matcher; use litellm_core_utils::settings::Lookup; +use veil::Redact; -#[derive(Clone, Debug, Default, PartialEq, Eq, Hash)] +#[derive(Clone, Redact, Default, PartialEq, Eq, Hash)] pub struct EnvironmentProxies { + #[redact] all: String, + #[redact] http: String, + #[redact] https: String, no: String, } @@ -25,16 +29,18 @@ impl EnvironmentProxies { } } - pub fn apply_to(&self, url: &reqwest::Url) -> bool { + pub(crate) fn matcher(&self) -> impl Fn(&reqwest::Url) -> bool + Send + Sync + use<> { let matcher = Matcher::builder() .all(self.all.clone()) .http(self.http.clone()) .https(self.https.clone()) .no(self.no.clone()) .build(); - url.as_str() - .parse::() - .is_ok_and(|uri| matcher.intercept(&uri).is_some()) + move |url| { + url.as_str() + .parse::() + .is_ok_and(|uri| matcher.intercept(&uri).is_some()) + } } pub(crate) fn reqwest_proxies(&self) -> Vec { @@ -82,7 +88,7 @@ mod tests { #[case] expected: bool, ) { let proxies = EnvironmentProxies::from_environment(&env_of(env)); - assert_eq!(proxies.apply_to(&url(target)), expected); + assert_eq!(proxies.matcher()(&url(target)), expected); } #[rstest] @@ -110,7 +116,19 @@ mod tests { ("REQUEST_METHOD", "GET"), ("HTTPS_PROXY", "http://proxy:3128"), ])); - assert!(proxies.apply_to(&url("https://api.test/"))); + assert!(proxies.matcher()(&url("https://api.test/"))); + } + + #[test] + fn debug_output_hides_proxy_credentials_but_shows_which_variables_are_set() { + let proxies = EnvironmentProxies::from_environment(&env_of(&[ + ("HTTPS_PROXY", "http://operator:hunter2@proxy.corp:3128"), + ("NO_PROXY", "internal.test"), + ])); + let debug = format!("{proxies:?}"); + assert!(!debug.contains("hunter2") && !debug.contains("operator")); + assert!(debug.contains("internal.test")); + assert_ne!(debug, format!("{:?}", EnvironmentProxies::default())); } #[test] From 1669213eb552fa69ad542c1673c7b7588e5f8ae0 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 19 Sep 2026 08:32:46 -0700 Subject: [PATCH 21/23] fix(rust): read OCR secrets from the process environment and decline when a secret manager is readable The OCR route called back into Python's get_secret_str for every env fallback. With no secret manager configured that is os.environ behind a GIL hop, and with one configured it blocked a tokio worker on vault I/O and also sent the Azure and GCP identity variables, which Python reads with os.getenv, to the vault. The other Rust routes already read the process environment. Read the process environment here too. When litellm would read secrets from a secret manager, decline the Rust route so the Python route serves the call with the vault-backed keys --- .../crates/python-bridge/python_settings.json | 3 + .../python-bridge/src/python_settings.rs | 74 +++---------------- .../python-bridge/src/routes/ocr/mod.rs | 73 ++++++++++++++++-- litellm/rust_bridge/settings.py | 11 ++- .../test_litellm/rust_bridge/test_settings.py | 25 ++++--- 5 files changed, 100 insertions(+), 86 deletions(-) diff --git a/litellm-rust/crates/python-bridge/python_settings.json b/litellm-rust/crates/python-bridge/python_settings.json index 4ad3edf682d..0af55083bef 100644 --- a/litellm-rust/crates/python-bridge/python_settings.json +++ b/litellm-rust/crates/python-bridge/python_settings.json @@ -19,5 +19,8 @@ "vertex_project", "vertex_location", "enable_azure_ad_token_refresh" + ], + "secret_manager": [ + "readable" ] } diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs index 83db4f02500..7ac23a05542 100644 --- a/litellm-rust/crates/python-bridge/src/python_settings.rs +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -1,4 +1,3 @@ -use litellm_core_utils::settings::Lookup; use pyo3::prelude::*; const MODULE: &str = "litellm.rust_bridge.settings"; @@ -8,17 +7,24 @@ pub(crate) enum PythonSettings { Http, UrlPolicy, ProviderDefaults, + SecretManager, } impl PythonSettings { #[cfg(test)] - pub(crate) const ALL: [Self; 3] = [Self::Http, Self::UrlPolicy, Self::ProviderDefaults]; + pub(crate) const ALL: [Self; 4] = [ + Self::Http, + Self::UrlPolicy, + Self::ProviderDefaults, + Self::SecretManager, + ]; pub(crate) fn name(self) -> &'static str { match self { Self::Http => "http_settings", Self::UrlPolicy => "url_policy", Self::ProviderDefaults => "provider_defaults", + Self::SecretManager => "secret_manager", } } @@ -32,23 +38,6 @@ impl PythonSettings { } } -pub(crate) struct PythonSecrets; - -impl Lookup for PythonSecrets { - fn get(&self, name: &str) -> Option { - Python::attach(|py| { - py.import(MODULE) - .and_then(|module| module.getattr("secret")?.call1((name,))) - .and_then(|value| value.extract::>()) - .unwrap_or_else(|error| { - let _ = - PythonSettings::warn(py, &format!("reading secret {name} failed: {error}")); - None - }) - }) - } -} - #[cfg(test)] pub(crate) const CONTRACT: &str = include_str!("../python_settings.json"); @@ -56,10 +45,9 @@ pub(crate) const CONTRACT: &str = include_str!("../python_settings.json"); mod tests { use std::{collections::BTreeSet, ffi::CString}; - use litellm_core_utils::settings::Lookup; use pyo3::{prelude::*, types::PyDict}; - use super::{CONTRACT, PythonSecrets, PythonSettings}; + use super::{CONTRACT, PythonSettings}; #[test] fn every_settings_group_is_in_the_python_contract() { @@ -83,48 +71,4 @@ mod tests { assert_eq!(read, declared); }); } - - #[test] - fn secrets_come_from_the_python_secret_reader_and_a_failed_read_is_unset() { - Python::initialize(); - Python::attach(|py| { - py.run( - c" -import sys -import types -settings = types.ModuleType('litellm.rust_bridge.settings') -settings.warnings = [] -def secret(name): - if name == 'BROKEN': - raise RuntimeError('vault down') - return {'MISTRAL_API_KEY': 'from-vault'}.get(name) -settings.secret = secret -settings.warn = settings.warnings.append -sys.modules.setdefault('litellm', types.ModuleType('litellm')) -sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bridge')) -sys.modules['litellm.rust_bridge.settings'] = settings -", - None, - None, - ) - .unwrap(); - }); - assert_eq!( - PythonSecrets.get("MISTRAL_API_KEY").as_deref(), - Some("from-vault") - ); - assert_eq!(PythonSecrets.get("ABSENT"), None); - assert_eq!(PythonSecrets.get("BROKEN"), None); - Python::attach(|py| { - let warnings: Vec = py - .import("litellm.rust_bridge.settings") - .unwrap() - .getattr("warnings") - .unwrap() - .extract() - .unwrap(); - assert_eq!(warnings.len(), 1); - assert!(warnings[0].contains("BROKEN") && warnings[0].contains("vault down")); - }); - } } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs index 785f6e48e13..9be3171f70b 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs @@ -10,17 +10,16 @@ use litellm_auth_gcp::VertexAuth; use litellm_callbacks_legacy::{LegacySurface, PublicCall, run_legacy_call}; use litellm_core::ocr::route::ocr_machine; use litellm_core_utils::settings::ProcessEnvironment; -use litellm_llms::base_llm::ocr::{handler::OcrClient, settings::OcrSettings}; +use litellm_llms::base_llm::ocr::{ + handler::OcrClient, + settings::{OcrSettings, Secrets}, +}; use pyo3::{ prelude::*, types::{PyDict, PyTuple}, }; -use crate::{ - errors::RustBridgeDeclined, - http, - python_settings::{PythonSecrets, PythonSettings}, -}; +use crate::{errors::RustBridgeDeclined, http, python_settings::PythonSettings}; const SURFACE: LegacySurface = LegacySurface { call_type: "ocr", @@ -42,6 +41,7 @@ fn run_ocr( kwargs: Bound<'_, PyDict>, asynchronous: bool, ) -> PyResult> { + let secrets = process_environment_secrets(&PythonSettings::SecretManager.read(py)?)?; let config = http::call_config(py, &kwargs, asynchronous)?; let client = OcrClient::new( http::pool(), @@ -49,7 +49,7 @@ fn run_ocr( http::url_policy(py)?, VERTEX_AUTH.clone(), ocr_settings(py)?, - Arc::new(PythonSecrets), + secrets, ) .map_err(|error| RustBridgeDeclined::new_err(error.to_string()))?; run_legacy_call( @@ -62,6 +62,21 @@ fn run_ocr( ) } +#[derive(FromPyObject)] +struct PythonSecretManager { + readable: bool, +} + +fn process_environment_secrets(secret_manager: &Bound<'_, PyAny>) -> PyResult { + let manager: PythonSecretManager = secret_manager.extract()?; + if manager.readable { + return Err(RustBridgeDeclined::new_err( + "a readable secret manager is configured and the Rust route only reads the process environment", + )); + } + Ok(Arc::new(ProcessEnvironment)) +} + #[derive(FromPyObject)] struct PythonProviderDefaults { vertex_project: Option, @@ -105,3 +120,47 @@ pub(crate) fn aocr( ) -> PyResult> { run_ocr(py, request, args, kwargs, true) } + +#[cfg(test)] +mod tests { + use pyo3::{prelude::*, types::PyDict}; + + use super::process_environment_secrets; + use crate::errors::RustBridgeDeclined; + + fn secret_manager<'py>(py: Python<'py>, readable: bool) -> Bound<'py, PyAny> { + let locals = PyDict::new(py); + locals.set_item("readable", readable).unwrap(); + py.run( + c"import types\nmanager = types.SimpleNamespace(readable=readable)", + Some(&locals), + Some(&locals), + ) + .unwrap(); + locals.get_item("manager").unwrap().unwrap() + } + + #[test] + fn a_readable_secret_manager_sends_the_call_back_to_python() { + Python::initialize(); + Python::attach(|py| { + let declined = process_environment_secrets(&secret_manager(py, true)) + .err() + .expect("the Rust route declines"); + assert!(declined.is_instance_of::(py)); + }); + } + + #[test] + fn without_a_readable_secret_manager_secrets_are_the_process_environment() { + Python::initialize(); + Python::attach(|py| { + let secrets = process_environment_secrets(&secret_manager(py, false)).unwrap(); + assert_eq!( + secrets.get("LITELLM_RUST_BRIDGE_UNSET_VARIABLE_FOR_TEST"), + None + ); + assert_eq!(secrets.get("PATH"), std::env::var("PATH").ok()); + }); + } +} diff --git a/litellm/rust_bridge/settings.py b/litellm/rust_bridge/settings.py index 037d6d9bd27..86450ffbb6b 100644 --- a/litellm/rust_bridge/settings.py +++ b/litellm/rust_bridge/settings.py @@ -31,16 +31,21 @@ class ProviderDefaults: enable_azure_ad_token_refresh: bool | None +@dataclass(frozen=True, slots=True) +class SecretManager: + readable: bool + + def warn(message: str) -> None: from litellm._logging import verbose_logger verbose_logger.warning("%s", message) -def secret(name: str) -> str | None: - from litellm.secret_managers.main import get_secret_str +def secret_manager() -> SecretManager: + from litellm.secret_managers.main import _should_read_secret_from_secret_manager - return get_secret_str(name) + return SecretManager(readable=_should_read_secret_from_secret_manager()) def provider_defaults() -> ProviderDefaults: diff --git a/tests/test_litellm/rust_bridge/test_settings.py b/tests/test_litellm/rust_bridge/test_settings.py index 44c5ec42b36..6b78ddad44b 100644 --- a/tests/test_litellm/rust_bridge/test_settings.py +++ b/tests/test_litellm/rust_bridge/test_settings.py @@ -11,6 +11,7 @@ import litellm from litellm.integrations.custom_secret_manager import CustomSecretManager from litellm.llms.custom_httpx.http_handler import default_user_agent from litellm.rust_bridge import settings +from litellm.secret_managers.main import get_secret_str from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem CONTRACT_PATH: Final = Path(__file__).parents[3] / "litellm-rust/crates/python-bridge/python_settings.json" @@ -23,6 +24,7 @@ def test_the_rust_contract_matches_the_returned_fields() -> None: "http_settings": [field.name for field in dataclasses.fields(settings.http_settings())], "url_policy": [field.name for field in dataclasses.fields(settings.url_policy())], "provider_defaults": [field.name for field in dataclasses.fields(settings.provider_defaults())], + "secret_manager": [field.name for field in dataclasses.fields(settings.secret_manager())], } @@ -101,25 +103,26 @@ class _VaultSecrets(CustomSecretManager): return self.secrets.get(secret_name) -def test_secret_prefers_the_secret_manager_and_falls_back_to_the_environment_on_a_miss( - monkeypatch: pytest.MonkeyPatch, +@pytest.mark.parametrize( + ("access_mode", "readable"), + [("read_only", True), ("read_and_write", True), ("write_only", False)], +) +def test_secret_manager_is_readable_only_when_litellm_would_read_secrets_from_it( + monkeypatch: pytest.MonkeyPatch, access_mode: str, readable: bool ) -> None: - monkeypatch.setenv("MISTRAL_API_KEY", "stale-env-key") - monkeypatch.setenv("REDUCTO_API_KEY", "env-only-key") + monkeypatch.setenv("MISTRAL_API_KEY", "env-key") monkeypatch.setattr(litellm, "secret_manager_client", _VaultSecrets({"MISTRAL_API_KEY": "vault-key"})) monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.CUSTOM) - monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode="read_only")) + monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode=access_mode)) - assert settings.secret("MISTRAL_API_KEY") == "vault-key" - assert settings.secret("REDUCTO_API_KEY") == "env-only-key" - assert settings.secret("ABSENT_KEY") is None + assert settings.secret_manager() == settings.SecretManager(readable=readable) + assert (get_secret_str("MISTRAL_API_KEY") == "vault-key") is readable -def test_secret_reads_the_environment_without_a_secret_manager(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("MISTRAL_API_KEY", "env-key") +def test_secret_manager_is_not_readable_without_a_client(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "secret_manager_client", None) - assert settings.secret("MISTRAL_API_KEY") == "env-key" + assert settings.secret_manager() == settings.SecretManager(readable=False) def test_provider_defaults_read_the_litellm_globals(monkeypatch: pytest.MonkeyPatch) -> None: From e2397e7dd3df6dae0331df14f46be9c09000a506 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 19 Sep 2026 08:37:27 -0700 Subject: [PATCH 22/23] fix(rust): drop http_proxy under CGI where environment names ignore case --- litellm-rust/crates/http/src/proxy.rs | 26 +++++++++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/litellm-rust/crates/http/src/proxy.rs b/litellm-rust/crates/http/src/proxy.rs index 7fedaff4418..eb960d8200d 100644 --- a/litellm-rust/crates/http/src/proxy.rs +++ b/litellm-rust/crates/http/src/proxy.rs @@ -15,6 +15,10 @@ pub struct EnvironmentProxies { impl EnvironmentProxies { pub fn from_environment(env: &impl Lookup) -> Self { + Self::resolve(env, cfg!(windows)) + } + + fn resolve(env: &impl Lookup, names_ignore_case: bool) -> Self { let lowercase_first = |upper: Option<&str>, lower: &str| { env.get(lower) .or_else(|| upper.and_then(|name| env.truthy(name))) @@ -23,7 +27,11 @@ impl EnvironmentProxies { let is_cgi = env.get("REQUEST_METHOD").is_some(); Self { all: lowercase_first(Some("ALL_PROXY"), "all_proxy"), - http: lowercase_first((!is_cgi).then_some("HTTP_PROXY"), "http_proxy"), + http: if is_cgi && names_ignore_case { + String::new() + } else { + lowercase_first((!is_cgi).then_some("HTTP_PROXY"), "http_proxy") + }, https: lowercase_first(Some("HTTPS_PROXY"), "https_proxy"), no: lowercase_first(Some("NO_PROXY"), "no_proxy"), } @@ -110,6 +118,22 @@ mod tests { ); } + #[test] + fn cgi_drops_http_proxy_entirely_where_variable_names_ignore_case() { + let windows_env = |name: &str| match name.to_ascii_uppercase().as_str() { + "REQUEST_METHOD" => Some("GET".to_string()), + "HTTP_PROXY" => Some("http://attacker:3128".to_string()), + "HTTPS_PROXY" => Some("http://proxy:3128".to_string()), + _ => None, + }; + let proxies = EnvironmentProxies::resolve(&windows_env, true); + assert!(!proxies.matcher()(&url("http://api.test/"))); + assert!(proxies.matcher()(&url("https://api.test/"))); + assert!(EnvironmentProxies::resolve(&windows_env, false).matcher()( + &url("http://api.test/") + )); + } + #[test] fn a_cgi_request_still_proxies_https_through_the_configured_proxy() { let proxies = EnvironmentProxies::from_environment(&env_of(&[ From 0d0c63dde126dbafcdbf1125335b07df0db911b2 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 19 Sep 2026 15:48:07 +0000 Subject: [PATCH 23/23] fix(rust): suppress private settings resolver lint Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/rust_bridge/settings.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/litellm/rust_bridge/settings.py b/litellm/rust_bridge/settings.py index 86450ffbb6b..3aa2d742862 100644 --- a/litellm/rust_bridge/settings.py +++ b/litellm/rust_bridge/settings.py @@ -43,7 +43,9 @@ def warn(message: str) -> None: def secret_manager() -> SecretManager: - from litellm.secret_managers.main import _should_read_secret_from_secret_manager + from litellm.secret_managers.main import ( + _should_read_secret_from_secret_manager, # pyright: ignore[reportPrivateUsage] # canonical resolver is private + ) return SecretManager(readable=_should_read_secret_from_secret_manager())