mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
Merge pull request #41897 from BerriAI/litellm_rust_http_pool_ocr
feat(rust): add litellm-http client pool and inject it into the OCR route
This commit is contained in:
yujonglee
GitHub
commit
b00d066ec2
29 changed files with 2409 additions and 91 deletions
19
litellm-rust/Cargo.lock
generated
19
litellm-rust/Cargo.lock
generated
|
|
@ -2050,8 +2050,10 @@ dependencies = [
|
|||
"futures-util",
|
||||
"litellm-auth",
|
||||
"litellm-auth-aws",
|
||||
"litellm-auth-gcp",
|
||||
"litellm-core-utils",
|
||||
"litellm-host",
|
||||
"litellm-http",
|
||||
"litellm-llms",
|
||||
"litellm-types",
|
||||
"mime_guess",
|
||||
|
|
@ -2129,6 +2131,20 @@ dependencies = [
|
|||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-http"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"http 1.4.2",
|
||||
"hyper-util",
|
||||
"reqwest 0.12.28",
|
||||
"rstest",
|
||||
"rustls 0.23.42",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
"webpki-roots",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-llms"
|
||||
version = "0.1.0"
|
||||
|
|
@ -2146,6 +2162,7 @@ dependencies = [
|
|||
"litellm-core-utils",
|
||||
"litellm-framing",
|
||||
"litellm-host",
|
||||
"litellm-http",
|
||||
"litellm-types",
|
||||
"reqwest 0.12.28",
|
||||
"rstest",
|
||||
|
|
@ -2167,9 +2184,11 @@ dependencies = [
|
|||
"criterion",
|
||||
"futures-util",
|
||||
"litellm-auth",
|
||||
"litellm-auth-gcp",
|
||||
"litellm-callbacks-legacy",
|
||||
"litellm-core",
|
||||
"litellm-host-python",
|
||||
"litellm-http",
|
||||
"litellm-llms",
|
||||
"litellm-token-counter",
|
||||
"litellm-types",
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ litellm-auth = { path = "crates/auth" }
|
|||
litellm-auth-aws = { path = "crates/auth-aws" }
|
||||
litellm-auth-azure = { path = "crates/auth-azure" }
|
||||
litellm-auth-gcp = { path = "crates/auth-gcp" }
|
||||
litellm-http = { path = "crates/http" }
|
||||
litellm-llms = { path = "crates/llms" }
|
||||
litellm-types = { path = "crates/types" }
|
||||
litellm-core-utils = { path = "crates/core-utils" }
|
||||
|
|
@ -26,6 +27,8 @@ litellm-token-counter = { path = "crates/token-counter" }
|
|||
litellm-host-python = { path = "crates/host-python" }
|
||||
|
||||
bytes = "1"
|
||||
http = "1"
|
||||
hyper-util = { version = "0.1.20", default-features = false, features = ["client-proxy"] }
|
||||
proptest = "1.7.0"
|
||||
pyo3 = "0.29.2"
|
||||
pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
|
||||
|
|
@ -49,6 +52,7 @@ base64 = "0.22"
|
|||
moka = { version = "0.12.16", features = ["future"] }
|
||||
strum = { version = "0.28.0", features = ["derive"] }
|
||||
url = "2.5.8"
|
||||
webpki-roots = "1"
|
||||
time = { version = "0.3.53", features = ["parsing"] }
|
||||
criterion = "0.8.2"
|
||||
fancy-regex = "0.19.2"
|
||||
|
|
|
|||
|
|
@ -35,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
|
||||
|
|
|
|||
|
|
@ -14,7 +14,3 @@ pub async fn perform(
|
|||
) -> Result<LiteLLMOcrResponse, Error> {
|
||||
litellm_host::run::run(ocr_machine(client.clone()), &LocalOcrHost::new(request)).await
|
||||
}
|
||||
|
||||
pub async fn ocr(request: LiteLLMOcrRequest) -> Result<LiteLLMOcrResponse, Error> {
|
||||
perform(&OcrClient::shared()?, request).await
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,16 +1,21 @@
|
|||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use litellm_auth_gcp::VertexAuth;
|
||||
use litellm_host::{
|
||||
event::{CallEvent, MachineEvent, WireRequest},
|
||||
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,
|
||||
custom_httpx::{
|
||||
llm_http_handler::OcrClient,
|
||||
media::{PublicDnsResolver, UrlPolicy},
|
||||
},
|
||||
};
|
||||
use rstest::rstest;
|
||||
use serde_json::{Value, json};
|
||||
|
|
@ -171,25 +176,24 @@ async fn facade_retains_native_response_when_requested() {
|
|||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn facade_uses_the_injected_http_client() {
|
||||
async fn ocr_client_uses_the_injected_http_pool_configuration() {
|
||||
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await;
|
||||
let mut default_headers = reqwest::header::HeaderMap::new();
|
||||
default_headers.insert(
|
||||
"x-transport-owner",
|
||||
reqwest::header::HeaderValue::from_static("host"),
|
||||
);
|
||||
let provider_http = reqwest::Client::builder()
|
||||
.default_headers(default_headers)
|
||||
.build()
|
||||
.unwrap();
|
||||
crate::ocr::client::perform(
|
||||
&OcrClient::new(provider_http).unwrap(),
|
||||
wire_request("mistral/model", &base, json!({})),
|
||||
let settings = HttpSettings {
|
||||
user_agent: Some("host-owned/1".into()),
|
||||
..HttpSettings::default()
|
||||
};
|
||||
let client = OcrClient::new(
|
||||
&HttpClientPool::new(Arc::new(PublicDnsResolver)),
|
||||
&Resolution::from(&settings).config,
|
||||
UrlPolicy::default(),
|
||||
VertexAuth::default(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
crate::ocr::client::perform(&client, wire_request("mistral/model", &base, json!({})))
|
||||
.await
|
||||
.unwrap();
|
||||
server.await.unwrap();
|
||||
assert!(seen.lock().unwrap()[0].contains("x-transport-owner: host"));
|
||||
assert!(seen.lock().unwrap()[0].contains("user-agent: host-owned/1"));
|
||||
}
|
||||
|
||||
fn event_name(event: &CallEvent) -> &'static str {
|
||||
|
|
|
|||
1
litellm-rust/crates/http/AGENTS.md
Normal file
1
litellm-rust/crates/http/AGENTS.md
Normal file
|
|
@ -0,0 +1 @@
|
|||
- https://github.com/BerriAI/litellm-docs/blob/main/docs/guides/security_settings.md
|
||||
18
litellm-rust/crates/http/Cargo.toml
Normal file
18
litellm-rust/crates/http/Cargo.toml
Normal file
|
|
@ -0,0 +1,18 @@
|
|||
[package]
|
||||
name = "litellm-http"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
http.workspace = true
|
||||
hyper-util.workspace = true
|
||||
reqwest.workspace = true
|
||||
rustls.workspace = true
|
||||
thiserror.workspace = true
|
||||
webpki-roots.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
rstest.workspace = true
|
||||
tokio.workspace = true
|
||||
296
litellm-rust/crates/http/src/config.rs
Normal file
296
litellm-rust/crates/http/src/config.rs
Normal file
|
|
@ -0,0 +1,296 @@
|
|||
use std::{
|
||||
net::{IpAddr, Ipv4Addr},
|
||||
path::PathBuf,
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use crate::{
|
||||
error::Error,
|
||||
settings::{HttpSettings, SslVerify, TcpKeepalive},
|
||||
tls::{CipherSelection, KeyExchangeGroup, Tls12CipherSuite, Unsupported},
|
||||
};
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
|
||||
pub enum Verify {
|
||||
Disabled,
|
||||
CaBundle(PathBuf),
|
||||
BuiltInRoots,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
|
||||
pub struct HttpClientConfig {
|
||||
pub verify: Verify,
|
||||
pub client_certificate: Option<PathBuf>,
|
||||
pub key_exchange_group: Option<KeyExchangeGroup>,
|
||||
pub tls12_cipher_suites: Option<Vec<Tls12CipherSuite>>,
|
||||
pub force_ipv4: bool,
|
||||
pub http2: bool,
|
||||
pub user_agent: Option<String>,
|
||||
pub trust_proxy_env: bool,
|
||||
pub connect_timeout: Duration,
|
||||
pub tcp_keepalive: Option<TcpKeepalive>,
|
||||
pub pool_idle_timeout: Duration,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct Resolution {
|
||||
pub config: HttpClientConfig,
|
||||
pub unsupported: Vec<Unsupported>,
|
||||
}
|
||||
|
||||
impl From<&HttpSettings> for Verify {
|
||||
fn from(settings: &HttpSettings) -> Self {
|
||||
match &settings.ssl_verify {
|
||||
Some(SslVerify::Disabled) => Self::Disabled,
|
||||
Some(SslVerify::CaBundle(path)) => Self::CaBundle(path.clone()),
|
||||
Some(SslVerify::Enabled) | None => settings
|
||||
.ssl_cert_file
|
||||
.clone()
|
||||
.map_or(Self::BuiltInRoots, Self::CaBundle),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<&HttpSettings> for Resolution {
|
||||
fn from(settings: &HttpSettings) -> Self {
|
||||
let curve = settings
|
||||
.ssl_ecdh_curve
|
||||
.as_deref()
|
||||
.map(str::parse::<KeyExchangeGroup>)
|
||||
.transpose();
|
||||
let ciphers = settings
|
||||
.ssl_security_level
|
||||
.as_deref()
|
||||
.map(CipherSelection::from)
|
||||
.unwrap_or_default();
|
||||
Self {
|
||||
config: HttpClientConfig {
|
||||
verify: Verify::from(settings),
|
||||
client_certificate: settings.ssl_certificate.clone(),
|
||||
key_exchange_group: curve.clone().ok().flatten(),
|
||||
tls12_cipher_suites: ciphers.tls12_cipher_suites,
|
||||
force_ipv4: settings.force_ipv4,
|
||||
http2: settings.http2,
|
||||
user_agent: settings.user_agent.clone(),
|
||||
trust_proxy_env: settings.trust_proxy_env,
|
||||
connect_timeout: settings.connect_timeout,
|
||||
tcp_keepalive: settings.tcp_keepalive,
|
||||
pool_idle_timeout: settings.pool_idle_timeout,
|
||||
},
|
||||
unsupported: curve.err().into_iter().chain(ciphers.unsupported).collect(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl TryFrom<&HttpClientConfig> for reqwest::ClientBuilder {
|
||||
type Error = Error;
|
||||
|
||||
fn try_from(config: &HttpClientConfig) -> Result<Self, Self::Error> {
|
||||
let base = reqwest::Client::builder()
|
||||
.use_preconfigured_tls(rustls::ClientConfig::try_from(config)?)
|
||||
.connect_timeout(config.connect_timeout)
|
||||
.pool_idle_timeout(config.pool_idle_timeout);
|
||||
let with_keepalive = match config.tcp_keepalive {
|
||||
None => base,
|
||||
Some(keepalive) => base
|
||||
.tcp_keepalive(keepalive.idle)
|
||||
.tcp_keepalive_interval(keepalive.interval)
|
||||
.tcp_keepalive_retries(keepalive.retries),
|
||||
};
|
||||
let with_address = if config.force_ipv4 {
|
||||
with_keepalive.local_address(IpAddr::V4(Ipv4Addr::UNSPECIFIED))
|
||||
} else {
|
||||
with_keepalive
|
||||
};
|
||||
let with_protocol = if config.http2 {
|
||||
with_address
|
||||
} else {
|
||||
with_address.http1_only()
|
||||
};
|
||||
let with_agent = match &config.user_agent {
|
||||
Some(agent) => with_protocol.user_agent(agent),
|
||||
None => with_protocol,
|
||||
};
|
||||
Ok(if config.trust_proxy_env {
|
||||
with_agent
|
||||
} else {
|
||||
with_agent.no_proxy()
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn settings(ssl_verify: Option<SslVerify>, ssl_cert_file: Option<&str>) -> HttpSettings {
|
||||
HttpSettings {
|
||||
ssl_verify,
|
||||
ssl_cert_file: ssl_cert_file.map(PathBuf::from),
|
||||
..HttpSettings::default()
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::default(settings(None, None), Verify::BuiltInRoots)]
|
||||
#[case::setting_disables(
|
||||
settings(Some(SslVerify::Disabled), Some("/env/roots.pem")),
|
||||
Verify::Disabled
|
||||
)]
|
||||
#[case::setting_bundle(
|
||||
settings(Some(SslVerify::CaBundle("/configured.pem".into())), Some("/env/roots.pem")),
|
||||
Verify::CaBundle("/configured.pem".into())
|
||||
)]
|
||||
#[case::enabled_uses_cert_file(
|
||||
settings(Some(SslVerify::Enabled), Some("/env/roots.pem")),
|
||||
Verify::CaBundle("/env/roots.pem".into())
|
||||
)]
|
||||
#[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] expected: Verify,
|
||||
) {
|
||||
let config = Resolution::from(&settings).config;
|
||||
assert_eq!(config.verify, expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::x25519("X25519", Some(KeyExchangeGroup::X25519))]
|
||||
#[case::openssl_p256("prime256v1", Some(KeyExchangeGroup::Secp256r1))]
|
||||
#[case::p384("secp384r1", Some(KeyExchangeGroup::Secp384r1))]
|
||||
fn ecdh_curve_selects_the_single_key_exchange_group(
|
||||
#[case] curve: &str,
|
||||
#[case] expected: Option<KeyExchangeGroup>,
|
||||
) {
|
||||
let settings = HttpSettings {
|
||||
ssl_ecdh_curve: Some(curve.into()),
|
||||
..HttpSettings::default()
|
||||
};
|
||||
let resolution = Resolution::from(&settings);
|
||||
assert_eq!(resolution.config.key_exchange_group, expected);
|
||||
assert_eq!(resolution.unsupported, []);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unsupported_ecdh_curve_keeps_the_defaults_and_is_reported() {
|
||||
let settings = HttpSettings {
|
||||
ssl_ecdh_curve: Some("secp521r1".into()),
|
||||
..HttpSettings::default()
|
||||
};
|
||||
let resolution = Resolution::from(&settings);
|
||||
assert_eq!(resolution.config.key_exchange_group, None);
|
||||
assert_eq!(
|
||||
resolution.unsupported,
|
||||
[Unsupported::EcdhCurve("secp521r1".into())]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn legacy_security_level_keeps_every_suite_and_is_reported_unsupported() {
|
||||
let settings = HttpSettings {
|
||||
ssl_security_level: Some("DEFAULT@SECLEVEL=1".into()),
|
||||
..HttpSettings::default()
|
||||
};
|
||||
let resolution = Resolution::from(&settings);
|
||||
assert_eq!(resolution.config.tls12_cipher_suites, None);
|
||||
assert_eq!(
|
||||
resolution.unsupported,
|
||||
[Unsupported::SecurityLevel("@SECLEVEL=1".into())]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn named_suites_restrict_tls12_and_unsupported_entries_are_reported() {
|
||||
let settings = HttpSettings {
|
||||
ssl_security_level: Some(
|
||||
"ECDHE-RSA-AES256-GCM-SHA384:ECDHE-ECDSA-AES128-GCM-SHA256:!aNULL:AES256-SHA@SECLEVEL=2"
|
||||
.into(),
|
||||
),
|
||||
..HttpSettings::default()
|
||||
};
|
||||
let resolution = Resolution::from(&settings);
|
||||
assert_eq!(
|
||||
resolution.config.tls12_cipher_suites,
|
||||
Some(vec![
|
||||
Tls12CipherSuite::EcdheEcdsaAes128Gcm,
|
||||
Tls12CipherSuite::EcdheRsaAes256Gcm
|
||||
])
|
||||
);
|
||||
assert_eq!(
|
||||
resolution.unsupported,
|
||||
[
|
||||
Unsupported::CipherToken("!aNULL".into()),
|
||||
Unsupported::CipherToken("AES256-SHA".into())
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn connection_settings_carry_over_unchanged() {
|
||||
let keepalive = TcpKeepalive {
|
||||
idle: Duration::from_secs(60),
|
||||
interval: Duration::from_secs(30),
|
||||
retries: 5,
|
||||
};
|
||||
let settings = HttpSettings {
|
||||
ssl_certificate: Some("/client.pem".into()),
|
||||
force_ipv4: true,
|
||||
http2: true,
|
||||
user_agent: Some("litellm/1.0".into()),
|
||||
trust_proxy_env: true,
|
||||
connect_timeout: Duration::from_secs(7),
|
||||
tcp_keepalive: Some(keepalive),
|
||||
pool_idle_timeout: Duration::from_secs(45),
|
||||
..HttpSettings::default()
|
||||
};
|
||||
let config = Resolution::from(&settings).config;
|
||||
assert_eq!(
|
||||
config,
|
||||
HttpClientConfig {
|
||||
verify: Verify::BuiltInRoots,
|
||||
client_certificate: Some("/client.pem".into()),
|
||||
key_exchange_group: None,
|
||||
tls12_cipher_suites: None,
|
||||
force_ipv4: true,
|
||||
http2: true,
|
||||
user_agent: Some("litellm/1.0".into()),
|
||||
trust_proxy_env: true,
|
||||
connect_timeout: Duration::from_secs(7),
|
||||
tcp_keepalive: Some(keepalive),
|
||||
pool_idle_timeout: Duration::from_secs(45),
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_ca_bundle_is_a_read_error() {
|
||||
let path = std::env::temp_dir().join("litellm-http-missing-bundle.pem");
|
||||
let config = HttpClientConfig {
|
||||
verify: Verify::CaBundle(path.clone()),
|
||||
..Resolution::from(&HttpSettings::default()).config
|
||||
};
|
||||
assert!(matches!(
|
||||
reqwest::ClientBuilder::try_from(&config),
|
||||
Err(Error::Read { path: reported, .. }) if reported == path
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_pem_ca_bundle_is_an_invalid_pem_error() {
|
||||
let path =
|
||||
std::env::temp_dir().join(format!("litellm-http-not-pem-{}.pem", std::process::id()));
|
||||
std::fs::write(&path, b"not a certificate").unwrap();
|
||||
let config = HttpClientConfig {
|
||||
verify: Verify::CaBundle(path.clone()),
|
||||
..Resolution::from(&HttpSettings::default()).config
|
||||
};
|
||||
let result = reqwest::ClientBuilder::try_from(&config).map(drop);
|
||||
std::fs::remove_file(&path).unwrap();
|
||||
assert!(matches!(
|
||||
result,
|
||||
Err(Error::InvalidPem { path: reported, .. }) if reported == path
|
||||
));
|
||||
}
|
||||
}
|
||||
17
litellm-rust/crates/http/src/error.rs
Normal file
17
litellm-rust/crates/http/src/error.rs
Normal file
|
|
@ -0,0 +1,17 @@
|
|||
use std::path::PathBuf;
|
||||
|
||||
#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)]
|
||||
pub enum Error {
|
||||
#[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<reqwest::Error> for Error {
|
||||
fn from(error: reqwest::Error) -> Self {
|
||||
Self::Client(error.without_url().to_string())
|
||||
}
|
||||
}
|
||||
13
litellm-rust/crates/http/src/lib.rs
Normal file
13
litellm-rust/crates/http/src/lib.rs
Normal file
|
|
@ -0,0 +1,13 @@
|
|||
mod config;
|
||||
mod error;
|
||||
mod pool;
|
||||
mod proxy;
|
||||
mod settings;
|
||||
mod tls;
|
||||
|
||||
pub use config::{HttpClientConfig, Resolution, Verify};
|
||||
pub use error::Error;
|
||||
pub use pool::{ClientVariant, HttpClientPool};
|
||||
pub use proxy::EnvironmentProxies;
|
||||
pub use settings::{HttpSettings, HttpSettingsLayer, SslVerify, TcpKeepalive};
|
||||
pub use tls::{KeyExchangeGroup, Tls12CipherSuite, Unsupported};
|
||||
325
litellm-rust/crates/http/src/pool.rs
Normal file
325
litellm-rust/crates/http/src/pool.rs
Normal file
|
|
@ -0,0 +1,325 @@
|
|||
use std::{
|
||||
collections::HashMap,
|
||||
sync::{Arc, Mutex, MutexGuard, PoisonError},
|
||||
time::{Duration, Instant},
|
||||
};
|
||||
|
||||
use reqwest::dns::Resolve;
|
||||
|
||||
use crate::{config::HttpClientConfig, error::Error};
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
|
||||
pub enum ClientVariant {
|
||||
Provider,
|
||||
NoRedirect,
|
||||
Media,
|
||||
UnpinnedMedia,
|
||||
}
|
||||
|
||||
const CLIENT_TTL: Duration = Duration::from_secs(3600);
|
||||
|
||||
struct PooledClient {
|
||||
client: reqwest::Client,
|
||||
built_at: Instant,
|
||||
}
|
||||
|
||||
type Clients = HashMap<(HttpClientConfig, ClientVariant), PooledClient>;
|
||||
|
||||
pub struct HttpClientPool {
|
||||
media_resolver: Arc<dyn Resolve>,
|
||||
ttl: Duration,
|
||||
clients: Mutex<Clients>,
|
||||
}
|
||||
|
||||
impl HttpClientPool {
|
||||
pub fn new(media_resolver: Arc<dyn Resolve>) -> Self {
|
||||
Self::with_ttl(media_resolver, CLIENT_TTL)
|
||||
}
|
||||
|
||||
pub fn with_ttl(media_resolver: Arc<dyn Resolve>, ttl: Duration) -> Self {
|
||||
Self {
|
||||
media_resolver,
|
||||
ttl,
|
||||
clients: Mutex::default(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn client(
|
||||
&self,
|
||||
config: &HttpClientConfig,
|
||||
variant: ClientVariant,
|
||||
) -> Result<reqwest::Client, Error> {
|
||||
let effective = match variant {
|
||||
ClientVariant::Media => HttpClientConfig {
|
||||
client_certificate: None,
|
||||
trust_proxy_env: false,
|
||||
..config.clone()
|
||||
},
|
||||
ClientVariant::UnpinnedMedia => HttpClientConfig {
|
||||
client_certificate: None,
|
||||
..config.clone()
|
||||
},
|
||||
ClientVariant::Provider | ClientVariant::NoRedirect => config.clone(),
|
||||
};
|
||||
let key = (effective, variant);
|
||||
if let Some(pooled) = self.lock().get(&key)
|
||||
&& pooled.built_at.elapsed() < self.ttl
|
||||
{
|
||||
return Ok(pooled.client.clone());
|
||||
}
|
||||
let client = self
|
||||
.apply(variant, reqwest::ClientBuilder::try_from(&key.0)?)
|
||||
.build()?;
|
||||
self.lock().insert(
|
||||
key,
|
||||
PooledClient {
|
||||
client: client.clone(),
|
||||
built_at: Instant::now(),
|
||||
},
|
||||
);
|
||||
Ok(client)
|
||||
}
|
||||
|
||||
fn lock(&self) -> MutexGuard<'_, Clients> {
|
||||
self.clients.lock().unwrap_or_else(PoisonError::into_inner)
|
||||
}
|
||||
|
||||
fn apply(
|
||||
&self,
|
||||
variant: ClientVariant,
|
||||
builder: reqwest::ClientBuilder,
|
||||
) -> reqwest::ClientBuilder {
|
||||
match variant {
|
||||
ClientVariant::Provider => builder,
|
||||
ClientVariant::NoRedirect | ClientVariant::UnpinnedMedia => {
|
||||
builder.redirect(reqwest::redirect::Policy::none())
|
||||
}
|
||||
ClientVariant::Media => builder
|
||||
.redirect(reqwest::redirect::Policy::none())
|
||||
.dns_resolver2(Arc::clone(&self.media_resolver)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::{
|
||||
net::SocketAddr,
|
||||
sync::atomic::{AtomicUsize, Ordering},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use reqwest::dns::{Addrs, Name, Resolving};
|
||||
use tokio::{
|
||||
io::{AsyncReadExt, AsyncWriteExt},
|
||||
net::TcpListener,
|
||||
};
|
||||
|
||||
use super::*;
|
||||
use crate::{HttpSettings, Resolution, Verify};
|
||||
|
||||
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)))
|
||||
}
|
||||
}
|
||||
|
||||
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()),
|
||||
..Resolution::from(&HttpSettings::default()).config
|
||||
}
|
||||
}
|
||||
|
||||
async fn serve(
|
||||
status_line: &'static str,
|
||||
) -> (SocketAddr, Arc<AtomicUsize>, Arc<Mutex<Vec<String>>>) {
|
||||
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()
|
||||
.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);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn expired_clients_are_rebuilt() {
|
||||
let (address, connections, _) = serve("HTTP/1.1 204 No Content").await;
|
||||
let url = format!("http://{address}");
|
||||
let pool = HttpClientPool::with_ttl(
|
||||
Arc::new(FixedResolver(([192, 0, 2, 1], 80).into())),
|
||||
Duration::ZERO,
|
||||
);
|
||||
get(&pool, &config("a"), ClientVariant::Provider, &url).await;
|
||||
get(&pool, &config("a"), ClientVariant::Provider, &url).await;
|
||||
assert_eq!(connections.load(Ordering::SeqCst), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn media_clients_are_shared_across_proxy_settings_they_never_use() {
|
||||
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] {
|
||||
let config = HttpClientConfig {
|
||||
trust_proxy_env,
|
||||
..config("a")
|
||||
};
|
||||
get(&pool, &config, ClientVariant::Media, &url).await;
|
||||
}
|
||||
assert_eq!(connections.load(Ordering::SeqCst), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn media_variant_never_loads_the_client_certificate() {
|
||||
let pool = pool();
|
||||
let with_identity = HttpClientConfig {
|
||||
client_certificate: Some(std::env::temp_dir().join("litellm-http-absent-client.pem")),
|
||||
..config("a")
|
||||
};
|
||||
assert!(
|
||||
pool.client(&with_identity, ClientVariant::Provider)
|
||||
.is_err()
|
||||
);
|
||||
assert!(pool.client(&with_identity, ClientVariant::Media).is_ok());
|
||||
assert!(
|
||||
pool.client(&with_identity, ClientVariant::UnpinnedMedia)
|
||||
.is_ok()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_failures_are_not_cached() {
|
||||
let pool = pool();
|
||||
let missing = HttpClientConfig {
|
||||
verify: Verify::CaBundle(std::env::temp_dir().join("litellm-http-absent.pem")),
|
||||
..config("a")
|
||||
};
|
||||
assert!(pool.client(&missing, ClientVariant::Provider).is_err());
|
||||
assert!(pool.client(&missing, ClientVariant::Provider).is_err());
|
||||
assert!(pool.client(&config("a"), ClientVariant::Provider).is_ok());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn provider_client_sends_the_configured_user_agent_over_http1() {
|
||||
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 = 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 (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");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unpinned_media_variant_uses_the_system_resolver_and_returns_redirects() {
|
||||
let (address, _, _) = serve("HTTP/1.1 302 Found").await;
|
||||
let pool = HttpClientPool::new(Arc::new(FixedResolver(([192, 0, 2, 1], 80).into())));
|
||||
let response = get(
|
||||
&pool,
|
||||
&config("a"),
|
||||
ClientVariant::UnpinnedMedia,
|
||||
&format!("http://localhost:{}/doc", address.port()),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(response.status(), 302);
|
||||
}
|
||||
|
||||
#[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()
|
||||
);
|
||||
}
|
||||
}
|
||||
15
litellm-rust/crates/http/src/proxy.rs
Normal file
15
litellm-rust/crates/http/src/proxy.rs
Normal file
|
|
@ -0,0 +1,15 @@
|
|||
use hyper_util::client::proxy::matcher::Matcher;
|
||||
|
||||
pub struct EnvironmentProxies(Matcher);
|
||||
|
||||
impl EnvironmentProxies {
|
||||
pub fn from_environment() -> Self {
|
||||
Self(Matcher::from_system())
|
||||
}
|
||||
|
||||
pub fn apply_to(&self, url: &reqwest::Url) -> bool {
|
||||
url.as_str()
|
||||
.parse::<http::Uri>()
|
||||
.is_ok_and(|uri| self.0.intercept(&uri).is_some())
|
||||
}
|
||||
}
|
||||
416
litellm-rust/crates/http/src/settings.rs
Normal file
416
litellm-rust/crates/http/src/settings.rs
Normal file
|
|
@ -0,0 +1,416 @@
|
|||
use std::{
|
||||
path::{Path, PathBuf},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
|
||||
pub enum SslVerify {
|
||||
Enabled,
|
||||
Disabled,
|
||||
CaBundle(PathBuf),
|
||||
}
|
||||
|
||||
impl SslVerify {
|
||||
pub fn parse(value: &str) -> Self {
|
||||
match value.trim().to_ascii_lowercase().as_str() {
|
||||
"true" => Self::Enabled,
|
||||
"false" => Self::Disabled,
|
||||
_ => Self::CaBundle(PathBuf::from(value)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
|
||||
pub struct TcpKeepalive {
|
||||
pub idle: Duration,
|
||||
pub interval: Duration,
|
||||
pub retries: u32,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, PartialEq, Eq)]
|
||||
pub struct HttpSettingsLayer {
|
||||
pub ssl_verify: Option<SslVerify>,
|
||||
pub ssl_cert_file: Option<PathBuf>,
|
||||
pub ssl_certificate: Option<PathBuf>,
|
||||
pub ssl_security_level: Option<String>,
|
||||
pub ssl_ecdh_curve: Option<String>,
|
||||
pub force_ipv4: Option<bool>,
|
||||
pub http2: Option<bool>,
|
||||
pub aiohttp_trust_env: Option<bool>,
|
||||
pub disable_aiohttp_trust_env: Option<bool>,
|
||||
pub disable_aiohttp_transport: Option<bool>,
|
||||
pub user_agent: Option<String>,
|
||||
pub tcp_keepalive: Option<TcpKeepalive>,
|
||||
pub pool_idle_timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
impl HttpSettingsLayer {
|
||||
pub fn from_environment(env: &(dyn Fn(&str) -> Option<String> + 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::<u32>().ok());
|
||||
let seconds = |name: &str, default: u32| {
|
||||
Duration::from_secs(u64::from(number(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"),
|
||||
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 {
|
||||
idle: seconds("AIOHTTP_TCP_KEEPIDLE", 60),
|
||||
interval: seconds("AIOHTTP_TCP_KEEPINTVL", 30),
|
||||
retries: number("AIOHTTP_TCP_KEEPCNT").unwrap_or(5),
|
||||
}),
|
||||
pool_idle_timeout: number("AIOHTTP_KEEPALIVE_TIMEOUT")
|
||||
.map(|timeout| Duration::from_secs(u64::from(timeout))),
|
||||
}
|
||||
}
|
||||
|
||||
fn or(self, lower: Self) -> Self {
|
||||
Self {
|
||||
ssl_verify: self.ssl_verify.or(lower.ssl_verify),
|
||||
ssl_cert_file: self.ssl_cert_file.or(lower.ssl_cert_file),
|
||||
ssl_certificate: self.ssl_certificate.or(lower.ssl_certificate),
|
||||
ssl_security_level: self.ssl_security_level.or(lower.ssl_security_level),
|
||||
ssl_ecdh_curve: self.ssl_ecdh_curve.or(lower.ssl_ecdh_curve),
|
||||
force_ipv4: self.force_ipv4.or(lower.force_ipv4),
|
||||
http2: self.http2.or(lower.http2),
|
||||
aiohttp_trust_env: self.aiohttp_trust_env.or(lower.aiohttp_trust_env),
|
||||
disable_aiohttp_trust_env: self
|
||||
.disable_aiohttp_trust_env
|
||||
.or(lower.disable_aiohttp_trust_env),
|
||||
disable_aiohttp_transport: self
|
||||
.disable_aiohttp_transport
|
||||
.or(lower.disable_aiohttp_transport),
|
||||
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),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct HttpSettings {
|
||||
pub ssl_verify: Option<SslVerify>,
|
||||
pub ssl_cert_file: Option<PathBuf>,
|
||||
pub ssl_certificate: Option<PathBuf>,
|
||||
pub ssl_security_level: Option<String>,
|
||||
pub ssl_ecdh_curve: Option<String>,
|
||||
pub force_ipv4: bool,
|
||||
pub http2: bool,
|
||||
pub user_agent: Option<String>,
|
||||
pub trust_proxy_env: bool,
|
||||
pub connect_timeout: Duration,
|
||||
pub tcp_keepalive: Option<TcpKeepalive>,
|
||||
pub pool_idle_timeout: Duration,
|
||||
}
|
||||
|
||||
impl Default for HttpSettings {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
ssl_verify: None,
|
||||
ssl_cert_file: None,
|
||||
ssl_certificate: None,
|
||||
ssl_security_level: None,
|
||||
ssl_ecdh_curve: None,
|
||||
force_ipv4: false,
|
||||
http2: false,
|
||||
user_agent: None,
|
||||
trust_proxy_env: true,
|
||||
connect_timeout: Duration::from_secs(10),
|
||||
tcp_keepalive: None,
|
||||
pool_idle_timeout: Duration::from_secs(120),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl HttpSettings {
|
||||
pub fn from_layers(
|
||||
highest_precedence_first: impl IntoIterator<Item = HttpSettingsLayer>,
|
||||
) -> Self {
|
||||
let merged = highest_precedence_first
|
||||
.into_iter()
|
||||
.reduce(HttpSettingsLayer::or)
|
||||
.unwrap_or_default();
|
||||
let defaults = Self::default();
|
||||
let http2 = merged.http2.unwrap_or(defaults.http2);
|
||||
Self {
|
||||
ssl_verify: merged.ssl_verify,
|
||||
ssl_cert_file: merged.ssl_cert_file,
|
||||
ssl_certificate: merged
|
||||
.ssl_certificate
|
||||
.filter(|path| !path.as_os_str().is_empty()),
|
||||
ssl_security_level: merged.ssl_security_level.filter(|level| !level.is_empty()),
|
||||
ssl_ecdh_curve: merged.ssl_ecdh_curve.filter(|curve| !curve.is_empty()),
|
||||
force_ipv4: merged.force_ipv4.unwrap_or(defaults.force_ipv4),
|
||||
http2,
|
||||
user_agent: merged.user_agent,
|
||||
trust_proxy_env: !merged.disable_aiohttp_trust_env.unwrap_or(false)
|
||||
|| merged.aiohttp_trust_env.unwrap_or(false)
|
||||
|| merged.disable_aiohttp_transport.unwrap_or(false)
|
||||
|| http2,
|
||||
tcp_keepalive: merged.tcp_keepalive,
|
||||
pool_idle_timeout: merged
|
||||
.pool_idle_timeout
|
||||
.unwrap_or(defaults.pool_idle_timeout),
|
||||
..defaults
|
||||
}
|
||||
}
|
||||
|
||||
pub fn without_missing_files(self, exists: &dyn Fn(&Path) -> bool) -> Self {
|
||||
Self {
|
||||
ssl_verify: match self.ssl_verify {
|
||||
Some(SslVerify::CaBundle(path)) if !exists(&path) => Some(SslVerify::Enabled),
|
||||
other => other,
|
||||
},
|
||||
ssl_cert_file: self.ssl_cert_file.filter(|path| exists(path)),
|
||||
..self
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn no_env(_: &str) -> Option<String> {
|
||||
None
|
||||
}
|
||||
|
||||
fn env_of(
|
||||
values: &'static [(&'static str, &'static str)],
|
||||
) -> impl Fn(&str) -> Option<String> + Sync {
|
||||
move |name| {
|
||||
values
|
||||
.iter()
|
||||
.find(|(key, _)| *key == name)
|
||||
.map(|(_, value)| value.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case("true", SslVerify::Enabled)]
|
||||
#[case(" True ", SslVerify::Enabled)]
|
||||
#[case("FALSE", SslVerify::Disabled)]
|
||||
#[case("/etc/ssl/bundle.pem", SslVerify::CaBundle("/etc/ssl/bundle.pem".into()))]
|
||||
fn ssl_verify_parses_bools_and_treats_anything_else_as_a_bundle_path(
|
||||
#[case] value: &str,
|
||||
#[case] expected: SslVerify,
|
||||
) {
|
||||
assert_eq!(SslVerify::parse(value), expected);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn higher_layers_override_lower_ones() {
|
||||
let configured = HttpSettingsLayer {
|
||||
ssl_verify: Some(SslVerify::Enabled),
|
||||
ssl_certificate: Some("/configured/client.pem".into()),
|
||||
ssl_security_level: Some("configured".into()),
|
||||
user_agent: Some("configured/1".into()),
|
||||
..HttpSettingsLayer::default()
|
||||
};
|
||||
let environment = HttpSettingsLayer::from_environment(&env_of(&[
|
||||
("SSL_VERIFY", "false"),
|
||||
("SSL_CERT_FILE", "/env/roots.pem"),
|
||||
("SSL_CERTIFICATE", "/env/client.pem"),
|
||||
("SSL_SECURITY_LEVEL", "DEFAULT@SECLEVEL=1"),
|
||||
("SSL_ECDH_CURVE", "X25519"),
|
||||
("LITELLM_USER_AGENT", "env/2"),
|
||||
]));
|
||||
let settings = HttpSettings::from_layers([environment, configured]);
|
||||
assert_eq!(settings.ssl_verify, Some(SslVerify::Disabled));
|
||||
assert_eq!(settings.ssl_cert_file, Some("/env/roots.pem".into()));
|
||||
assert_eq!(settings.ssl_certificate, Some("/env/client.pem".into()));
|
||||
assert_eq!(
|
||||
settings.ssl_security_level.as_deref(),
|
||||
Some("DEFAULT@SECLEVEL=1")
|
||||
);
|
||||
assert_eq!(settings.ssl_ecdh_curve.as_deref(), Some("X25519"));
|
||||
assert_eq!(settings.user_agent.as_deref(), Some("env/2"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_explicit_false_in_a_higher_layer_beats_a_lower_true() {
|
||||
let higher = HttpSettingsLayer {
|
||||
http2: Some(false),
|
||||
force_ipv4: Some(false),
|
||||
..HttpSettingsLayer::default()
|
||||
};
|
||||
let lower = HttpSettingsLayer {
|
||||
http2: Some(true),
|
||||
force_ipv4: Some(true),
|
||||
..HttpSettingsLayer::default()
|
||||
};
|
||||
let settings = HttpSettings::from_layers([higher, lower]);
|
||||
assert!(!settings.http2);
|
||||
assert!(!settings.force_ipv4);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_empty_environment_is_an_empty_layer_so_lower_layers_and_defaults_apply() {
|
||||
assert_eq!(
|
||||
HttpSettingsLayer::from_environment(&no_env),
|
||||
HttpSettingsLayer::default()
|
||||
);
|
||||
let configured = HttpSettingsLayer {
|
||||
ssl_verify: Some(SslVerify::CaBundle("/configured/roots.pem".into())),
|
||||
http2: Some(true),
|
||||
user_agent: Some("configured/1".into()),
|
||||
..HttpSettingsLayer::default()
|
||||
};
|
||||
assert_eq!(
|
||||
HttpSettings::from_layers([HttpSettingsLayer::default(), configured]),
|
||||
HttpSettings {
|
||||
ssl_verify: Some(SslVerify::CaBundle("/configured/roots.pem".into())),
|
||||
http2: true,
|
||||
user_agent: Some("configured/1".into()),
|
||||
..HttpSettings::default()
|
||||
}
|
||||
);
|
||||
assert_eq!(HttpSettings::from_layers([]), HttpSettings::default());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_environment_values_clear_the_setting_like_python_truthiness() {
|
||||
let configured = HttpSettingsLayer {
|
||||
ssl_certificate: Some("/configured/client.pem".into()),
|
||||
ssl_security_level: Some("configured".into()),
|
||||
ssl_ecdh_curve: Some("X25519".into()),
|
||||
..HttpSettingsLayer::default()
|
||||
};
|
||||
let environment = HttpSettingsLayer::from_environment(&env_of(&[
|
||||
("SSL_CERTIFICATE", ""),
|
||||
("SSL_SECURITY_LEVEL", ""),
|
||||
("SSL_ECDH_CURVE", ""),
|
||||
]));
|
||||
let settings = HttpSettings::from_layers([environment, configured]);
|
||||
assert_eq!(settings.ssl_certificate, None);
|
||||
assert_eq!(settings.ssl_security_level, None);
|
||||
assert_eq!(settings.ssl_ecdh_curve, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn socket_keepalive_follows_the_aiohttp_variables_with_python_defaults() {
|
||||
let tuned = HttpSettings::from_layers([HttpSettingsLayer::from_environment(&env_of(&[
|
||||
("AIOHTTP_SO_KEEPALIVE", "True"),
|
||||
("AIOHTTP_TCP_KEEPIDLE", "45"),
|
||||
("AIOHTTP_KEEPALIVE_TIMEOUT", "30"),
|
||||
]))]);
|
||||
assert_eq!(
|
||||
tuned.tcp_keepalive,
|
||||
Some(TcpKeepalive {
|
||||
idle: Duration::from_secs(45),
|
||||
interval: Duration::from_secs(30),
|
||||
retries: 5,
|
||||
})
|
||||
);
|
||||
assert_eq!(tuned.pool_idle_timeout, Duration::from_secs(30));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn socket_keepalive_stays_off_unless_enabled() {
|
||||
let settings = HttpSettings::from_layers([HttpSettingsLayer::from_environment(&env_of(
|
||||
&[("AIOHTTP_TCP_KEEPIDLE", "45")],
|
||||
))]);
|
||||
assert_eq!(settings.tcp_keepalive, None);
|
||||
assert_eq!(settings.pool_idle_timeout, Duration::from_secs(120));
|
||||
}
|
||||
|
||||
fn proxy_flags(
|
||||
aiohttp_trust_env: bool,
|
||||
disable_aiohttp_trust_env: bool,
|
||||
disable_aiohttp_transport: bool,
|
||||
http2: bool,
|
||||
) -> HttpSettingsLayer {
|
||||
HttpSettingsLayer {
|
||||
aiohttp_trust_env: Some(aiohttp_trust_env),
|
||||
disable_aiohttp_trust_env: Some(disable_aiohttp_trust_env),
|
||||
disable_aiohttp_transport: Some(disable_aiohttp_transport),
|
||||
http2: Some(http2),
|
||||
..HttpSettingsLayer::default()
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::aiohttp_default(proxy_flags(false, false, false, false), true)]
|
||||
#[case::aiohttp_opted_out(proxy_flags(false, true, false, false), false)]
|
||||
#[case::session_trust_env_beats_opt_out(proxy_flags(true, true, false, false), true)]
|
||||
#[case::http2_uses_httpx(proxy_flags(false, true, false, true), true)]
|
||||
#[case::aiohttp_disabled(proxy_flags(false, true, true, false), true)]
|
||||
fn environment_proxies_apply_unless_the_aiohttp_transport_opts_out(
|
||||
#[case] layer: HttpSettingsLayer,
|
||||
#[case] expected: bool,
|
||||
) {
|
||||
assert_eq!(HttpSettings::from_layers([layer]).trust_proxy_env, expected);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_proxy_opt_out_in_one_source_still_yields_to_trust_env_from_another() {
|
||||
let environment =
|
||||
HttpSettingsLayer::from_environment(&env_of(&[("DISABLE_AIOHTTP_TRUST_ENV", "true")]));
|
||||
let configured = HttpSettingsLayer {
|
||||
aiohttp_trust_env: Some(true),
|
||||
..HttpSettingsLayer::default()
|
||||
};
|
||||
assert!(!HttpSettings::from_layers([environment.clone()]).trust_proxy_env);
|
||||
assert!(HttpSettings::from_layers([environment, configured]).trust_proxy_env);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_files_fall_back_to_default_verification() {
|
||||
let settings = HttpSettings {
|
||||
ssl_verify: Some(SslVerify::CaBundle("/absent/roots.pem".into())),
|
||||
ssl_cert_file: Some("/absent/env.pem".into()),
|
||||
..HttpSettings::default()
|
||||
}
|
||||
.without_missing_files(&|_| false);
|
||||
assert_eq!(settings.ssl_verify, Some(SslVerify::Enabled));
|
||||
assert_eq!(settings.ssl_cert_file, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn existing_files_are_kept() {
|
||||
let settings = HttpSettings {
|
||||
ssl_verify: Some(SslVerify::CaBundle("/present/roots.pem".into())),
|
||||
ssl_cert_file: Some("/present/env.pem".into()),
|
||||
..HttpSettings::default()
|
||||
};
|
||||
assert_eq!(settings.clone().without_missing_files(&|_| true), settings);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case("true", Some(true))]
|
||||
#[case("True", Some(true))]
|
||||
#[case("false", None)]
|
||||
#[case("1", None)]
|
||||
fn boolean_switches_only_turn_on_for_true(
|
||||
#[case] value: &'static str,
|
||||
#[case] expected: Option<bool>,
|
||||
) {
|
||||
let env = move |name: &str| match name {
|
||||
"LITELLM_HTTP2"
|
||||
| "AIOHTTP_TRUST_ENV"
|
||||
| "DISABLE_AIOHTTP_TRANSPORT"
|
||||
| "DISABLE_AIOHTTP_TRUST_ENV" => Some(value.to_string()),
|
||||
_ => None,
|
||||
};
|
||||
let layer = HttpSettingsLayer::from_environment(&env);
|
||||
assert_eq!(layer.http2, expected);
|
||||
assert_eq!(layer.aiohttp_trust_env, expected);
|
||||
assert_eq!(layer.disable_aiohttp_transport, expected);
|
||||
assert_eq!(layer.disable_aiohttp_trust_env, expected);
|
||||
}
|
||||
}
|
||||
411
litellm-rust/crates/http/src/tls.rs
Normal file
411
litellm-rust/crates/http/src/tls.rs
Normal file
|
|
@ -0,0 +1,411 @@
|
|||
use std::{fmt, path::Path, str::FromStr, sync::Arc};
|
||||
|
||||
use rustls::{
|
||||
CipherSuite, ClientConfig, DigitallySignedStruct, RootCertStore, SignatureScheme,
|
||||
client::danger::{HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier},
|
||||
crypto::{CryptoProvider, SupportedKxGroup, ring},
|
||||
pki_types::{CertificateDer, PrivateKeyDer, ServerName, UnixTime, pem::PemObject},
|
||||
};
|
||||
|
||||
use crate::{
|
||||
config::{HttpClientConfig, Verify},
|
||||
error::Error,
|
||||
};
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
|
||||
pub enum KeyExchangeGroup {
|
||||
X25519,
|
||||
Secp256r1,
|
||||
Secp384r1,
|
||||
}
|
||||
|
||||
impl FromStr for KeyExchangeGroup {
|
||||
type Err = Unsupported;
|
||||
|
||||
fn from_str(name: &str) -> Result<Self, Self::Err> {
|
||||
match name.trim().to_ascii_lowercase().as_str() {
|
||||
"x25519" => Ok(Self::X25519),
|
||||
"prime256v1" | "secp256r1" | "p-256" => Ok(Self::Secp256r1),
|
||||
"secp384r1" | "p-384" => Ok(Self::Secp384r1),
|
||||
_ => Err(Unsupported::EcdhCurve(name.to_owned())),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl KeyExchangeGroup {
|
||||
fn supported(self) -> &'static dyn SupportedKxGroup {
|
||||
match self {
|
||||
Self::X25519 => ring::kx_group::X25519,
|
||||
Self::Secp256r1 => ring::kx_group::SECP256R1,
|
||||
Self::Secp384r1 => ring::kx_group::SECP384R1,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
|
||||
pub enum Tls12CipherSuite {
|
||||
EcdheEcdsaAes128Gcm,
|
||||
EcdheEcdsaAes256Gcm,
|
||||
EcdheEcdsaChacha20,
|
||||
EcdheRsaAes128Gcm,
|
||||
EcdheRsaAes256Gcm,
|
||||
EcdheRsaChacha20,
|
||||
}
|
||||
|
||||
impl FromStr for Tls12CipherSuite {
|
||||
type Err = Unsupported;
|
||||
|
||||
fn from_str(name: &str) -> Result<Self, Self::Err> {
|
||||
match name {
|
||||
"ECDHE-ECDSA-AES128-GCM-SHA256" => Ok(Self::EcdheEcdsaAes128Gcm),
|
||||
"ECDHE-ECDSA-AES256-GCM-SHA384" => Ok(Self::EcdheEcdsaAes256Gcm),
|
||||
"ECDHE-ECDSA-CHACHA20-POLY1305" => Ok(Self::EcdheEcdsaChacha20),
|
||||
"ECDHE-RSA-AES128-GCM-SHA256" => Ok(Self::EcdheRsaAes128Gcm),
|
||||
"ECDHE-RSA-AES256-GCM-SHA384" => Ok(Self::EcdheRsaAes256Gcm),
|
||||
"ECDHE-RSA-CHACHA20-POLY1305" => Ok(Self::EcdheRsaChacha20),
|
||||
_ => Err(Unsupported::CipherToken(name.to_owned())),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Tls12CipherSuite {
|
||||
fn suite(self) -> CipherSuite {
|
||||
match self {
|
||||
Self::EcdheEcdsaAes128Gcm => CipherSuite::TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,
|
||||
Self::EcdheEcdsaAes256Gcm => CipherSuite::TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384,
|
||||
Self::EcdheEcdsaChacha20 => CipherSuite::TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305_SHA256,
|
||||
Self::EcdheRsaAes128Gcm => CipherSuite::TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256,
|
||||
Self::EcdheRsaAes256Gcm => CipherSuite::TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384,
|
||||
Self::EcdheRsaChacha20 => CipherSuite::TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305_SHA256,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Hash, thiserror::Error)]
|
||||
pub enum Unsupported {
|
||||
#[error(
|
||||
"ssl_ecdh_curve {0:?} is not supported: rustls with ring only offers X25519, prime256v1 and secp384r1, so the default key exchange groups are used"
|
||||
)]
|
||||
EcdhCurve(String),
|
||||
#[error(
|
||||
"ssl_security_level {0:?} is not supported: rustls has one fixed security level, comparable to OpenSSL level 2, so legacy servers that need a lower level cannot be reached"
|
||||
)]
|
||||
SecurityLevel(String),
|
||||
#[error(
|
||||
"ssl_security_level entry {0:?} is not supported: rustls only offers ECDHE AEAD cipher suites, so the entry is ignored"
|
||||
)]
|
||||
CipherToken(String),
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub(crate) struct CipherSelection {
|
||||
pub(crate) tls12_cipher_suites: Option<Vec<Tls12CipherSuite>>,
|
||||
pub(crate) unsupported: Vec<Unsupported>,
|
||||
}
|
||||
|
||||
enum CipherToken {
|
||||
Suite(Tls12CipherSuite),
|
||||
EverySuite,
|
||||
Ordering,
|
||||
Unsupported(Unsupported),
|
||||
}
|
||||
|
||||
impl From<&str> for CipherToken {
|
||||
fn from(token: &str) -> Self {
|
||||
match token {
|
||||
"DEFAULT" | "ALL" | "HIGH" => Self::EverySuite,
|
||||
"@STRENGTH" | "@SECLEVEL=2" => Self::Ordering,
|
||||
level if level.starts_with("@SECLEVEL=") => {
|
||||
Self::Unsupported(Unsupported::SecurityLevel(level.to_owned()))
|
||||
}
|
||||
name => name.parse().map_or_else(Self::Unsupported, Self::Suite),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<&str> for CipherSelection {
|
||||
fn from(value: &str) -> Self {
|
||||
let tokens: Vec<CipherToken> = tokenize(value)
|
||||
.iter()
|
||||
.map(|token| CipherToken::from(token.as_str()))
|
||||
.collect();
|
||||
let every_suite = tokens
|
||||
.iter()
|
||||
.any(|token| matches!(token, CipherToken::EverySuite));
|
||||
let mut suites: Vec<Tls12CipherSuite> = tokens
|
||||
.iter()
|
||||
.filter_map(|token| match token {
|
||||
CipherToken::Suite(suite) => Some(*suite),
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
suites.sort_unstable();
|
||||
suites.dedup();
|
||||
CipherSelection {
|
||||
tls12_cipher_suites: (!every_suite && !suites.is_empty()).then_some(suites),
|
||||
unsupported: tokens
|
||||
.into_iter()
|
||||
.filter_map(|token| match token {
|
||||
CipherToken::Unsupported(unsupported) => Some(unsupported),
|
||||
_ => None,
|
||||
})
|
||||
.collect(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn tokenize(value: &str) -> Vec<String> {
|
||||
value
|
||||
.split([':', ',', ' '])
|
||||
.flat_map(|entry| match entry.split_once('@') {
|
||||
Some((name, command)) => vec![name.to_owned(), format!("@{command}")],
|
||||
None => vec![entry.to_owned()],
|
||||
})
|
||||
.filter(|token| !token.is_empty())
|
||||
.collect()
|
||||
}
|
||||
|
||||
impl TryFrom<&HttpClientConfig> for ClientConfig {
|
||||
type Error = Error;
|
||||
|
||||
fn try_from(config: &HttpClientConfig) -> Result<Self, Self::Error> {
|
||||
let base = ring::default_provider();
|
||||
let provider = Arc::new(CryptoProvider {
|
||||
kx_groups: config
|
||||
.key_exchange_group
|
||||
.map_or_else(|| base.kx_groups.clone(), |group| vec![group.supported()]),
|
||||
cipher_suites: base
|
||||
.cipher_suites
|
||||
.iter()
|
||||
.copied()
|
||||
.filter(|suite| {
|
||||
suite.tls13().is_some()
|
||||
|| config.tls12_cipher_suites.as_ref().is_none_or(|allowed| {
|
||||
allowed.iter().any(|a| a.suite() == suite.suite())
|
||||
})
|
||||
})
|
||||
.collect(),
|
||||
..base
|
||||
});
|
||||
let builder = ClientConfig::builder_with_provider(Arc::clone(&provider))
|
||||
.with_safe_default_protocol_versions()
|
||||
.map_err(|error| Error::Client(error.to_string()))?;
|
||||
let verified = match &config.verify {
|
||||
Verify::Disabled => builder
|
||||
.dangerous()
|
||||
.with_custom_certificate_verifier(Arc::new(NoVerification(provider))),
|
||||
Verify::BuiltInRoots => builder.with_root_certificates(RootCertStore {
|
||||
roots: webpki_roots::TLS_SERVER_ROOTS.to_vec(),
|
||||
}),
|
||||
Verify::CaBundle(path) => builder.with_root_certificates(bundle_roots(path)?),
|
||||
};
|
||||
let mut tls = match &config.client_certificate {
|
||||
None => verified.with_no_client_auth(),
|
||||
Some(path) => {
|
||||
let (chain, key) = identity(path)?;
|
||||
verified
|
||||
.with_client_auth_cert(chain, key)
|
||||
.map_err(|error| invalid_pem(path, error))?
|
||||
}
|
||||
};
|
||||
tls.alpn_protocols = if config.http2 {
|
||||
vec![b"h2".to_vec(), b"http/1.1".to_vec()]
|
||||
} else {
|
||||
vec![b"http/1.1".to_vec()]
|
||||
};
|
||||
Ok(tls)
|
||||
}
|
||||
}
|
||||
|
||||
fn bundle_roots(path: &Path) -> Result<RootCertStore, Error> {
|
||||
let certificates = certificates(path)?;
|
||||
if certificates.is_empty() {
|
||||
return Err(invalid_pem(path, "no certificates found"));
|
||||
}
|
||||
let mut store = RootCertStore::empty();
|
||||
for certificate in certificates {
|
||||
store
|
||||
.add(certificate)
|
||||
.map_err(|error| invalid_pem(path, error))?;
|
||||
}
|
||||
Ok(store)
|
||||
}
|
||||
|
||||
fn identity(path: &Path) -> Result<(Vec<CertificateDer<'static>>, PrivateKeyDer<'static>), Error> {
|
||||
let chain = certificates(path)?;
|
||||
if chain.is_empty() {
|
||||
return Err(invalid_pem(path, "no certificates found"));
|
||||
}
|
||||
let key =
|
||||
PrivateKeyDer::from_pem_slice(&read(path)?).map_err(|error| invalid_pem(path, error))?;
|
||||
Ok((chain, key))
|
||||
}
|
||||
|
||||
fn certificates(path: &Path) -> Result<Vec<CertificateDer<'static>>, Error> {
|
||||
CertificateDer::pem_slice_iter(&read(path)?)
|
||||
.collect::<Result<_, _>>()
|
||||
.map_err(|error| invalid_pem(path, error))
|
||||
}
|
||||
|
||||
fn read(path: &Path) -> Result<Vec<u8>, Error> {
|
||||
std::fs::read(path).map_err(|error| Error::Read {
|
||||
path: path.to_path_buf(),
|
||||
message: error.to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
fn invalid_pem(path: &Path, message: impl fmt::Display) -> Error {
|
||||
Error::InvalidPem {
|
||||
path: path.to_path_buf(),
|
||||
message: message.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct NoVerification(Arc<CryptoProvider>);
|
||||
|
||||
impl ServerCertVerifier for NoVerification {
|
||||
fn verify_server_cert(
|
||||
&self,
|
||||
_end_entity: &CertificateDer<'_>,
|
||||
_intermediates: &[CertificateDer<'_>],
|
||||
_server_name: &ServerName<'_>,
|
||||
_ocsp_response: &[u8],
|
||||
_now: UnixTime,
|
||||
) -> Result<ServerCertVerified, rustls::Error> {
|
||||
Ok(ServerCertVerified::assertion())
|
||||
}
|
||||
|
||||
fn verify_tls12_signature(
|
||||
&self,
|
||||
_message: &[u8],
|
||||
_cert: &CertificateDer<'_>,
|
||||
_dss: &DigitallySignedStruct,
|
||||
) -> Result<HandshakeSignatureValid, rustls::Error> {
|
||||
Ok(HandshakeSignatureValid::assertion())
|
||||
}
|
||||
|
||||
fn verify_tls13_signature(
|
||||
&self,
|
||||
_message: &[u8],
|
||||
_cert: &CertificateDer<'_>,
|
||||
_dss: &DigitallySignedStruct,
|
||||
) -> Result<HandshakeSignatureValid, rustls::Error> {
|
||||
Ok(HandshakeSignatureValid::assertion())
|
||||
}
|
||||
|
||||
fn supported_verify_schemes(&self) -> Vec<SignatureScheme> {
|
||||
self.0.signature_verification_algorithms.supported_schemes()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
use rustls::NamedGroup;
|
||||
|
||||
use super::*;
|
||||
use crate::{HttpSettings, Resolution};
|
||||
|
||||
fn config(settings: HttpSettings) -> HttpClientConfig {
|
||||
Resolution::from(&settings).config
|
||||
}
|
||||
|
||||
fn offered_groups(tls: &ClientConfig) -> Vec<NamedGroup> {
|
||||
tls.crypto_provider()
|
||||
.kx_groups
|
||||
.iter()
|
||||
.map(|group| group.name())
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn offered_tls12_suites(tls: &ClientConfig) -> Vec<CipherSuite> {
|
||||
tls.crypto_provider()
|
||||
.cipher_suites
|
||||
.iter()
|
||||
.filter(|suite| suite.tls13().is_none())
|
||||
.map(|suite| suite.suite())
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case("X25519", NamedGroup::X25519)]
|
||||
#[case("prime256v1", NamedGroup::secp256r1)]
|
||||
#[case("secp384r1", NamedGroup::secp384r1)]
|
||||
fn ecdh_curve_is_the_only_key_exchange_group_offered(
|
||||
#[case] curve: &str,
|
||||
#[case] expected: NamedGroup,
|
||||
) {
|
||||
let tls = ClientConfig::try_from(&config(HttpSettings {
|
||||
ssl_ecdh_curve: Some(curve.into()),
|
||||
..HttpSettings::default()
|
||||
}))
|
||||
.unwrap();
|
||||
assert_eq!(offered_groups(&tls), [expected]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn default_settings_offer_every_group_and_suite_of_the_provider() {
|
||||
let tls = ClientConfig::try_from(&config(HttpSettings::default())).unwrap();
|
||||
let provider = ring::default_provider();
|
||||
assert_eq!(offered_groups(&tls).len(), provider.kx_groups.len());
|
||||
assert_eq!(
|
||||
tls.crypto_provider().cipher_suites.len(),
|
||||
provider.cipher_suites.len()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn named_suites_are_the_only_tls12_suites_offered_and_tls13_stays() {
|
||||
let tls = ClientConfig::try_from(&config(HttpSettings {
|
||||
ssl_security_level: Some("ECDHE-RSA-AES256-GCM-SHA384".into()),
|
||||
..HttpSettings::default()
|
||||
}))
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
offered_tls12_suites(&tls),
|
||||
[CipherSuite::TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384]
|
||||
);
|
||||
assert!(
|
||||
tls.crypto_provider()
|
||||
.cipher_suites
|
||||
.iter()
|
||||
.any(|suite| suite.tls13().is_some())
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case(true, &[b"h2".as_slice(), b"http/1.1".as_slice()])]
|
||||
#[case(false, &[b"http/1.1".as_slice()])]
|
||||
fn alpn_offers_h2_only_when_http2_is_on(#[case] http2: bool, #[case] expected: &[&[u8]]) {
|
||||
let tls = ClientConfig::try_from(&config(HttpSettings {
|
||||
http2,
|
||||
..HttpSettings::default()
|
||||
}))
|
||||
.unwrap();
|
||||
assert_eq!(tls.alpn_protocols, expected);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn client_certificate_without_a_private_key_is_an_invalid_pem_error() {
|
||||
let path = std::env::temp_dir().join(format!(
|
||||
"litellm-http-cert-without-key-{}.pem",
|
||||
std::process::id()
|
||||
));
|
||||
std::fs::write(
|
||||
&path,
|
||||
b"-----BEGIN CERTIFICATE-----\nAA==\n-----END CERTIFICATE-----\n",
|
||||
)
|
||||
.unwrap();
|
||||
let result = ClientConfig::try_from(&HttpClientConfig {
|
||||
client_certificate: Some(path.clone()),
|
||||
..config(HttpSettings::default())
|
||||
})
|
||||
.map(drop);
|
||||
std::fs::remove_file(&path).unwrap();
|
||||
assert!(matches!(
|
||||
result,
|
||||
Err(Error::InvalidPem { path: reported, .. }) if reported == path
|
||||
));
|
||||
}
|
||||
}
|
||||
|
|
@ -17,6 +17,7 @@ litellm-auth-azure.workspace = true
|
|||
litellm-auth-gcp.workspace = true
|
||||
litellm-host.workspace = true
|
||||
litellm-framing.workspace = true
|
||||
litellm-http.workspace = true
|
||||
base64.workspace = true
|
||||
bytes.workspace = true
|
||||
data-url = "0.3.2"
|
||||
|
|
|
|||
|
|
@ -21,7 +21,6 @@ use crate::{
|
|||
|
||||
pub const OCR_RESPONSE_MAX_BYTES: usize = 64 * 1024 * 1024;
|
||||
pub const OCR_HTTP_TIMEOUT_SECS: u64 = 600;
|
||||
pub const OCR_CONNECT_TIMEOUT_SECS: u64 = 10;
|
||||
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;
|
||||
|
|
|
|||
|
|
@ -1,9 +1,8 @@
|
|||
use std::{sync::OnceLock, time::Duration};
|
||||
|
||||
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 serde::{Serialize, de::DeserializeOwned};
|
||||
use serde_json::Value;
|
||||
|
||||
|
|
@ -11,14 +10,13 @@ use crate::{
|
|||
base_llm::ocr::{
|
||||
error::Error,
|
||||
transformation::{
|
||||
BaseOcrConfig, DecodedOcrResponse, LiteLLMOcrResponse, OCR_CONNECT_TIMEOUT_SECS,
|
||||
OcrDocument, OcrResponseContext, PreparedOcrRequest, decode_request_value,
|
||||
decode_response,
|
||||
BaseOcrConfig, DecodedOcrResponse, LiteLLMOcrResponse, OcrDocument, OcrResponseContext,
|
||||
PreparedOcrRequest, decode_request_value, decode_response,
|
||||
},
|
||||
},
|
||||
custom_httpx::{
|
||||
http_handler::{HeaderPolicy, execute_http_request, with_headers},
|
||||
media::MediaFetcher,
|
||||
media::{MediaFetcher, UrlPolicy},
|
||||
transport,
|
||||
},
|
||||
};
|
||||
|
|
@ -40,30 +38,20 @@ pub struct OcrClient {
|
|||
}
|
||||
|
||||
impl OcrClient {
|
||||
pub fn new(provider_http: reqwest::Client) -> Result<Self, transport::Error> {
|
||||
let document_fetcher = MediaFetcher::new().map_err(transport::Error::from)?;
|
||||
pub fn new(
|
||||
pool: &HttpClientPool,
|
||||
config: &HttpClientConfig,
|
||||
url_policy: UrlPolicy,
|
||||
vertex_auth: VertexAuth,
|
||||
) -> Result<Self, litellm_http::Error> {
|
||||
Ok(Self {
|
||||
provider_http,
|
||||
polling_http: no_redirect_http()?,
|
||||
document_fetcher,
|
||||
vertex_auth: VertexAuth::default(),
|
||||
provider_http: pool.client(config, ClientVariant::Provider)?,
|
||||
polling_http: pool.client(config, ClientVariant::NoRedirect)?,
|
||||
document_fetcher: MediaFetcher::new(pool, config, url_policy)?,
|
||||
vertex_auth,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn shared() -> Result<Self, Error> {
|
||||
static CLIENT: OnceLock<Result<OcrClient, transport::Error>> = OnceLock::new();
|
||||
let client = CLIENT
|
||||
.get_or_init(|| {
|
||||
reqwest::Client::builder()
|
||||
.connect_timeout(Duration::from_secs(OCR_CONNECT_TIMEOUT_SECS))
|
||||
.build()
|
||||
.map_err(transport::Error::from)
|
||||
.and_then(OcrClient::new)
|
||||
})
|
||||
.clone()?;
|
||||
Ok(client)
|
||||
}
|
||||
|
||||
pub fn provider_http(&self) -> &reqwest::Client {
|
||||
&self.provider_http
|
||||
}
|
||||
|
|
@ -84,21 +72,16 @@ impl OcrClient {
|
|||
pub fn for_test(provider_http: reqwest::Client, document_http: reqwest::Client) -> Self {
|
||||
Self {
|
||||
provider_http,
|
||||
polling_http: no_redirect_http().expect("test polling client builds"),
|
||||
polling_http: reqwest::Client::builder()
|
||||
.redirect(reqwest::redirect::Policy::none())
|
||||
.build()
|
||||
.expect("test polling client builds"),
|
||||
document_fetcher: MediaFetcher::for_test(document_http),
|
||||
vertex_auth: VertexAuth::default(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn no_redirect_http() -> Result<reqwest::Client, transport::Error> {
|
||||
reqwest::Client::builder()
|
||||
.connect_timeout(Duration::from_secs(OCR_CONNECT_TIMEOUT_SECS))
|
||||
.redirect(reqwest::redirect::Policy::none())
|
||||
.build()
|
||||
.map_err(transport::Error::from)
|
||||
}
|
||||
|
||||
/// Rust counterpart of `BaseLLMHTTPHandler.async_ocr`: prepare the provider request,
|
||||
/// send it, and hand the response to the config for normalization.
|
||||
pub async fn ocr<C: BaseOcrConfig>(
|
||||
|
|
@ -296,6 +279,8 @@ pub fn body_document(body: &Value) -> Result<OcrDocument, Error> {
|
|||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::time::Duration;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
|
|
|
|||
|
|
@ -7,13 +7,12 @@ use std::{
|
|||
time::Duration,
|
||||
};
|
||||
|
||||
use litellm_http::{ClientVariant, EnvironmentProxies, HttpClientConfig, HttpClientPool};
|
||||
use reqwest::{
|
||||
Url,
|
||||
dns::{Addrs, Name, Resolve, Resolving},
|
||||
};
|
||||
|
||||
const MEDIA_CONNECT_TIMEOUT_SECS: u64 = 10;
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum Error {
|
||||
#[error("media URL rejected by network policy")]
|
||||
|
|
@ -36,10 +35,45 @@ pub enum Error {
|
|||
Transport(#[from] crate::custom_httpx::transport::Error),
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct UrlPolicy {
|
||||
pub validate: bool,
|
||||
pub allowed_hosts: Vec<String>,
|
||||
}
|
||||
|
||||
impl Default for UrlPolicy {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
validate: true,
|
||||
allowed_hosts: Vec::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl UrlPolicy {
|
||||
fn allows(&self, host: &str, port: u16) -> bool {
|
||||
let host = normalize_host(host);
|
||||
let with_port = format!("{host}:{port}");
|
||||
self.allowed_hosts
|
||||
.iter()
|
||||
.map(|entry| normalize_host(entry))
|
||||
.any(|entry| entry == host || entry == with_port)
|
||||
}
|
||||
}
|
||||
|
||||
fn normalize_host(host: &str) -> String {
|
||||
host.to_ascii_lowercase().trim_end_matches('.').to_owned()
|
||||
}
|
||||
|
||||
type ProxyMatch = Arc<dyn Fn(&Url) -> bool + Send + Sync>;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct MediaFetcher {
|
||||
client: reqwest::Client,
|
||||
pinned: reqwest::Client,
|
||||
unpinned: reqwest::Client,
|
||||
uses_proxy: ProxyMatch,
|
||||
address_resolver: Arc<dyn AddressResolver>,
|
||||
url_policy: UrlPolicy,
|
||||
allow_private_network: bool,
|
||||
}
|
||||
|
||||
|
|
@ -63,26 +97,39 @@ pub struct DownloadedMedia {
|
|||
}
|
||||
|
||||
impl MediaFetcher {
|
||||
pub fn new() -> Result<Self, reqwest::Error> {
|
||||
Self::with_resolvers(Arc::new(PublicDnsResolver), Arc::new(SystemAddressResolver))
|
||||
pub fn new(
|
||||
pool: &HttpClientPool,
|
||||
config: &HttpClientConfig,
|
||||
url_policy: UrlPolicy,
|
||||
) -> Result<Self, litellm_http::Error> {
|
||||
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)
|
||||
};
|
||||
Self::with_resolution(
|
||||
pool,
|
||||
config,
|
||||
url_policy,
|
||||
Arc::new(SystemAddressResolver),
|
||||
uses_proxy,
|
||||
)
|
||||
}
|
||||
|
||||
fn with_resolvers<R>(
|
||||
transport_resolver: Arc<R>,
|
||||
fn with_resolution(
|
||||
pool: &HttpClientPool,
|
||||
config: &HttpClientConfig,
|
||||
url_policy: UrlPolicy,
|
||||
address_resolver: Arc<dyn AddressResolver>,
|
||||
) -> Result<Self, reqwest::Error>
|
||||
where
|
||||
R: Resolve + 'static,
|
||||
{
|
||||
let client = reqwest::Client::builder()
|
||||
.connect_timeout(Duration::from_secs(MEDIA_CONNECT_TIMEOUT_SECS))
|
||||
.redirect(reqwest::redirect::Policy::none())
|
||||
.no_proxy()
|
||||
.dns_resolver(transport_resolver)
|
||||
.build()?;
|
||||
uses_proxy: ProxyMatch,
|
||||
) -> Result<Self, litellm_http::Error> {
|
||||
Ok(Self {
|
||||
client,
|
||||
pinned: pool.client(config, ClientVariant::Media)?,
|
||||
unpinned: pool.client(config, ClientVariant::UnpinnedMedia)?,
|
||||
uses_proxy,
|
||||
address_resolver,
|
||||
url_policy,
|
||||
allow_private_network: false,
|
||||
})
|
||||
}
|
||||
|
|
@ -90,8 +137,11 @@ impl MediaFetcher {
|
|||
#[cfg(any(test, feature = "test-support"))]
|
||||
pub fn for_test(client: reqwest::Client) -> Self {
|
||||
Self {
|
||||
client,
|
||||
pinned: client.clone(),
|
||||
unpinned: client,
|
||||
uses_proxy: Arc::new(|_| false),
|
||||
address_resolver: Arc::new(AllowPrivateResolver),
|
||||
url_policy: UrlPolicy::default(),
|
||||
allow_private_network: true,
|
||||
}
|
||||
}
|
||||
|
|
@ -112,9 +162,9 @@ impl MediaFetcher {
|
|||
) -> Result<DownloadedMedia, Error> {
|
||||
let mut redirects_followed = 0;
|
||||
loop {
|
||||
self.validate_url(&url).await?;
|
||||
let mut response = self
|
||||
.client
|
||||
.client_for(&url)
|
||||
.await?
|
||||
.get(url.clone())
|
||||
.send()
|
||||
.await
|
||||
|
|
@ -161,7 +211,10 @@ impl MediaFetcher {
|
|||
}
|
||||
}
|
||||
|
||||
async fn validate_url(&self, url: &Url) -> Result<(), Error> {
|
||||
async fn client_for(&self, url: &Url) -> Result<&reqwest::Client, Error> {
|
||||
if !self.url_policy.validate {
|
||||
return Ok(&self.unpinned);
|
||||
}
|
||||
if !matches!(url.scheme(), "http" | "https")
|
||||
|| !url.username().is_empty()
|
||||
|| url.password().is_some()
|
||||
|
|
@ -170,12 +223,28 @@ impl MediaFetcher {
|
|||
}
|
||||
let host = url.host_str().ok_or(Error::BlockedUrl)?;
|
||||
if self.allow_private_network {
|
||||
return Ok(());
|
||||
}
|
||||
if let Ok(ip) = host.parse::<IpAddr>() {
|
||||
return (!is_blocked_ip(ip)).then_some(()).ok_or(Error::BlockedUrl);
|
||||
return Ok(&self.pinned);
|
||||
}
|
||||
let port = url.port_or_known_default().ok_or(Error::BlockedUrl)?;
|
||||
if self.url_policy.allows(host, port) {
|
||||
return Ok(&self.unpinned);
|
||||
}
|
||||
self.validate_host(host, port).await?;
|
||||
Ok(if (self.uses_proxy)(url) {
|
||||
&self.unpinned
|
||||
} else {
|
||||
&self.pinned
|
||||
})
|
||||
}
|
||||
|
||||
async fn validate_host(&self, host: &str, port: u16) -> Result<(), Error> {
|
||||
if let Ok(ip) = host
|
||||
.trim_start_matches('[')
|
||||
.trim_end_matches(']')
|
||||
.parse::<IpAddr>()
|
||||
{
|
||||
return (!is_blocked_ip(ip)).then_some(()).ok_or(Error::BlockedUrl);
|
||||
}
|
||||
let addresses = self
|
||||
.address_resolver
|
||||
.resolve(host, port)
|
||||
|
|
@ -236,7 +305,7 @@ fn is_blocked_ip(ip: IpAddr) -> bool {
|
|||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct PublicDnsResolver;
|
||||
pub struct PublicDnsResolver;
|
||||
|
||||
struct SystemAddressResolver;
|
||||
|
||||
|
|
@ -281,6 +350,7 @@ impl Resolve for PublicDnsResolver {
|
|||
mod tests {
|
||||
use std::collections::HashSet;
|
||||
|
||||
use litellm_http::{HttpSettings, Resolution};
|
||||
use tokio::{
|
||||
io::{AsyncReadExt, AsyncWriteExt},
|
||||
net::TcpListener,
|
||||
|
|
@ -364,13 +434,35 @@ mod tests {
|
|||
address: SocketAddr,
|
||||
blocked_hosts: HashSet<&'static str>,
|
||||
) -> MediaFetcher {
|
||||
MediaFetcher::with_resolvers(
|
||||
Arc::new(LoopbackDnsResolver(address)),
|
||||
fetcher(address, blocked_hosts, UrlPolicy::default(), false)
|
||||
}
|
||||
|
||||
fn fetcher(
|
||||
pinned_address: SocketAddr,
|
||||
blocked_hosts: HashSet<&'static str>,
|
||||
url_policy: UrlPolicy,
|
||||
uses_proxy: bool,
|
||||
) -> MediaFetcher {
|
||||
let direct = HttpClientConfig {
|
||||
trust_proxy_env: false,
|
||||
..Resolution::from(&HttpSettings::default()).config
|
||||
};
|
||||
MediaFetcher::with_resolution(
|
||||
&HttpClientPool::new(Arc::new(LoopbackDnsResolver(pinned_address))),
|
||||
&direct,
|
||||
url_policy,
|
||||
Arc::new(TestAddressResolver { blocked_hosts }),
|
||||
Arc::new(move |_| uses_proxy),
|
||||
)
|
||||
.expect("test fetcher builds")
|
||||
}
|
||||
|
||||
const UNROUTABLE: SocketAddr =
|
||||
SocketAddr::new(IpAddr::V4(std::net::Ipv4Addr::new(192, 0, 2, 1)), 9);
|
||||
|
||||
const OK_RESPONSE: &[u8] =
|
||||
b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok";
|
||||
|
||||
fn policy(max_bytes: u64, max_redirects: usize) -> DownloadPolicy {
|
||||
DownloadPolicy {
|
||||
timeout: Duration::from_secs(1),
|
||||
|
|
@ -542,12 +634,90 @@ mod tests {
|
|||
|
||||
#[tokio::test]
|
||||
async fn rejects_url_credentials_before_network_access() {
|
||||
let fetcher = MediaFetcher::new().expect("media fetcher builds");
|
||||
let fetcher = MediaFetcher::new(
|
||||
&HttpClientPool::new(Arc::new(PublicDnsResolver)),
|
||||
&Resolution::from(&HttpSettings::default()).config,
|
||||
UrlPolicy::default(),
|
||||
)
|
||||
.expect("media fetcher builds");
|
||||
let url =
|
||||
Url::parse("https://user:password@8.8.8.8/document").expect("credentialed URL parses");
|
||||
assert!(matches!(
|
||||
fetcher.validate_url(&url).await,
|
||||
fetcher.fetch(url, policy(1, 0)).await,
|
||||
Err(Error::BlockedUrl)
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn allowlisted_private_host_is_fetched_without_the_pinned_resolver() {
|
||||
let (url, server, _) = serve_named("localhost", vec![OK_RESPONSE]).await;
|
||||
let port = url.port().expect("test URL has a port");
|
||||
let allowed = UrlPolicy {
|
||||
validate: true,
|
||||
allowed_hosts: vec![format!("LOCALHOST:{port}")],
|
||||
};
|
||||
let media = fetcher(UNROUTABLE, HashSet::from(["localhost"]), allowed, false)
|
||||
.fetch(url, policy(2, 0))
|
||||
.await
|
||||
.expect("allowlisted host downloads");
|
||||
server.await.expect("server completes");
|
||||
assert_eq!(media.bytes, b"ok");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn allowlist_entry_for_another_port_does_not_open_the_host() {
|
||||
let (url, _server, _) = serve_named("localhost", vec![OK_RESPONSE]).await;
|
||||
let other_port = UrlPolicy {
|
||||
validate: true,
|
||||
allowed_hosts: vec!["localhost:1".into()],
|
||||
};
|
||||
let result = fetcher(UNROUTABLE, HashSet::from(["localhost"]), other_port, false)
|
||||
.fetch(url, policy(2, 0))
|
||||
.await;
|
||||
assert!(matches!(result, Err(Error::BlockedUrl)));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn validation_off_fetches_private_hosts_and_follows_redirects() {
|
||||
let (url, server, _) = serve_named(
|
||||
"localhost",
|
||||
vec![
|
||||
b"HTTP/1.1 302 Found\r\nLocation: /moved\r\nContent-Length: 0\r\nConnection: close\r\n\r\n",
|
||||
OK_RESPONSE,
|
||||
],
|
||||
)
|
||||
.await;
|
||||
let off = UrlPolicy {
|
||||
validate: false,
|
||||
allowed_hosts: Vec::new(),
|
||||
};
|
||||
let media = fetcher(UNROUTABLE, HashSet::from(["localhost"]), off, false)
|
||||
.fetch(url, policy(2, 1))
|
||||
.await
|
||||
.expect("unvalidated download succeeds");
|
||||
let requests = server.await.expect("server completes");
|
||||
assert_eq!(media.bytes, b"ok");
|
||||
assert!(requests[1].starts_with("GET /moved "));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn proxied_urls_skip_the_pinned_resolver_but_keep_the_address_check() {
|
||||
let (url, server, _) = serve_named("localhost", vec![OK_RESPONSE]).await;
|
||||
let media = fetcher(UNROUTABLE, HashSet::new(), UrlPolicy::default(), true)
|
||||
.fetch(url.clone(), policy(2, 0))
|
||||
.await
|
||||
.expect("public host behind a proxy downloads");
|
||||
server.await.expect("server completes");
|
||||
assert_eq!(media.bytes, b"ok");
|
||||
|
||||
let blocked = fetcher(
|
||||
UNROUTABLE,
|
||||
HashSet::from(["localhost"]),
|
||||
UrlPolicy::default(),
|
||||
true,
|
||||
)
|
||||
.fetch(url, policy(2, 0))
|
||||
.await;
|
||||
assert!(matches!(blocked, Err(Error::BlockedUrl)));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ pub enum Error {
|
|||
impl Error {
|
||||
pub fn from_reqwest_before_dispatch(error: reqwest::Error) -> Self {
|
||||
let before_dispatch = !error.is_timeout() && (error.is_connect() || error.is_builder());
|
||||
let message = error.without_url().to_string();
|
||||
let message = describe(error);
|
||||
if before_dispatch {
|
||||
Self::Connect(message)
|
||||
} else {
|
||||
|
|
@ -22,10 +22,18 @@ impl Error {
|
|||
|
||||
impl From<reqwest::Error> for Error {
|
||||
fn from(error: reqwest::Error) -> Self {
|
||||
Self::Network(error.without_url().to_string())
|
||||
Self::Network(describe(error))
|
||||
}
|
||||
}
|
||||
|
||||
fn describe(error: reqwest::Error) -> String {
|
||||
let error = error.without_url();
|
||||
std::iter::successors(std::error::Error::source(&error), |cause| cause.source())
|
||||
.fold(error.to_string(), |message, cause| {
|
||||
format!("{message}: {cause}")
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
#[tokio::test]
|
||||
|
|
@ -47,6 +55,32 @@ mod tests {
|
|||
assert!(!error.to_string().contains("private"));
|
||||
}
|
||||
|
||||
fn root_cause(error: &dyn std::error::Error) -> Option<String> {
|
||||
match error.source() {
|
||||
Some(cause) => root_cause(cause).or_else(|| Some(cause.to_string())),
|
||||
None => None,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn network_error_message_names_the_underlying_cause() {
|
||||
let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind");
|
||||
let address = listener.local_addr().expect("address");
|
||||
drop(listener);
|
||||
let error = reqwest::Client::builder()
|
||||
.no_proxy()
|
||||
.build()
|
||||
.expect("client")
|
||||
.get(format!("http://{address}/private?api_key=secret"))
|
||||
.send()
|
||||
.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();
|
||||
assert!(message.contains(&root_cause), "{message}");
|
||||
assert!(!message.contains("secret"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn request_timeout_is_not_safe_to_retry_as_a_connect_failure() {
|
||||
use std::time::Duration;
|
||||
|
|
|
|||
|
|
@ -20,6 +20,8 @@ bytes.workspace = true
|
|||
litellm-auth.workspace = true
|
||||
litellm-callbacks-legacy.workspace = true
|
||||
litellm-core.workspace = true
|
||||
litellm-auth-gcp.workspace = true
|
||||
litellm-http.workspace = true
|
||||
litellm-llms.workspace = true
|
||||
litellm-types.workspace = true
|
||||
litellm-host-python.workspace = true
|
||||
|
|
|
|||
18
litellm-rust/crates/python-bridge/python_settings.json
Normal file
18
litellm-rust/crates/python-bridge/python_settings.json
Normal file
|
|
@ -0,0 +1,18 @@
|
|||
{
|
||||
"http_settings": [
|
||||
"ssl_verify",
|
||||
"ssl_certificate",
|
||||
"ssl_security_level",
|
||||
"ssl_ecdh_curve",
|
||||
"force_ipv4",
|
||||
"http2",
|
||||
"aiohttp_trust_env",
|
||||
"disable_aiohttp_trust_env",
|
||||
"disable_aiohttp_transport",
|
||||
"user_agent"
|
||||
],
|
||||
"url_policy": [
|
||||
"user_url_validation",
|
||||
"user_url_allowed_hosts"
|
||||
]
|
||||
}
|
||||
354
litellm-rust/crates/python-bridge/src/http.rs
Normal file
354
litellm-rust/crates/python-bridge/src/http.rs
Normal file
|
|
@ -0,0 +1,354 @@
|
|||
use std::{
|
||||
collections::HashSet,
|
||||
path::{Path, PathBuf},
|
||||
sync::{Arc, LazyLock, Mutex, PoisonError},
|
||||
};
|
||||
|
||||
use litellm_http::{
|
||||
HttpClientConfig, HttpClientPool, HttpSettings, HttpSettingsLayer, Resolution, SslVerify,
|
||||
Unsupported,
|
||||
};
|
||||
use litellm_llms::custom_httpx::media::{PublicDnsResolver, UrlPolicy};
|
||||
use pyo3::{prelude::*, types::PyDict};
|
||||
|
||||
use crate::{errors::RustBridgeDeclined, python_settings::PythonSettings};
|
||||
|
||||
static POOL: LazyLock<HttpClientPool> =
|
||||
LazyLock::new(|| HttpClientPool::new(Arc::new(PublicDnsResolver)));
|
||||
|
||||
static REPORTED_UNSUPPORTED: LazyLock<Mutex<HashSet<Unsupported>>> = LazyLock::new(Mutex::default);
|
||||
|
||||
pub(crate) fn pool() -> &'static HttpClientPool {
|
||||
&POOL
|
||||
}
|
||||
|
||||
pub(crate) fn call_config(
|
||||
py: Python<'_>,
|
||||
kwargs: &Bound<'_, PyDict>,
|
||||
asynchronous: bool,
|
||||
) -> PyResult<HttpClientConfig> {
|
||||
let settings = HttpSettings::from_layers([
|
||||
for_call(call_ssl_verify(kwargs)?, asynchronous),
|
||||
HttpSettingsLayer::from_environment(&|name| std::env::var(name).ok()),
|
||||
configured(&PythonSettings::Http.read(py)?)?,
|
||||
])
|
||||
.without_missing_files(&|path: &Path| path.exists());
|
||||
let resolution = Resolution::from(&settings);
|
||||
for unsupported in unreported(&REPORTED_UNSUPPORTED, resolution.unsupported) {
|
||||
PythonSettings::warn(py, &unsupported.to_string())?;
|
||||
}
|
||||
Ok(resolution.config)
|
||||
}
|
||||
|
||||
fn unreported(
|
||||
reported: &Mutex<HashSet<Unsupported>>,
|
||||
unsupported: Vec<Unsupported>,
|
||||
) -> Vec<Unsupported> {
|
||||
let mut reported = reported.lock().unwrap_or_else(PoisonError::into_inner);
|
||||
unsupported
|
||||
.into_iter()
|
||||
.filter(|unsupported| reported.insert(unsupported.clone()))
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(crate) fn url_policy(py: Python<'_>) -> PyResult<UrlPolicy> {
|
||||
let policy: PythonUrlPolicy =
|
||||
PythonSettings::UrlPolicy
|
||||
.read(py)?
|
||||
.extract()
|
||||
.map_err(|error: PyErr| {
|
||||
RustBridgeDeclined::new_err(format!(
|
||||
"litellm URL policy cannot be used by the Rust route: {error}"
|
||||
))
|
||||
})?;
|
||||
Ok(UrlPolicy {
|
||||
validate: policy.user_url_validation,
|
||||
allowed_hosts: policy.user_url_allowed_hosts,
|
||||
})
|
||||
}
|
||||
|
||||
fn call_ssl_verify(kwargs: &Bound<'_, PyDict>) -> PyResult<Option<SslVerify>> {
|
||||
Ok(kwargs
|
||||
.get_item("ssl_verify")?
|
||||
.and_then(|value| ssl_verify(&value)))
|
||||
}
|
||||
|
||||
fn for_call(call_ssl_verify: Option<SslVerify>, asynchronous: bool) -> HttpSettingsLayer {
|
||||
HttpSettingsLayer {
|
||||
ssl_verify: call_ssl_verify,
|
||||
disable_aiohttp_transport: (!asynchronous).then_some(true),
|
||||
..HttpSettingsLayer::default()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(FromPyObject)]
|
||||
struct PythonUrlPolicy {
|
||||
user_url_validation: bool,
|
||||
user_url_allowed_hosts: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(FromPyObject)]
|
||||
struct PythonHttpSettings<'py> {
|
||||
ssl_verify: Bound<'py, PyAny>,
|
||||
ssl_certificate: Option<String>,
|
||||
ssl_security_level: Option<String>,
|
||||
ssl_ecdh_curve: Option<String>,
|
||||
force_ipv4: bool,
|
||||
http2: bool,
|
||||
aiohttp_trust_env: bool,
|
||||
disable_aiohttp_trust_env: bool,
|
||||
disable_aiohttp_transport: bool,
|
||||
user_agent: String,
|
||||
}
|
||||
|
||||
fn configured(value: &Bound<'_, PyAny>) -> PyResult<HttpSettingsLayer> {
|
||||
let python: PythonHttpSettings = value.extract().map_err(|error: PyErr| {
|
||||
RustBridgeDeclined::new_err(format!(
|
||||
"litellm HTTP settings cannot be used by the Rust route: {error}"
|
||||
))
|
||||
})?;
|
||||
Ok(HttpSettingsLayer {
|
||||
ssl_verify: 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: Some(python.force_ipv4),
|
||||
http2: Some(python.http2),
|
||||
aiohttp_trust_env: Some(python.aiohttp_trust_env),
|
||||
disable_aiohttp_trust_env: Some(python.disable_aiohttp_trust_env),
|
||||
disable_aiohttp_transport: Some(python.disable_aiohttp_transport),
|
||||
user_agent: Some(python.user_agent),
|
||||
..HttpSettingsLayer::default()
|
||||
})
|
||||
}
|
||||
|
||||
fn ssl_verify(value: &Bound<'_, PyAny>) -> Option<SslVerify> {
|
||||
if let Ok(enabled) = value.extract::<bool>() {
|
||||
return Some(if enabled {
|
||||
SslVerify::Enabled
|
||||
} else {
|
||||
SslVerify::Disabled
|
||||
});
|
||||
}
|
||||
value
|
||||
.extract::<String>()
|
||||
.ok()
|
||||
.map(|path| SslVerify::parse(&path))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use litellm_http::Verify;
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
use crate::python_settings::CONTRACT;
|
||||
|
||||
fn python_settings<'py>(py: Python<'py>, overrides: &str) -> Bound<'py, PyAny> {
|
||||
let source = format!(
|
||||
"
|
||||
import json
|
||||
import types
|
||||
defaults = dict(
|
||||
ssl_verify=True,
|
||||
ssl_certificate=None,
|
||||
ssl_security_level=None,
|
||||
ssl_ecdh_curve=None,
|
||||
force_ipv4=False,
|
||||
http2=False,
|
||||
aiohttp_trust_env=False,
|
||||
disable_aiohttp_trust_env=False,
|
||||
disable_aiohttp_transport=False,
|
||||
user_agent='litellm/test',
|
||||
)
|
||||
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();
|
||||
py.run(&source, Some(&locals), Some(&locals)).unwrap();
|
||||
locals.get_item("settings").unwrap().unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn default_python_settings_resolve_to_default_settings_with_verification_on() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let layer = configured(&python_settings(py, "")).unwrap();
|
||||
assert_eq!(
|
||||
HttpSettings::from_layers([layer]),
|
||||
HttpSettings {
|
||||
ssl_verify: Some(SslVerify::Enabled),
|
||||
user_agent: Some("litellm/test".into()),
|
||||
..HttpSettings::default()
|
||||
}
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn python_settings_flow_into_the_configured_layer() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let layer = configured(&python_settings(
|
||||
py,
|
||||
"
|
||||
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,
|
||||
disable_aiohttp_trust_env=True,
|
||||
disable_aiohttp_transport=True,
|
||||
user_agent='litellm/9.9.9',
|
||||
",
|
||||
))
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
layer,
|
||||
HttpSettingsLayer {
|
||||
ssl_verify: Some(SslVerify::CaBundle("/etc/ssl/corp.pem".into())),
|
||||
ssl_certificate: Some("/etc/ssl/client.pem".into()),
|
||||
ssl_security_level: Some("2".into()),
|
||||
ssl_ecdh_curve: Some("X25519".into()),
|
||||
force_ipv4: Some(true),
|
||||
http2: Some(true),
|
||||
aiohttp_trust_env: Some(true),
|
||||
disable_aiohttp_trust_env: Some(true),
|
||||
disable_aiohttp_transport: Some(true),
|
||||
user_agent: Some("litellm/9.9.9".into()),
|
||||
..HttpSettingsLayer::default()
|
||||
}
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn user_agent_environment_variable_beats_the_python_default() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let settings = HttpSettings::from_layers([
|
||||
HttpSettingsLayer::from_environment(&|name| {
|
||||
(name == "LITELLM_USER_AGENT").then(|| "operator/1".to_string())
|
||||
}),
|
||||
configured(&python_settings(py, "")).unwrap(),
|
||||
]);
|
||||
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 layer = configured(&python_settings(py, overrides)).unwrap();
|
||||
let config = Resolution::from(&HttpSettings::from_layers([layer])).config;
|
||||
assert_eq!(config.verify, expected);
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ssl_context_global_is_ignored_so_environment_and_defaults_apply() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let layer = configured(&python_settings(py, "ssl_verify=object()")).unwrap();
|
||||
assert_eq!(layer.ssl_verify, None);
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unsupported_settings_are_reported_once_per_process() {
|
||||
let reported = Mutex::default();
|
||||
let curve = Unsupported::EcdhCurve("secp521r1".into());
|
||||
let level = Unsupported::SecurityLevel("@SECLEVEL=1".into());
|
||||
assert_eq!(
|
||||
unreported(&reported, vec![curve.clone(), level.clone()]),
|
||||
[curve.clone(), level]
|
||||
);
|
||||
assert_eq!(unreported(&reported, vec![curve]), []);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mistyped_python_settings_decline_instead_of_raising() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let error = configured(&python_settings(py, "force_ipv4='yes'")).unwrap_err();
|
||||
assert!(error.is_instance_of::<RustBridgeDeclined>(py));
|
||||
});
|
||||
}
|
||||
|
||||
fn configured_ssl_verify(ssl_verify: SslVerify) -> HttpSettingsLayer {
|
||||
HttpSettingsLayer {
|
||||
ssl_verify: Some(ssl_verify),
|
||||
..HttpSettingsLayer::default()
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn call_ssl_verify_beats_the_configured_value() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let kwargs = PyDict::new(py);
|
||||
kwargs.set_item("ssl_verify", false).unwrap();
|
||||
let call = for_call(call_ssl_verify(&kwargs).unwrap(), true);
|
||||
let settings =
|
||||
HttpSettings::from_layers([call, configured_ssl_verify(SslVerify::Enabled)]);
|
||||
assert_eq!(settings.ssl_verify, Some(SslVerify::Disabled));
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn absent_call_ssl_verify_keeps_the_configured_value() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let kwargs = PyDict::new(py);
|
||||
kwargs.set_item("ssl_verify", py.None()).unwrap();
|
||||
let call = for_call(call_ssl_verify(&kwargs).unwrap(), true);
|
||||
let settings =
|
||||
HttpSettings::from_layers([call, configured_ssl_verify(SslVerify::Disabled)]);
|
||||
assert_eq!(settings.ssl_verify, Some(SslVerify::Disabled));
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn live_ssl_context_argument_is_ignored_so_the_configured_value_applies() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let kwargs = PyDict::new(py);
|
||||
kwargs
|
||||
.set_item("ssl_verify", py.eval(c"object()", None, None).unwrap())
|
||||
.unwrap();
|
||||
let call = for_call(call_ssl_verify(&kwargs).unwrap(), true);
|
||||
let settings =
|
||||
HttpSettings::from_layers([call, configured_ssl_verify(SslVerify::Disabled)]);
|
||||
assert_eq!(settings.ssl_verify, Some(SslVerify::Disabled));
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::asynchronous(true, false)]
|
||||
#[case::synchronous(false, true)]
|
||||
fn synchronous_calls_honor_environment_proxies_even_when_aiohttp_opts_out(
|
||||
#[case] asynchronous: bool,
|
||||
#[case] expected: bool,
|
||||
) {
|
||||
let opted_out = HttpSettingsLayer {
|
||||
disable_aiohttp_trust_env: Some(true),
|
||||
disable_aiohttp_transport: Some(false),
|
||||
..HttpSettingsLayer::default()
|
||||
};
|
||||
let settings = HttpSettings::from_layers([for_call(None, asynchronous), opted_out]);
|
||||
assert_eq!(settings.trust_proxy_env, expected);
|
||||
}
|
||||
}
|
||||
|
|
@ -1,7 +1,9 @@
|
|||
mod credentials;
|
||||
mod diagnostics;
|
||||
mod errors;
|
||||
mod http;
|
||||
mod marshal;
|
||||
mod python_settings;
|
||||
mod routes;
|
||||
mod token_counter;
|
||||
|
||||
|
|
|
|||
65
litellm-rust/crates/python-bridge/src/python_settings.rs
Normal file
65
litellm-rust/crates/python-bridge/src/python_settings.rs
Normal file
|
|
@ -0,0 +1,65 @@
|
|||
use pyo3::prelude::*;
|
||||
|
||||
const MODULE: &str = "litellm.rust_bridge.settings";
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub(crate) enum PythonSettings {
|
||||
Http,
|
||||
UrlPolicy,
|
||||
}
|
||||
|
||||
impl PythonSettings {
|
||||
#[cfg(test)]
|
||||
pub(crate) const ALL: [Self; 2] = [Self::Http, Self::UrlPolicy];
|
||||
|
||||
pub(crate) fn name(self) -> &'static str {
|
||||
match self {
|
||||
Self::Http => "http_settings",
|
||||
Self::UrlPolicy => "url_policy",
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn read(self, py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
|
||||
py.import(MODULE)?.getattr(self.name())?.call0()
|
||||
}
|
||||
|
||||
pub(crate) fn warn(py: Python<'_>, message: &str) -> PyResult<()> {
|
||||
py.import(MODULE)?.getattr("warn")?.call1((message,))?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) const CONTRACT: &str = include_str!("../python_settings.json");
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::{collections::BTreeSet, ffi::CString};
|
||||
|
||||
use pyo3::{prelude::*, types::PyDict};
|
||||
|
||||
use super::{CONTRACT, PythonSettings};
|
||||
|
||||
#[test]
|
||||
fn every_settings_group_is_in_the_python_contract() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let locals = PyDict::new(py);
|
||||
locals.set_item("contract", CONTRACT).unwrap();
|
||||
let source = CString::new("import json\nkeys = list(json.loads(contract))").unwrap();
|
||||
py.run(&source, Some(&locals), Some(&locals)).unwrap();
|
||||
let declared: BTreeSet<String> = locals
|
||||
.get_item("keys")
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.extract::<Vec<String>>()
|
||||
.unwrap()
|
||||
.into_iter()
|
||||
.collect();
|
||||
let read: BTreeSet<String> = PythonSettings::ALL
|
||||
.map(|group| group.name().to_owned())
|
||||
.into();
|
||||
assert_eq!(read, declared);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
@ -3,7 +3,10 @@ mod errors;
|
|||
mod host;
|
||||
mod project;
|
||||
|
||||
use std::sync::LazyLock;
|
||||
|
||||
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;
|
||||
|
|
@ -12,6 +15,8 @@ use pyo3::{
|
|||
types::{PyDict, PyTuple},
|
||||
};
|
||||
|
||||
use crate::{errors::RustBridgeDeclined, http};
|
||||
|
||||
const SURFACE: LegacySurface = LegacySurface {
|
||||
call_type: "ocr",
|
||||
input_description: "OCR document processing",
|
||||
|
|
@ -23,6 +28,8 @@ const ASYNC_SURFACE: LegacySurface = LegacySurface {
|
|||
..SURFACE
|
||||
};
|
||||
|
||||
static VERTEX_AUTH: LazyLock<VertexAuth> = LazyLock::new(VertexAuth::default);
|
||||
|
||||
fn run_ocr(
|
||||
py: Python<'_>,
|
||||
request: Bound<'_, PyAny>,
|
||||
|
|
@ -30,7 +37,14 @@ fn run_ocr(
|
|||
kwargs: Bound<'_, PyDict>,
|
||||
asynchronous: bool,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
let client = OcrClient::shared().map_err(errors::to_pyerr)?;
|
||||
let config = http::call_config(py, &kwargs, asynchronous)?;
|
||||
let client = OcrClient::new(
|
||||
http::pool(),
|
||||
&config,
|
||||
http::url_policy(py)?,
|
||||
VERTEX_AUTH.clone(),
|
||||
)
|
||||
.map_err(|error| RustBridgeDeclined::new_err(error.to_string()))?;
|
||||
run_legacy_call(
|
||||
py,
|
||||
if asynchronous { ASYNC_SURFACE } else { SURFACE },
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
57
litellm/rust_bridge/settings.py
Normal file
57
litellm/rust_bridge/settings.py
Normal file
|
|
@ -0,0 +1,57 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
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
|
||||
disable_aiohttp_trust_env: bool
|
||||
disable_aiohttp_transport: bool
|
||||
user_agent: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class UrlPolicy:
|
||||
user_url_validation: bool
|
||||
user_url_allowed_hosts: Sequence[str]
|
||||
|
||||
|
||||
def warn(message: str) -> None:
|
||||
from litellm._logging import verbose_logger
|
||||
|
||||
verbose_logger.warning("%s", message)
|
||||
|
||||
|
||||
def url_policy() -> UrlPolicy:
|
||||
import litellm
|
||||
|
||||
return UrlPolicy(
|
||||
user_url_validation=litellm.user_url_validation,
|
||||
user_url_allowed_hosts=litellm.user_url_allowed_hosts,
|
||||
)
|
||||
|
||||
|
||||
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,
|
||||
disable_aiohttp_trust_env=litellm.disable_aiohttp_trust_env,
|
||||
disable_aiohttp_transport=litellm.disable_aiohttp_transport,
|
||||
user_agent=default_user_agent(),
|
||||
)
|
||||
75
tests/test_litellm/rust_bridge/test_settings.py
Normal file
75
tests/test_litellm/rust_bridge/test_settings.py
Normal file
|
|
@ -0,0 +1,75 @@
|
|||
import dataclasses
|
||||
import logging
|
||||
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())],
|
||||
"url_policy": [field.name for field in dataclasses.fields(settings.url_policy())],
|
||||
}
|
||||
|
||||
|
||||
def test_url_policy_reads_the_litellm_globals(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "user_url_validation", False)
|
||||
monkeypatch.setattr(litellm, "user_url_allowed_hosts", ["docs.internal:8443"])
|
||||
|
||||
assert settings.url_policy() == settings.UrlPolicy(
|
||||
user_url_validation=False,
|
||||
user_url_allowed_hosts=["docs.internal:8443"],
|
||||
)
|
||||
|
||||
|
||||
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)
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_trust_env", True)
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", 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,
|
||||
disable_aiohttp_trust_env=True,
|
||||
disable_aiohttp_transport=True,
|
||||
user_agent=default_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
|
||||
|
||||
|
||||
def test_warn_reaches_the_litellm_logger(caplog: pytest.LogCaptureFixture) -> None:
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM"):
|
||||
settings.warn("ssl_ecdh_curve 'secp521r1' is not supported")
|
||||
|
||||
assert [record.getMessage() for record in caplog.records] == ["ssl_ecdh_curve 'secp521r1' is not supported"]
|
||||
|
|
@ -39,6 +39,7 @@ async def test_proxy_metadata_remains_python_owned(ocr_server: RecordingServer)
|
|||
)
|
||||
events: Final = await recorder.wait_for_async("async_log_success_event")
|
||||
assert response.pages[0].markdown == "native OCR response"
|
||||
assert response._hidden_params["additional_headers"]["x-litellm-rust"] == "true"
|
||||
assert events[0].kwargs["litellm_params"]["metadata"]["user_api_key_auth"].user_id == "ocr-user"
|
||||
assert "metadata" not in ocr_server.requests[0].body
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue