diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index 1a769ec9708..ab04fb8d4ae 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -15,8 +15,6 @@ futures-util.workspace = true base64.workspace = true litellm-auth.workspace = true litellm-auth-aws.workspace = true -litellm-auth-gcp.workspace = true -litellm-http.workspace = true litellm-llms.workspace = true moka.workspace = true mime_guess = "2.0.5" @@ -37,6 +35,8 @@ url.workspace = true 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/ocr/client.rs b/litellm-rust/crates/core/src/ocr/client.rs index f0b24623a88..c7b4751bd9e 100644 --- a/litellm-rust/crates/core/src/ocr/client.rs +++ b/litellm-rust/crates/core/src/ocr/client.rs @@ -1,5 +1,3 @@ -use litellm_auth_gcp::VertexAuth; -use litellm_http::{HttpClientConfig, HttpClientPool}; use litellm_llms::{ base_llm::ocr::{error::Error, transformation::LiteLLMOcrResponse}, custom_httpx::llm_http_handler::OcrClient, @@ -16,12 +14,3 @@ pub async fn perform( ) -> Result { litellm_host::run::run(ocr_machine(client.clone()), &LocalOcrHost::new(request)).await } - -pub async fn ocr( - pool: &HttpClientPool, - config: &HttpClientConfig, - vertex_auth: VertexAuth, - request: LiteLLMOcrRequest, -) -> Result { - perform(&OcrClient::new(pool, config, vertex_auth)?, request).await -} diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index 2e414c58541..2ae162d964f 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -6,13 +6,13 @@ use litellm_host::{ host::{Host, HostOp, HostResult}, machine::{HostFailure, Machine, MachineStep}, }; -use litellm_http::{HttpClientConfig, HttpClientPool, HttpSettings, Verify}; +use litellm_http::{HttpClientConfig, HttpClientPool, HttpSettings}; use litellm_llms::{ base_llm::ocr::{ error::Error as OcrError, transformation::{LiteLLMOcrResponse, OCR_RESPONSE_MAX_BYTES, OcrTransportConfig}, }, - custom_httpx::llm_http_handler::OcrClient, + custom_httpx::{llm_http_handler::OcrClient, media::PublicDnsResolver}, }; use rstest::rstest; use serde_json::{Value, json}; @@ -173,48 +173,25 @@ async fn facade_retains_native_response_when_requested() { } #[tokio::test] -async fn facade_uses_the_injected_http_pool_configuration() { +async fn ocr_client_uses_the_injected_http_pool_configuration() { let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; let settings = HttpSettings { user_agent: Some("host-owned/1".into()), ..HttpSettings::default() }; - let config = HttpClientConfig::resolve(&settings, None).unwrap(); - crate::ocr::client::ocr( - &HttpClientPool::new(), - &config, + let client = OcrClient::new( + &HttpClientPool::new(Arc::new(PublicDnsResolver)), + &HttpClientConfig::resolve(&settings).unwrap(), VertexAuth::default(), - wire_request("mistral/model", &base, json!({})), ) - .await .unwrap(); + crate::ocr::client::perform(&client, wire_request("mistral/model", &base, json!({}))) + .await + .unwrap(); server.await.unwrap(); assert!(seen.lock().unwrap()[0].contains("user-agent: host-owned/1")); } -#[tokio::test] -async fn unbuildable_http_configuration_fails_before_dispatch() { - let (base, _seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let config = HttpClientConfig { - verify: Verify::CaBundle(std::env::temp_dir().join("litellm-ocr-missing-bundle.pem")), - ..HttpClientConfig::resolve(&HttpSettings::default(), None).unwrap() - }; - let error = crate::ocr::client::ocr( - &HttpClientPool::new(), - &config, - VertexAuth::default(), - wire_request("mistral/model", &base, json!({})), - ) - .await - .unwrap_err(); - server.abort(); - assert!(matches!( - error, - OcrError::Transport(litellm_llms::custom_httpx::transport::Error::Connect(_)) - )); - assert!(error.to_string().contains("litellm-ocr-missing-bundle.pem")); -} - fn event_name(event: &CallEvent) -> &'static str { match event { CallEvent::Started { .. } => "started", diff --git a/litellm-rust/crates/http/src/config.rs b/litellm-rust/crates/http/src/config.rs index 86f1a9b43b6..f27092c1fb5 100644 --- a/litellm-rust/crates/http/src/config.rs +++ b/litellm-rust/crates/http/src/config.rs @@ -4,28 +4,10 @@ use std::{ time::Duration, }; -use crate::settings::{HttpSettings, SslVerify}; - -#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)] -pub enum Error { - #[error("{setting} cannot be expressed with rustls: {reason}")] - Unsupported { - setting: &'static str, - reason: String, - }, - #[error("could not read {}: {message}", path.display())] - Read { path: PathBuf, message: String }, - #[error("{} is not a PEM file: {message}", path.display())] - InvalidPem { path: PathBuf, message: String }, - #[error("could not build the HTTP client: {0}")] - Client(String), -} - -impl From for Error { - fn from(error: reqwest::Error) -> Self { - Self::Client(error.without_url().to_string()) - } -} +use crate::{ + error::Error, + settings::{HttpSettings, SslVerify}, +}; #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub enum Verify { @@ -45,17 +27,13 @@ pub struct HttpClientConfig { pub user_agent: Option, pub trust_proxy_env: bool, pub connect_timeout: Duration, - pub request_timeout: Option, } impl HttpClientConfig { - /// Port of `get_ssl_verify` + `get_ssl_configuration`: the per-call value wins, then the - /// configured (environment-overlaid) `ssl_verify`, then `SSL_CERT_FILE`, then the built-in - /// roots. Settings rustls has no equivalent for are an error instead of a silent no-op. - pub fn resolve( - settings: &HttpSettings, - per_call_ssl_verify: Option<&SslVerify>, - ) -> Result { + /// Port of `get_ssl_verify` + `get_ssl_configuration`: the configured (environment-overlaid) + /// `ssl_verify`, then `SSL_CERT_FILE`, then the built-in roots. Settings rustls has no + /// equivalent for are an error instead of a silent no-op. + pub fn resolve(settings: &HttpSettings) -> Result { if let Some(level) = &settings.ssl_security_level { return Err(Error::Unsupported { setting: "ssl_security_level", @@ -68,7 +46,7 @@ impl HttpClientConfig { reason: format!("key exchange group {curve:?} is fixed by the rustls provider"), }); } - let verify = match per_call_ssl_verify.or(settings.ssl_verify.as_ref()) { + let verify = match &settings.ssl_verify { Some(SslVerify::Disabled) => Verify::Disabled, Some(SslVerify::CaBundle(path)) => Verify::CaBundle(path.clone()), Some(SslVerify::Enabled) | None => settings @@ -84,7 +62,6 @@ impl HttpClientConfig { user_agent: settings.user_agent.clone(), trust_proxy_env: settings.trust_proxy_env, connect_timeout: settings.connect_timeout, - request_timeout: settings.request_timeout, }) } @@ -141,14 +118,10 @@ impl HttpClientConfig { Some(agent) => with_protocol.user_agent(agent), None => with_protocol, }; - let with_proxy = if self.trust_proxy_env { + Ok(if self.trust_proxy_env { with_agent } else { with_agent.no_proxy() - }; - Ok(match self.request_timeout { - Some(timeout) => with_proxy.timeout(timeout), - None => with_proxy, }) } } @@ -180,49 +153,25 @@ mod tests { } #[rstest] - #[case::default(settings(None, None), None, Verify::BuiltInRoots)] + #[case::default(settings(None, None), Verify::BuiltInRoots)] #[case::setting_disables( settings(Some(SslVerify::Disabled), Some("/env/roots.pem")), - None, Verify::Disabled )] #[case::setting_bundle( settings(Some(SslVerify::CaBundle("/configured.pem".into())), Some("/env/roots.pem")), - None, Verify::CaBundle("/configured.pem".into()) )] #[case::enabled_uses_cert_file( settings(Some(SslVerify::Enabled), Some("/env/roots.pem")), - None, Verify::CaBundle("/env/roots.pem".into()) )] - #[case::unset_uses_cert_file(settings(None, Some("/env/roots.pem")), None, Verify::CaBundle("/env/roots.pem".into()))] - #[case::per_call_beats_setting( - settings(Some(SslVerify::Disabled), None), - Some(SslVerify::Enabled), - Verify::BuiltInRoots - )] - #[case::per_call_disables( - settings(Some(SslVerify::CaBundle("/configured.pem".into())), Some("/env/roots.pem")), - Some(SslVerify::Disabled), - Verify::Disabled - )] - #[case::per_call_bundle( - settings(None, Some("/env/roots.pem")), - Some(SslVerify::CaBundle("/call.pem".into())), - Verify::CaBundle("/call.pem".into()) - )] - #[case::per_call_enabled_still_honours_cert_file( - settings(Some(SslVerify::Disabled), Some("/env/roots.pem")), - Some(SslVerify::Enabled), - Verify::CaBundle("/env/roots.pem".into()) - )] - fn verify_follows_per_call_then_setting_then_cert_file( + #[case::unset_uses_cert_file(settings(None, Some("/env/roots.pem")), Verify::CaBundle("/env/roots.pem".into()))] + fn verify_follows_setting_then_cert_file( #[case] settings: HttpSettings, - #[case] per_call: Option, #[case] expected: Verify, ) { - let config = HttpClientConfig::resolve(&settings, per_call.as_ref()).unwrap(); + let config = HttpClientConfig::resolve(&settings).unwrap(); assert_eq!(config.verify, expected); } @@ -233,7 +182,7 @@ mod tests { ..HttpSettings::default() } .with_environment(&|name: &str| (name == "SSL_VERIFY").then(|| "true".to_string())); - let config = HttpClientConfig::resolve(&settings, None).unwrap(); + let config = HttpClientConfig::resolve(&settings).unwrap(); assert_eq!(config.verify, Verify::BuiltInRoots); } @@ -244,7 +193,7 @@ mod tests { ..HttpSettings::default() }; assert!(matches!( - HttpClientConfig::resolve(&settings, None), + HttpClientConfig::resolve(&settings), Err(Error::Unsupported { setting: "ssl_security_level", .. @@ -259,7 +208,7 @@ mod tests { ..HttpSettings::default() }; assert!(matches!( - HttpClientConfig::resolve(&settings, None), + HttpClientConfig::resolve(&settings), Err(Error::Unsupported { setting: "ssl_ecdh_curve", .. @@ -276,10 +225,9 @@ mod tests { user_agent: Some("litellm/1.0".into()), trust_proxy_env: true, connect_timeout: Duration::from_secs(7), - request_timeout: Some(Duration::from_secs(70)), ..HttpSettings::default() }; - let config = HttpClientConfig::resolve(&settings, None).unwrap(); + let config = HttpClientConfig::resolve(&settings).unwrap(); assert_eq!( config, HttpClientConfig { @@ -290,7 +238,6 @@ mod tests { user_agent: Some("litellm/1.0".into()), trust_proxy_env: true, connect_timeout: Duration::from_secs(7), - request_timeout: Some(Duration::from_secs(70)), } ); } @@ -300,7 +247,7 @@ mod tests { let path = std::env::temp_dir().join("litellm-http-missing-bundle.pem"); let config = HttpClientConfig { verify: Verify::CaBundle(path.clone()), - ..HttpClientConfig::resolve(&HttpSettings::default(), None).unwrap() + ..HttpClientConfig::resolve(&HttpSettings::default()).unwrap() }; assert!(matches!( config.client_builder(), @@ -315,7 +262,7 @@ mod tests { std::fs::write(&path, b"not a certificate").unwrap(); let config = HttpClientConfig { verify: Verify::CaBundle(path.clone()), - ..HttpClientConfig::resolve(&HttpSettings::default(), None).unwrap() + ..HttpClientConfig::resolve(&HttpSettings::default()).unwrap() }; let result = config.client_builder().map(drop); std::fs::remove_file(&path).unwrap(); diff --git a/litellm-rust/crates/http/src/error.rs b/litellm-rust/crates/http/src/error.rs new file mode 100644 index 00000000000..27899f06cf1 --- /dev/null +++ b/litellm-rust/crates/http/src/error.rs @@ -0,0 +1,22 @@ +use std::path::PathBuf; + +#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)] +pub enum Error { + #[error("{setting} cannot be expressed with rustls: {reason}")] + Unsupported { + setting: &'static str, + reason: String, + }, + #[error("could not read {}: {message}", path.display())] + Read { path: PathBuf, message: String }, + #[error("{} is not a PEM file: {message}", path.display())] + InvalidPem { path: PathBuf, message: String }, + #[error("could not build the HTTP client: {0}")] + Client(String), +} + +impl From for Error { + fn from(error: reqwest::Error) -> Self { + Self::Client(error.without_url().to_string()) + } +} diff --git a/litellm-rust/crates/http/src/lib.rs b/litellm-rust/crates/http/src/lib.rs index d62dd768fe1..9c88e3101a7 100644 --- a/litellm-rust/crates/http/src/lib.rs +++ b/litellm-rust/crates/http/src/lib.rs @@ -3,9 +3,11 @@ //! caches `reqwest::Client`s per resolved configuration. mod config; +mod error; mod pool; mod settings; -pub use config::{Error, HttpClientConfig, Verify}; +pub use config::{HttpClientConfig, Verify}; +pub use error::Error; pub use pool::{ClientVariant, HttpClientPool}; pub use settings::{HttpSettings, SslVerify}; diff --git a/litellm-rust/crates/http/src/pool.rs b/litellm-rust/crates/http/src/pool.rs index 065097556e1..03c01e968ac 100644 --- a/litellm-rust/crates/http/src/pool.rs +++ b/litellm-rust/crates/http/src/pool.rs @@ -1,112 +1,174 @@ use std::{ collections::HashMap, - sync::{Mutex, PoisonError}, + sync::{Arc, Mutex, PoisonError}, }; -use crate::config::{Error, HttpClientConfig}; +use reqwest::dns::Resolve; + +use crate::{config::HttpClientConfig, error::Error}; /// The client shapes routes need; each is the shared base plus one policy. #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] pub enum ClientVariant { Provider, NoRedirect, - /// Media downloads: no redirects (the fetcher validates each hop) and never a proxy. + /// Media downloads: no redirects (the fetcher validates each hop), never a proxy, and the + /// pool's media resolver. Media, } -impl ClientVariant { - fn apply(self, builder: reqwest::ClientBuilder) -> reqwest::ClientBuilder { - match self { - Self::Provider => builder, - Self::NoRedirect => builder.redirect(reqwest::redirect::Policy::none()), - Self::Media => builder - .redirect(reqwest::redirect::Policy::none()) - .no_proxy(), - } - } -} - /// Counterpart of `get_async_httpx_client`: one `reqwest::Client` per resolved configuration /// and variant, built on first use and shared afterwards. -#[derive(Default)] pub struct HttpClientPool { + media_resolver: Arc, clients: Mutex>, } impl HttpClientPool { - pub fn new() -> Self { - Self::default() + pub fn new(media_resolver: Arc) -> Self { + Self { + media_resolver, + clients: Mutex::default(), + } } pub fn client( &self, config: &HttpClientConfig, variant: ClientVariant, - ) -> Result { - self.client_with(config, variant, |builder| builder) - } - - /// Like [`Self::client`], with a caller hook for builder options that are not plain values - /// (a DNS resolver, for example). The hook only runs when the client is first built. - pub fn client_with( - &self, - config: &HttpClientConfig, - variant: ClientVariant, - customize: impl FnOnce(reqwest::ClientBuilder) -> reqwest::ClientBuilder, ) -> Result { let key = (config.clone(), variant); let mut clients = self.clients.lock().unwrap_or_else(PoisonError::into_inner); if let Some(client) = clients.get(&key) { return Ok(client.clone()); } - let client = customize(variant.apply(config.client_builder()?)).build()?; + let client = self.apply(variant, config.client_builder()?).build()?; clients.insert(key, client.clone()); Ok(client) } + + fn apply( + &self, + variant: ClientVariant, + builder: reqwest::ClientBuilder, + ) -> reqwest::ClientBuilder { + match variant { + ClientVariant::Provider => builder, + ClientVariant::NoRedirect => builder.redirect(reqwest::redirect::Policy::none()), + ClientVariant::Media => builder + .redirect(reqwest::redirect::Policy::none()) + .no_proxy() + .dns_resolver2(Arc::clone(&self.media_resolver)), + } + } } #[cfg(test)] mod tests { - use std::{cell::Cell, time::Duration}; + use std::{ + net::SocketAddr, + sync::atomic::{AtomicUsize, Ordering}, + time::Duration, + }; - use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use reqwest::dns::{Addrs, Name, Resolving}; + use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::TcpListener, + }; use super::*; use crate::{HttpSettings, Verify}; - fn config(user_agent: &str) -> HttpClientConfig { - HttpClientConfig { - user_agent: Some(user_agent.into()), - ..HttpClientConfig::resolve(&HttpSettings::default(), None).unwrap() + struct FixedResolver(SocketAddr); + + impl Resolve for FixedResolver { + fn resolve(&self, _: Name) -> Resolving { + let addrs: Addrs = Box::new(std::iter::once(self.0)); + Box::pin(std::future::ready(Ok(addrs))) } } - #[test] - fn clients_are_built_once_per_config_and_variant() { - let pool = HttpClientPool::new(); - let builds = Cell::new(0); - let build = |config: &HttpClientConfig, variant| { - pool.client_with(config, variant, |builder| { - builds.set(builds.get() + 1); - builder - }) + fn pool() -> HttpClientPool { + HttpClientPool::new(Arc::new(FixedResolver(([192, 0, 2, 1], 80).into()))) + } + + fn config(user_agent: &str) -> HttpClientConfig { + HttpClientConfig { + user_agent: Some(user_agent.into()), + ..HttpClientConfig::resolve(&HttpSettings::default()).unwrap() + } + } + + /// Answers every request on every connection with `status_line` and counts connections, + /// so a reused client shows up as a reused keep-alive connection. + async fn serve( + status_line: &'static str, + ) -> (SocketAddr, Arc, Arc>>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let connections = Arc::new(AtomicUsize::new(0)); + let requests = Arc::new(Mutex::new(Vec::new())); + let (accepted, seen) = (Arc::clone(&connections), Arc::clone(&requests)); + tokio::spawn(async move { + loop { + let (mut socket, _) = listener.accept().await.unwrap(); + accepted.fetch_add(1, Ordering::SeqCst); + let seen = Arc::clone(&seen); + tokio::spawn(async move { + let mut buffer = vec![0u8; 4096]; + while let Ok(read) = socket.read(&mut buffer).await { + if read == 0 { + return; + } + seen.lock() + .unwrap() + .push(String::from_utf8_lossy(&buffer[..read]).into_owned()); + let response = format!( + "{status_line}\r\nLocation: /elsewhere\r\nContent-Length: 0\r\n\r\n" + ); + if socket.write_all(response.as_bytes()).await.is_err() { + return; + } + } + }); + } + }); + (address, connections, requests) + } + + async fn get( + pool: &HttpClientPool, + config: &HttpClientConfig, + variant: ClientVariant, + url: &str, + ) -> reqwest::Response { + pool.client(config, variant) .unwrap() - }; - build(&config("a"), ClientVariant::Provider); - build(&config("a"), ClientVariant::Provider); - assert_eq!(builds.get(), 1); - build(&config("a"), ClientVariant::NoRedirect); - assert_eq!(builds.get(), 2); - build(&config("b"), ClientVariant::Provider); - assert_eq!(builds.get(), 3); - build(&config("b"), ClientVariant::Provider); - build(&config("a"), ClientVariant::NoRedirect); - assert_eq!(builds.get(), 3); + .get(url) + .timeout(Duration::from_secs(5)) + .send() + .await + .unwrap() + } + + #[tokio::test] + async fn clients_are_shared_per_config_and_variant() { + let (address, connections, _) = serve("HTTP/1.1 204 No Content").await; + let url = format!("http://{address}"); + let pool = pool(); + get(&pool, &config("a"), ClientVariant::Provider, &url).await; + get(&pool, &config("a"), ClientVariant::Provider, &url).await; + assert_eq!(connections.load(Ordering::SeqCst), 1); + get(&pool, &config("a"), ClientVariant::NoRedirect, &url).await; + assert_eq!(connections.load(Ordering::SeqCst), 2); + get(&pool, &config("b"), ClientVariant::Provider, &url).await; + assert_eq!(connections.load(Ordering::SeqCst), 3); } #[test] fn build_failures_are_not_cached() { - let pool = HttpClientPool::new(); + let pool = pool(); let missing = HttpClientConfig { verify: Verify::CaBundle(std::env::temp_dir().join("litellm-http-absent.pem")), ..config("a") @@ -116,57 +178,52 @@ mod tests { assert!(pool.client(&config("a"), ClientVariant::Provider).is_ok()); } - async fn serve_once(status_line: &'static str) -> (String, tokio::task::JoinHandle) { - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base = format!("http://{}", listener.local_addr().unwrap()); - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.unwrap(); - let mut request = vec![0u8; 4096]; - let read = socket.read(&mut request).await.unwrap(); - socket - .write_all( - format!("{status_line}\r\nLocation: /elsewhere\r\nContent-Length: 0\r\nConnection: close\r\n\r\n") - .as_bytes(), - ) - .await - .unwrap(); - String::from_utf8_lossy(&request[..read]).into_owned() - }); - (base, server) - } - #[tokio::test] async fn provider_client_sends_the_configured_user_agent_over_http1() { - let (base, server) = serve_once("HTTP/1.1 204 No Content").await; - let config = HttpClientConfig { - connect_timeout: Duration::from_secs(2), - ..config("litellm-test/9") - }; - let response = HttpClientPool::new() - .client(&config, ClientVariant::Provider) - .unwrap() - .get(&base) - .send() - .await - .unwrap(); + let (address, _, requests) = serve("HTTP/1.1 204 No Content").await; + let response = get( + &pool(), + &config("litellm-test/9"), + ClientVariant::Provider, + &format!("http://{address}"), + ) + .await; assert_eq!(response.status(), 204); assert_eq!(response.version(), reqwest::Version::HTTP_11); - let request = server.await.unwrap(); + let request = requests.lock().unwrap()[0].clone(); assert!(request.contains("user-agent: litellm-test/9"), "{request}"); } #[tokio::test] async fn no_redirect_variant_returns_the_redirect_instead_of_following_it() { - let (base, server) = serve_once("HTTP/1.1 302 Found").await; - let response = HttpClientPool::new() - .client(&config("a"), ClientVariant::NoRedirect) - .unwrap() - .get(&base) - .send() - .await - .unwrap(); + let (address, _, _) = serve("HTTP/1.1 302 Found").await; + let response = get( + &pool(), + &config("a"), + ClientVariant::NoRedirect, + &format!("http://{address}"), + ) + .await; assert_eq!(response.status(), 302); assert_eq!(response.headers()["location"], "/elsewhere"); - server.await.unwrap(); + } + + #[tokio::test] + async fn media_variant_resolves_through_the_injected_resolver() { + let (address, _, requests) = 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()); + let response = get(&pool, &config("a"), ClientVariant::Media, &url).await; + assert_eq!(response.status(), 204); + assert!(requests.lock().unwrap()[0].contains("host: media.invalid")); + assert!( + pool.client(&config("a"), ClientVariant::Provider) + .unwrap() + .get(&url) + .timeout(Duration::from_secs(5)) + .send() + .await + .is_err() + ); } } diff --git a/litellm-rust/crates/http/src/settings.rs b/litellm-rust/crates/http/src/settings.rs index 4d936e89046..45aab0d6fa5 100644 --- a/litellm-rust/crates/http/src/settings.rs +++ b/litellm-rust/crates/http/src/settings.rs @@ -31,7 +31,6 @@ pub struct HttpSettings { pub user_agent: Option, pub trust_proxy_env: bool, pub connect_timeout: Duration, - pub request_timeout: Option, } impl Default for HttpSettings { @@ -47,7 +46,6 @@ impl Default for HttpSettings { user_agent: None, trust_proxy_env: false, connect_timeout: Duration::from_secs(5), - request_timeout: None, } } } diff --git a/litellm-rust/crates/llms/src/custom_httpx/media.rs b/litellm-rust/crates/llms/src/custom_httpx/media.rs index aeac4894683..a1c4fe68734 100644 --- a/litellm-rust/crates/llms/src/custom_httpx/media.rs +++ b/litellm-rust/crates/llms/src/custom_httpx/media.rs @@ -66,26 +66,15 @@ impl MediaFetcher { pool: &HttpClientPool, config: &HttpClientConfig, ) -> Result { - Self::with_resolvers( - pool, - config, - Arc::new(PublicDnsResolver), - Arc::new(SystemAddressResolver), - ) + Self::with_address_resolver(pool, config, Arc::new(SystemAddressResolver)) } - fn with_resolvers( + fn with_address_resolver( pool: &HttpClientPool, config: &HttpClientConfig, - transport_resolver: Arc, address_resolver: Arc, - ) -> Result - where - R: Resolve + 'static, - { - let client = pool.client_with(config, ClientVariant::Media, |builder| { - builder.dns_resolver(transport_resolver) - })?; + ) -> Result { + let client = pool.client(config, ClientVariant::Media)?; Ok(Self { client, address_resolver, @@ -242,7 +231,7 @@ fn is_blocked_ip(ip: IpAddr) -> bool { } #[derive(Default)] -struct PublicDnsResolver; +pub struct PublicDnsResolver; struct SystemAddressResolver; @@ -371,10 +360,9 @@ mod tests { address: SocketAddr, blocked_hosts: HashSet<&'static str>, ) -> MediaFetcher { - MediaFetcher::with_resolvers( - &HttpClientPool::new(), - &HttpClientConfig::resolve(&HttpSettings::default(), None).unwrap(), - Arc::new(LoopbackDnsResolver(address)), + MediaFetcher::with_address_resolver( + &HttpClientPool::new(Arc::new(LoopbackDnsResolver(address))), + &HttpClientConfig::resolve(&HttpSettings::default()).unwrap(), Arc::new(TestAddressResolver { blocked_hosts }), ) .expect("test fetcher builds") @@ -552,9 +540,8 @@ mod tests { #[tokio::test] async fn rejects_url_credentials_before_network_access() { let fetcher = MediaFetcher::new( - &HttpClientPool::new(), - &HttpClientConfig::resolve(&litellm_http::HttpSettings::default(), None) - .expect("default settings resolve"), + &HttpClientPool::new(Arc::new(PublicDnsResolver)), + &HttpClientConfig::resolve(&HttpSettings::default()).expect("default settings resolve"), ) .expect("media fetcher builds"); let url = diff --git a/litellm-rust/crates/python-bridge/python_settings.json b/litellm-rust/crates/python-bridge/python_settings.json new file mode 100644 index 00000000000..a6cd959c6de --- /dev/null +++ b/litellm-rust/crates/python-bridge/python_settings.json @@ -0,0 +1,12 @@ +{ + "http_settings": [ + "ssl_verify", + "ssl_certificate", + "ssl_security_level", + "ssl_ecdh_curve", + "force_ipv4", + "http2", + "aiohttp_trust_env", + "user_agent" + ] +} diff --git a/litellm-rust/crates/python-bridge/src/http.rs b/litellm-rust/crates/python-bridge/src/http.rs index d2e05c3b949..de6fb5bb96d 100644 --- a/litellm-rust/crates/python-bridge/src/http.rs +++ b/litellm-rust/crates/python-bridge/src/http.rs @@ -1,11 +1,16 @@ -use std::{path::PathBuf, sync::LazyLock}; +use std::{ + path::PathBuf, + sync::{Arc, LazyLock}, +}; use litellm_http::{HttpClientConfig, HttpClientPool, HttpSettings, SslVerify}; +use litellm_llms::custom_httpx::media::PublicDnsResolver; use pyo3::{prelude::*, types::PyDict}; -use crate::errors::RustBridgeDeclined; +use crate::{errors::RustBridgeDeclined, python_settings::PythonSettings}; -static POOL: LazyLock = LazyLock::new(HttpClientPool::new); +static POOL: LazyLock = + LazyLock::new(|| HttpClientPool::new(Arc::new(PublicDnsResolver))); /// Keyword arguments that carry a live Python HTTP client or session. They cannot cross into /// Rust, so a call that supplies one stays on the Python path. @@ -15,21 +20,16 @@ pub(crate) fn pool() -> &'static HttpClientPool { &POOL } -/// The client configuration for one call: the process settings from the `litellm` module and -/// the environment, narrowed by the call's own `ssl_verify`. +/// The client configuration for one call: the `litellm.*` HTTP settings with the environment +/// overlaid, the same way `http_handler.py` combines them. pub(crate) fn call_config( py: Python<'_>, kwargs: &Bound<'_, PyDict>, ) -> PyResult { decline_live_clients(kwargs)?; - let settings = settings(py.import("litellm")?.as_any())? + let settings = settings(&PythonSettings::Http.read(py)?)? .with_environment(&|name| std::env::var(name).ok()); - let per_call = kwargs - .get_item("ssl_verify")? - .map(|value| ssl_verify(&value, "ssl_verify")) - .transpose()? - .flatten(); - HttpClientConfig::resolve(&settings, per_call.as_ref()) + HttpClientConfig::resolve(&settings) .map_err(|error| RustBridgeDeclined::new_err(error.to_string())) } @@ -44,45 +44,47 @@ pub(crate) fn decline_live_clients(kwargs: &Bound<'_, PyDict>) -> PyResult<()> { Ok(()) } -/// Read the `litellm.*` globals `http_handler.py` consults. `globals` is the `litellm` module in -/// production and any attribute holder in tests. -pub(crate) fn settings(globals: &Bound<'_, PyAny>) -> PyResult { +#[derive(FromPyObject)] +struct PythonHttpSettings<'py> { + ssl_verify: Bound<'py, PyAny>, + ssl_certificate: Option, + ssl_security_level: Option, + ssl_ecdh_curve: Option, + force_ipv4: bool, + http2: bool, + aiohttp_trust_env: bool, + user_agent: String, +} + +fn settings(value: &Bound<'_, PyAny>) -> PyResult { + let python: PythonHttpSettings = value.extract()?; Ok(HttpSettings { - ssl_verify: ssl_verify(&globals.getattr("ssl_verify")?, "litellm.ssl_verify")?, - ssl_certificate: optional_path(globals, "ssl_certificate")?, - ssl_security_level: globals.getattr("ssl_security_level")?.extract()?, - ssl_ecdh_curve: globals.getattr("ssl_ecdh_curve")?.extract()?, - force_ipv4: globals.getattr("force_ipv4")?.extract()?, - http2: globals.getattr("http2")?.extract()?, - trust_proxy_env: globals.getattr("aiohttp_trust_env")?.extract()?, + ssl_verify: Some(ssl_verify(&python.ssl_verify)?), + ssl_certificate: python.ssl_certificate.map(PathBuf::from), + ssl_security_level: python.ssl_security_level, + ssl_ecdh_curve: python.ssl_ecdh_curve, + force_ipv4: python.force_ipv4, + http2: python.http2, + user_agent: Some(python.user_agent), + trust_proxy_env: python.aiohttp_trust_env, ..HttpSettings::default() }) } -fn optional_path(globals: &Bound<'_, PyAny>, name: &str) -> PyResult> { - Ok(globals - .getattr(name)? - .extract::>()? - .map(PathBuf::from)) -} - -fn ssl_verify(value: &Bound<'_, PyAny>, name: &str) -> PyResult> { - if value.is_none() { - return Ok(None); - } +fn ssl_verify(value: &Bound<'_, PyAny>) -> PyResult { if let Ok(enabled) = value.extract::() { - return Ok(Some(if enabled { + return Ok(if enabled { SslVerify::Enabled } else { SslVerify::Disabled - })); + }); } if let Ok(path) = value.extract::() { - return Ok(Some(SslVerify::CaBundle(PathBuf::from(path)))); + return Ok(SslVerify::parse(&path)); } - Err(RustBridgeDeclined::new_err(format!( - "{name} is a live Python object and cannot be used by the Rust route" - ))) + Err(RustBridgeDeclined::new_err( + "litellm.ssl_verify is a live Python object and cannot be used by the Rust route", + )) } #[cfg(test)] @@ -91,18 +93,16 @@ mod tests { use rstest::rstest; use super::*; + use crate::python_settings::CONTRACT; - fn eval<'py>(py: Python<'py>, source: &std::ffi::CStr) -> Bound<'py, PyDict> { - let locals = PyDict::new(py); - py.run(source, Some(&locals), Some(&locals)).unwrap(); - locals - } - - fn globals<'py>(py: Python<'py>, overrides: &str) -> Bound<'py, PyAny> { + /// A stand-in for `http_settings()` carrying exactly the fields the contract declares, so a + /// field Rust reads but Python does not return fails here. + fn python_settings<'py>(py: Python<'py>, overrides: &str) -> Bound<'py, PyAny> { let source = format!( " +import json import types -globals = types.SimpleNamespace( +defaults = dict( ssl_verify=True, ssl_certificate=None, ssl_security_level=None, @@ -110,23 +110,29 @@ globals = types.SimpleNamespace( force_ipv4=False, http2=False, aiohttp_trust_env=False, + user_agent='litellm/test', ) -{overrides} +defaults.update(dict({overrides})) +settings = types.SimpleNamespace(**{{name: defaults[name] for name in json.loads(contract)['http_settings']}}) " ); + let locals = PyDict::new(py); + locals.set_item("contract", CONTRACT).unwrap(); let source = std::ffi::CString::new(source).unwrap(); - eval(py, &source).get_item("globals").unwrap().unwrap() + py.run(&source, Some(&locals), Some(&locals)).unwrap(); + locals.get_item("settings").unwrap().unwrap() } #[test] - fn default_globals_produce_default_settings_with_verification_on() { + fn default_python_settings_produce_default_settings_with_verification_on() { Python::initialize(); Python::attach(|py| { - let settings = settings(&globals(py, "")).unwrap(); + let settings = settings(&python_settings(py, "")).unwrap(); assert_eq!( settings, HttpSettings { ssl_verify: Some(SslVerify::Enabled), + user_agent: Some("litellm/test".into()), ..HttpSettings::default() } ); @@ -134,19 +140,20 @@ globals = types.SimpleNamespace( } #[test] - fn globals_flow_into_settings() { + fn python_settings_flow_into_settings() { Python::initialize(); Python::attach(|py| { - let settings = settings(&globals( + let settings = settings(&python_settings( py, " -globals.ssl_verify = '/etc/ssl/corp.pem' -globals.ssl_certificate = '/etc/ssl/client.pem' -globals.ssl_security_level = '2' -globals.ssl_ecdh_curve = 'X25519' -globals.force_ipv4 = True -globals.http2 = True -globals.aiohttp_trust_env = True +ssl_verify='/etc/ssl/corp.pem', +ssl_certificate='/etc/ssl/client.pem', +ssl_security_level='2', +ssl_ecdh_curve='X25519', +force_ipv4=True, +http2=True, +aiohttp_trust_env=True, +user_agent='litellm/9.9.9', ", )) .unwrap(); @@ -159,6 +166,7 @@ globals.aiohttp_trust_env = True ssl_ecdh_curve: Some("X25519".into()), force_ipv4: true, http2: true, + user_agent: Some("litellm/9.9.9".into()), trust_proxy_env: true, ..HttpSettings::default() } @@ -167,13 +175,32 @@ globals.aiohttp_trust_env = True } #[test] - fn disabled_verification_global_resolves_to_disabled() { + fn user_agent_environment_variable_beats_the_python_default() { Python::initialize(); Python::attach(|py| { - let settings = settings(&globals(py, "globals.ssl_verify = False")).unwrap(); - assert_eq!(settings.ssl_verify, Some(SslVerify::Disabled)); - let config = HttpClientConfig::resolve(&settings, None).unwrap(); - assert_eq!(config.verify, Verify::Disabled); + let settings = settings(&python_settings(py, "")) + .unwrap() + .with_environment(&|name| { + (name == "LITELLM_USER_AGENT").then(|| "operator/1".to_string()) + }); + assert_eq!(settings.user_agent.as_deref(), Some("operator/1")); + }); + } + + #[rstest] + #[case::disabled("ssl_verify=False", Verify::Disabled)] + #[case::disabled_string("ssl_verify='False'", Verify::Disabled)] + #[case::enabled_string("ssl_verify='true'", Verify::BuiltInRoots)] + #[case::bundle("ssl_verify='/tmp/ca.pem'", Verify::CaBundle("/tmp/ca.pem".into()))] + fn ssl_verify_global_resolves_like_get_ssl_verify( + #[case] overrides: &str, + #[case] expected: Verify, + ) { + Python::initialize(); + Python::attach(|py| { + let settings = settings(&python_settings(py, overrides)).unwrap(); + let config = HttpClientConfig::resolve(&settings).unwrap(); + assert_eq!(config.verify, expected); }); } @@ -181,7 +208,7 @@ globals.aiohttp_trust_env = True fn ssl_context_global_declines_instead_of_being_dropped() { Python::initialize(); Python::attach(|py| { - let error = settings(&globals(py, "globals.ssl_verify = object()")).unwrap_err(); + let error = settings(&python_settings(py, "ssl_verify=object()")).unwrap_err(); assert!(error.is_instance_of::(py)); assert!(error.value(py).to_string().contains("litellm.ssl_verify")); }); @@ -215,20 +242,4 @@ globals.aiohttp_trust_env = True decline_live_clients(&kwargs).unwrap(); }); } - - #[rstest] - #[case::disabled(c"False", Some(SslVerify::Disabled))] - #[case::enabled(c"True", Some(SslVerify::Enabled))] - #[case::bundle(c"'/tmp/ca.pem'", Some(SslVerify::CaBundle("/tmp/ca.pem".into())))] - #[case::unset(c"None", None)] - fn per_call_ssl_verify_values_project( - #[case] source: &std::ffi::CStr, - #[case] expected: Option, - ) { - Python::initialize(); - Python::attach(|py| { - let value = py.eval(source, None, None).unwrap(); - assert_eq!(ssl_verify(&value, "ssl_verify").unwrap(), expected); - }); - } } diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 11cb0a7f655..7eba0d201be 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -3,6 +3,7 @@ mod diagnostics; mod errors; mod http; mod marshal; +mod python_settings; mod routes; mod token_counter; diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs new file mode 100644 index 00000000000..0c6554f0970 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -0,0 +1,48 @@ +use pyo3::prelude::*; + +const MODULE: &str = "litellm.rust_bridge.settings"; + +/// Every group of `litellm.*` module globals the native routes read. Environment overrides are +/// applied on the Rust side, so each function returns only what the Python process configured. +/// A group is deleted once Rust owns loading that configuration, so this enum only shrinks. +/// +/// `litellm/rust_bridge/settings.py` is the only Python module behind it, and +/// `python_settings.json` pins the fields each function returns on both sides. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum PythonSettings { + Http, +} + +impl PythonSettings { + #[cfg(test)] + pub(crate) const ALL: [Self; 1] = [Self::Http]; + + pub(crate) fn name(self) -> &'static str { + match self { + Self::Http => "http_settings", + } + } + + pub(crate) fn read(self, py: Python<'_>) -> PyResult> { + py.import(MODULE)?.getattr(self.name())?.call0() + } +} + +#[cfg(test)] +pub(crate) const CONTRACT: &str = include_str!("../python_settings.json"); + +#[cfg(test)] +mod tests { + use std::collections::BTreeSet; + + use super::{CONTRACT, PythonSettings}; + + #[test] + fn every_settings_group_is_in_the_python_contract() { + let contract: serde_json::Map = + serde_json::from_str(CONTRACT).unwrap(); + let declared: BTreeSet<&str> = contract.keys().map(String::as_str).collect(); + let read: BTreeSet<&str> = PythonSettings::ALL.map(PythonSettings::name).into(); + assert_eq!(read, declared); + } +} diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 05dff0cb9d8..6b90394043f 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -150,7 +150,11 @@ def get_default_headers() -> dict: if user_agent is not None: return {"User-Agent": user_agent} - return {"User-Agent": f"litellm/{version}"} + return {"User-Agent": default_user_agent()} + + +def default_user_agent() -> str: + return f"litellm/{version}" # Initialize headers (User-Agent) diff --git a/litellm/rust_bridge/settings.py b/litellm/rust_bridge/settings.py new file mode 100644 index 00000000000..a8229b12d13 --- /dev/null +++ b/litellm/rust_bridge/settings.py @@ -0,0 +1,37 @@ +"""The `litellm.*` module globals the native routes read. + +Environment variables that override these are applied in Rust, so nothing here reads `os.environ`. +`litellm-rust/crates/python-bridge/python_settings.json` pins the fields each function returns. +""" + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True, slots=True) +class HttpSettings: + ssl_verify: bool | str + ssl_certificate: str | None + ssl_security_level: str | None + ssl_ecdh_curve: str | None + force_ipv4: bool + http2: bool + aiohttp_trust_env: bool + user_agent: str + + +def http_settings() -> HttpSettings: + import litellm + from litellm.llms.custom_httpx.http_handler import default_user_agent + + return HttpSettings( + ssl_verify=litellm.ssl_verify, + ssl_certificate=litellm.ssl_certificate, + ssl_security_level=litellm.ssl_security_level, + ssl_ecdh_curve=litellm.ssl_ecdh_curve, + force_ipv4=litellm.force_ipv4, + http2=litellm.http2, + aiohttp_trust_env=litellm.aiohttp_trust_env, + user_agent=default_user_agent(), + ) diff --git a/tests/test_litellm/rust_bridge/test_settings.py b/tests/test_litellm/rust_bridge/test_settings.py new file mode 100644 index 00000000000..618fa400136 --- /dev/null +++ b/tests/test_litellm/rust_bridge/test_settings.py @@ -0,0 +1,50 @@ +import dataclasses +from pathlib import Path +from typing import Final + +import pytest +from pydantic import TypeAdapter + +import litellm +from litellm.llms.custom_httpx.http_handler import default_user_agent +from litellm.rust_bridge import settings + +CONTRACT_PATH: Final = Path(__file__).parents[3] / "litellm-rust/crates/python-bridge/python_settings.json" + + +def test_the_rust_contract_matches_the_returned_fields() -> None: + contract: Final = TypeAdapter(dict[str, list[str]]).validate_json(CONTRACT_PATH.read_text()) + + assert contract == {"http_settings": [field.name for field in dataclasses.fields(settings.http_settings())]} + + +def test_http_settings_reads_the_litellm_globals(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "ssl_verify", "/etc/ssl/corp.pem") + monkeypatch.setattr(litellm, "ssl_certificate", "/etc/ssl/client.pem") + monkeypatch.setattr(litellm, "ssl_security_level", "DEFAULT@SECLEVEL=1") + monkeypatch.setattr(litellm, "ssl_ecdh_curve", "X25519") + monkeypatch.setattr(litellm, "force_ipv4", True) + monkeypatch.setattr(litellm, "http2", True) + monkeypatch.setattr(litellm, "aiohttp_trust_env", True) + + assert settings.http_settings() == settings.HttpSettings( + ssl_verify="/etc/ssl/corp.pem", + ssl_certificate="/etc/ssl/client.pem", + ssl_security_level="DEFAULT@SECLEVEL=1", + ssl_ecdh_curve="X25519", + force_ipv4=True, + http2=True, + aiohttp_trust_env=True, + user_agent=settings.http_settings().user_agent, + ) + + +def test_http_settings_ignores_environment_overrides(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_USER_AGENT", "operator/1") + monkeypatch.setenv("SSL_VERIFY", "false") + monkeypatch.setattr(litellm, "ssl_verify", True) + + result: Final = settings.http_settings() + + assert result.user_agent == default_user_agent() + assert result.ssl_verify is True