From caf37c8b6f21bfc84752baaafc568deeaf3b1996 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 18 Sep 2026 20:36:07 -0700 Subject: [PATCH 01/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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 0a000217229880d612a63e483fd6eefb71370ee6 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 19 Sep 2026 07:33:02 -0700 Subject: [PATCH 07/12] 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 0074b943a65087b4c7fde9897703ab5109e35d05 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 19 Sep 2026 08:30:11 -0700 Subject: [PATCH 08/12] 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 09/12] 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 10/12] 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 11/12] 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 12/12] 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())