mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
feat(rust): add litellm-http client pool and inject it into the OCR route
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yujong Lee
parent
8e93031c19
commit
abb9618971
19 changed files with 1037 additions and 62 deletions
13
litellm-rust/Cargo.lock
generated
13
litellm-rust/Cargo.lock
generated
|
|
@ -2058,6 +2058,7 @@ dependencies = [
|
|||
"litellm-auth-aws",
|
||||
"litellm-callbacks",
|
||||
"litellm-core-utils",
|
||||
"litellm-http",
|
||||
"litellm-llms",
|
||||
"litellm-types",
|
||||
"mime_guess",
|
||||
|
|
@ -2125,6 +2126,16 @@ dependencies = [
|
|||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-http"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"reqwest 0.12.28",
|
||||
"rstest",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-llms"
|
||||
version = "0.1.0"
|
||||
|
|
@ -2142,6 +2153,7 @@ dependencies = [
|
|||
"litellm-callbacks",
|
||||
"litellm-core-utils",
|
||||
"litellm-framing",
|
||||
"litellm-http",
|
||||
"litellm-types",
|
||||
"reqwest 0.12.28",
|
||||
"rstest",
|
||||
|
|
@ -2166,6 +2178,7 @@ dependencies = [
|
|||
"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" }
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ futures-util.workspace = true
|
|||
base64.workspace = true
|
||||
litellm-auth.workspace = true
|
||||
litellm-auth-aws.workspace = true
|
||||
litellm-http.workspace = true
|
||||
litellm-llms.workspace = true
|
||||
moka.workspace = true
|
||||
mime_guess = "2.0.5"
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
use litellm_http::{HttpClientConfig, HttpClientPool};
|
||||
use litellm_llms::{
|
||||
base_llm::ocr::{error::Error, transformation::LiteLLMOcrResponse},
|
||||
custom_httpx::llm_http_handler::OcrClient,
|
||||
|
|
@ -15,6 +16,10 @@ pub async fn perform(
|
|||
litellm_callbacks::run::run(ocr_machine(client.clone()), &LocalOcrHost::new(request)).await
|
||||
}
|
||||
|
||||
pub async fn ocr(request: LiteLLMOcrRequest) -> Result<LiteLLMOcrResponse, Error> {
|
||||
perform(&OcrClient::shared()?, request).await
|
||||
pub async fn ocr(
|
||||
pool: &HttpClientPool,
|
||||
config: &HttpClientConfig,
|
||||
request: LiteLLMOcrRequest,
|
||||
) -> Result<LiteLLMOcrResponse, Error> {
|
||||
perform(&OcrClient::new(pool, config)?, request).await
|
||||
}
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ use litellm_callbacks::{
|
|||
host::{Host, HostOp, HostResult},
|
||||
machine::{HostFailure, Machine, MachineStep},
|
||||
};
|
||||
use litellm_http::{HttpClientConfig, HttpClientPool, HttpSettings, Verify};
|
||||
use litellm_llms::{
|
||||
base_llm::ocr::{
|
||||
error::Error as OcrError,
|
||||
|
|
@ -171,25 +172,44 @@ async fn facade_retains_native_response_when_requested() {
|
|||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn facade_uses_the_injected_http_client() {
|
||||
async fn facade_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(),
|
||||
let settings = HttpSettings {
|
||||
user_agent: Some("host-owned/1".into()),
|
||||
..HttpSettings::default()
|
||||
};
|
||||
let config = HttpClientConfig::resolve(&settings, None).unwrap();
|
||||
crate::ocr::client::ocr(
|
||||
&HttpClientPool::new(),
|
||||
&config,
|
||||
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"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unbuildable_http_configuration_fails_before_dispatch() {
|
||||
let (base, _seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await;
|
||||
let config = HttpClientConfig {
|
||||
verify: Verify::CaBundle(std::env::temp_dir().join("litellm-ocr-missing-bundle.pem")),
|
||||
..HttpClientConfig::resolve(&HttpSettings::default(), None).unwrap()
|
||||
};
|
||||
let error = crate::ocr::client::ocr(
|
||||
&HttpClientPool::new(),
|
||||
&config,
|
||||
wire_request("mistral/model", &base, json!({})),
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
server.abort();
|
||||
assert!(matches!(
|
||||
error,
|
||||
OcrError::Transport(litellm_llms::custom_httpx::transport::Error::Connect(_))
|
||||
));
|
||||
assert!(error.to_string().contains("litellm-ocr-missing-bundle.pem"));
|
||||
}
|
||||
|
||||
fn event_name(event: &CallEvent) -> &'static str {
|
||||
|
|
|
|||
14
litellm-rust/crates/http/Cargo.toml
Normal file
14
litellm-rust/crates/http/Cargo.toml
Normal file
|
|
@ -0,0 +1,14 @@
|
|||
[package]
|
||||
name = "litellm-http"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
reqwest.workspace = true
|
||||
thiserror.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
rstest.workspace = true
|
||||
tokio.workspace = true
|
||||
327
litellm-rust/crates/http/src/config.rs
Normal file
327
litellm-rust/crates/http/src/config.rs
Normal file
|
|
@ -0,0 +1,327 @@
|
|||
use std::{
|
||||
net::{IpAddr, Ipv4Addr},
|
||||
path::{Path, PathBuf},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use crate::settings::{HttpSettings, SslVerify};
|
||||
|
||||
#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)]
|
||||
pub enum Error {
|
||||
#[error("{setting} cannot be expressed with rustls: {reason}")]
|
||||
Unsupported {
|
||||
setting: &'static str,
|
||||
reason: String,
|
||||
},
|
||||
#[error("could not read {}: {message}", path.display())]
|
||||
Read { path: PathBuf, message: String },
|
||||
#[error("{} is not a PEM file: {message}", path.display())]
|
||||
InvalidPem { path: PathBuf, message: String },
|
||||
#[error("could not build the HTTP client: {0}")]
|
||||
Client(String),
|
||||
}
|
||||
|
||||
impl From<reqwest::Error> for Error {
|
||||
fn from(error: reqwest::Error) -> Self {
|
||||
Self::Client(error.without_url().to_string())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
|
||||
pub enum Verify {
|
||||
Disabled,
|
||||
CaBundle(PathBuf),
|
||||
BuiltInRoots,
|
||||
}
|
||||
|
||||
/// One fully resolved client configuration. Every field is a plain value so the pool can
|
||||
/// key cached clients on it.
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
|
||||
pub struct HttpClientConfig {
|
||||
pub verify: Verify,
|
||||
pub client_certificate: Option<PathBuf>,
|
||||
pub force_ipv4: bool,
|
||||
pub http2: bool,
|
||||
pub user_agent: Option<String>,
|
||||
pub trust_proxy_env: bool,
|
||||
pub connect_timeout: Duration,
|
||||
pub request_timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
impl HttpClientConfig {
|
||||
/// Port of `get_ssl_verify` + `get_ssl_configuration`: the per-call value wins, then the
|
||||
/// configured (environment-overlaid) `ssl_verify`, then `SSL_CERT_FILE`, then the built-in
|
||||
/// roots. Settings rustls has no equivalent for are an error instead of a silent no-op.
|
||||
pub fn resolve(
|
||||
settings: &HttpSettings,
|
||||
per_call_ssl_verify: Option<&SslVerify>,
|
||||
) -> Result<Self, Error> {
|
||||
if let Some(level) = &settings.ssl_security_level {
|
||||
return Err(Error::Unsupported {
|
||||
setting: "ssl_security_level",
|
||||
reason: format!("OpenSSL cipher string {level:?} has no rustls equivalent"),
|
||||
});
|
||||
}
|
||||
if let Some(curve) = &settings.ssl_ecdh_curve {
|
||||
return Err(Error::Unsupported {
|
||||
setting: "ssl_ecdh_curve",
|
||||
reason: format!("key exchange group {curve:?} is fixed by the rustls provider"),
|
||||
});
|
||||
}
|
||||
let verify = match per_call_ssl_verify.or(settings.ssl_verify.as_ref()) {
|
||||
Some(SslVerify::Disabled) => Verify::Disabled,
|
||||
Some(SslVerify::CaBundle(path)) => Verify::CaBundle(path.clone()),
|
||||
Some(SslVerify::Enabled) | None => settings
|
||||
.ssl_cert_file
|
||||
.clone()
|
||||
.map_or(Verify::BuiltInRoots, Verify::CaBundle),
|
||||
};
|
||||
Ok(Self {
|
||||
verify,
|
||||
client_certificate: settings.ssl_certificate.clone(),
|
||||
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,
|
||||
request_timeout: settings.request_timeout,
|
||||
})
|
||||
}
|
||||
|
||||
/// A builder carrying every shared setting; variants add their own policy on top.
|
||||
pub fn client_builder(&self) -> Result<reqwest::ClientBuilder, Error> {
|
||||
let base = reqwest::Client::builder().connect_timeout(self.connect_timeout);
|
||||
let with_roots = match &self.verify {
|
||||
Verify::Disabled => base.danger_accept_invalid_certs(true),
|
||||
Verify::BuiltInRoots => base,
|
||||
Verify::CaBundle(path) => {
|
||||
let pem = read(path)?;
|
||||
let certificates =
|
||||
reqwest::Certificate::from_pem_bundle(&pem).map_err(|error| {
|
||||
Error::InvalidPem {
|
||||
path: path.clone(),
|
||||
message: error.without_url().to_string(),
|
||||
}
|
||||
})?;
|
||||
if certificates.is_empty() {
|
||||
return Err(Error::InvalidPem {
|
||||
path: path.clone(),
|
||||
message: "no certificates found".into(),
|
||||
});
|
||||
}
|
||||
certificates.into_iter().fold(
|
||||
base.tls_built_in_root_certs(false),
|
||||
|builder, certificate| builder.add_root_certificate(certificate),
|
||||
)
|
||||
}
|
||||
};
|
||||
let with_identity = match &self.client_certificate {
|
||||
None => with_roots,
|
||||
Some(path) => {
|
||||
let identity = reqwest::Identity::from_pem(&read(path)?).map_err(|error| {
|
||||
Error::InvalidPem {
|
||||
path: path.clone(),
|
||||
message: error.without_url().to_string(),
|
||||
}
|
||||
})?;
|
||||
with_roots.identity(identity)
|
||||
}
|
||||
};
|
||||
let with_address = if self.force_ipv4 {
|
||||
with_identity.local_address(IpAddr::V4(Ipv4Addr::UNSPECIFIED))
|
||||
} else {
|
||||
with_identity
|
||||
};
|
||||
let with_protocol = if self.http2 {
|
||||
with_address
|
||||
} else {
|
||||
with_address.http1_only()
|
||||
};
|
||||
let with_agent = match &self.user_agent {
|
||||
Some(agent) => with_protocol.user_agent(agent),
|
||||
None => with_protocol,
|
||||
};
|
||||
let with_proxy = if self.trust_proxy_env {
|
||||
with_agent
|
||||
} else {
|
||||
with_agent.no_proxy()
|
||||
};
|
||||
Ok(match self.request_timeout {
|
||||
Some(timeout) => with_proxy.timeout(timeout),
|
||||
None => with_proxy,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
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(),
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn no_env(_: &str) -> Option<String> {
|
||||
None
|
||||
}
|
||||
|
||||
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()
|
||||
}
|
||||
.with_environment(&no_env)
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::default(settings(None, None), None, Verify::BuiltInRoots)]
|
||||
#[case::setting_disables(
|
||||
settings(Some(SslVerify::Disabled), Some("/env/roots.pem")),
|
||||
None,
|
||||
Verify::Disabled
|
||||
)]
|
||||
#[case::setting_bundle(
|
||||
settings(Some(SslVerify::CaBundle("/configured.pem".into())), Some("/env/roots.pem")),
|
||||
None,
|
||||
Verify::CaBundle("/configured.pem".into())
|
||||
)]
|
||||
#[case::enabled_uses_cert_file(
|
||||
settings(Some(SslVerify::Enabled), Some("/env/roots.pem")),
|
||||
None,
|
||||
Verify::CaBundle("/env/roots.pem".into())
|
||||
)]
|
||||
#[case::unset_uses_cert_file(settings(None, Some("/env/roots.pem")), None, Verify::CaBundle("/env/roots.pem".into()))]
|
||||
#[case::per_call_beats_setting(
|
||||
settings(Some(SslVerify::Disabled), None),
|
||||
Some(SslVerify::Enabled),
|
||||
Verify::BuiltInRoots
|
||||
)]
|
||||
#[case::per_call_disables(
|
||||
settings(Some(SslVerify::CaBundle("/configured.pem".into())), Some("/env/roots.pem")),
|
||||
Some(SslVerify::Disabled),
|
||||
Verify::Disabled
|
||||
)]
|
||||
#[case::per_call_bundle(
|
||||
settings(None, Some("/env/roots.pem")),
|
||||
Some(SslVerify::CaBundle("/call.pem".into())),
|
||||
Verify::CaBundle("/call.pem".into())
|
||||
)]
|
||||
#[case::per_call_enabled_still_honours_cert_file(
|
||||
settings(Some(SslVerify::Disabled), Some("/env/roots.pem")),
|
||||
Some(SslVerify::Enabled),
|
||||
Verify::CaBundle("/env/roots.pem".into())
|
||||
)]
|
||||
fn verify_follows_per_call_then_setting_then_cert_file(
|
||||
#[case] settings: HttpSettings,
|
||||
#[case] per_call: Option<SslVerify>,
|
||||
#[case] expected: Verify,
|
||||
) {
|
||||
let config = HttpClientConfig::resolve(&settings, per_call.as_ref()).unwrap();
|
||||
assert_eq!(config.verify, expected);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ssl_verify_environment_variable_beats_the_configured_setting() {
|
||||
let settings = HttpSettings {
|
||||
ssl_verify: Some(SslVerify::Disabled),
|
||||
..HttpSettings::default()
|
||||
}
|
||||
.with_environment(&|name: &str| (name == "SSL_VERIFY").then(|| "true".to_string()));
|
||||
let config = HttpClientConfig::resolve(&settings, None).unwrap();
|
||||
assert_eq!(config.verify, Verify::BuiltInRoots);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cipher_strings_are_rejected_rather_than_ignored() {
|
||||
let settings = HttpSettings {
|
||||
ssl_security_level: Some("DEFAULT@SECLEVEL=1".into()),
|
||||
..HttpSettings::default()
|
||||
};
|
||||
assert!(matches!(
|
||||
HttpClientConfig::resolve(&settings, None),
|
||||
Err(Error::Unsupported {
|
||||
setting: "ssl_security_level",
|
||||
..
|
||||
})
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ecdh_curves_are_rejected_rather_than_ignored() {
|
||||
let settings = HttpSettings {
|
||||
ssl_ecdh_curve: Some("X25519".into()),
|
||||
..HttpSettings::default()
|
||||
};
|
||||
assert!(matches!(
|
||||
HttpClientConfig::resolve(&settings, None),
|
||||
Err(Error::Unsupported {
|
||||
setting: "ssl_ecdh_curve",
|
||||
..
|
||||
})
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn connection_settings_carry_over_unchanged() {
|
||||
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),
|
||||
request_timeout: Some(Duration::from_secs(70)),
|
||||
..HttpSettings::default()
|
||||
};
|
||||
let config = HttpClientConfig::resolve(&settings, None).unwrap();
|
||||
assert_eq!(
|
||||
config,
|
||||
HttpClientConfig {
|
||||
verify: Verify::BuiltInRoots,
|
||||
client_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),
|
||||
request_timeout: Some(Duration::from_secs(70)),
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[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()),
|
||||
..HttpClientConfig::resolve(&HttpSettings::default(), None).unwrap()
|
||||
};
|
||||
assert!(matches!(
|
||||
config.client_builder(),
|
||||
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()),
|
||||
..HttpClientConfig::resolve(&HttpSettings::default(), None).unwrap()
|
||||
};
|
||||
let result = config.client_builder().map(drop);
|
||||
std::fs::remove_file(&path).unwrap();
|
||||
assert!(matches!(
|
||||
result,
|
||||
Err(Error::InvalidPem { path: reported, .. }) if reported == path
|
||||
));
|
||||
}
|
||||
}
|
||||
11
litellm-rust/crates/http/src/lib.rs
Normal file
11
litellm-rust/crates/http/src/lib.rs
Normal file
|
|
@ -0,0 +1,11 @@
|
|||
//! Rust counterpart of `litellm/llms/custom_httpx/http_handler.py`: the plain HTTP settings
|
||||
//! LiteLLM exposes, their resolution into one typed client configuration, and a pool that
|
||||
//! caches `reqwest::Client`s per resolved configuration.
|
||||
|
||||
mod config;
|
||||
mod pool;
|
||||
mod settings;
|
||||
|
||||
pub use config::{Error, HttpClientConfig, Verify};
|
||||
pub use pool::{ClientVariant, HttpClientPool};
|
||||
pub use settings::{HttpSettings, SslVerify};
|
||||
172
litellm-rust/crates/http/src/pool.rs
Normal file
172
litellm-rust/crates/http/src/pool.rs
Normal file
|
|
@ -0,0 +1,172 @@
|
|||
use std::{
|
||||
collections::HashMap,
|
||||
sync::{Mutex, PoisonError},
|
||||
};
|
||||
|
||||
use crate::config::{Error, HttpClientConfig};
|
||||
|
||||
/// The client shapes routes need; each is the shared base plus one policy.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
|
||||
pub enum ClientVariant {
|
||||
Provider,
|
||||
NoRedirect,
|
||||
/// Media downloads: no redirects (the fetcher validates each hop) and never a proxy.
|
||||
Media,
|
||||
}
|
||||
|
||||
impl ClientVariant {
|
||||
fn apply(self, builder: reqwest::ClientBuilder) -> reqwest::ClientBuilder {
|
||||
match self {
|
||||
Self::Provider => builder,
|
||||
Self::NoRedirect => builder.redirect(reqwest::redirect::Policy::none()),
|
||||
Self::Media => builder
|
||||
.redirect(reqwest::redirect::Policy::none())
|
||||
.no_proxy(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Counterpart of `get_async_httpx_client`: one `reqwest::Client` per resolved configuration
|
||||
/// and variant, built on first use and shared afterwards.
|
||||
#[derive(Default)]
|
||||
pub struct HttpClientPool {
|
||||
clients: Mutex<HashMap<(HttpClientConfig, ClientVariant), reqwest::Client>>,
|
||||
}
|
||||
|
||||
impl HttpClientPool {
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub fn client(
|
||||
&self,
|
||||
config: &HttpClientConfig,
|
||||
variant: ClientVariant,
|
||||
) -> Result<reqwest::Client, Error> {
|
||||
self.client_with(config, variant, |builder| builder)
|
||||
}
|
||||
|
||||
/// Like [`Self::client`], with a caller hook for builder options that are not plain values
|
||||
/// (a DNS resolver, for example). The hook only runs when the client is first built.
|
||||
pub fn client_with(
|
||||
&self,
|
||||
config: &HttpClientConfig,
|
||||
variant: ClientVariant,
|
||||
customize: impl FnOnce(reqwest::ClientBuilder) -> reqwest::ClientBuilder,
|
||||
) -> Result<reqwest::Client, Error> {
|
||||
let key = (config.clone(), variant);
|
||||
let mut clients = self.clients.lock().unwrap_or_else(PoisonError::into_inner);
|
||||
if let Some(client) = clients.get(&key) {
|
||||
return Ok(client.clone());
|
||||
}
|
||||
let client = customize(variant.apply(config.client_builder()?)).build()?;
|
||||
clients.insert(key, client.clone());
|
||||
Ok(client)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::{cell::Cell, time::Duration};
|
||||
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
|
||||
use super::*;
|
||||
use crate::{HttpSettings, Verify};
|
||||
|
||||
fn config(user_agent: &str) -> HttpClientConfig {
|
||||
HttpClientConfig {
|
||||
user_agent: Some(user_agent.into()),
|
||||
..HttpClientConfig::resolve(&HttpSettings::default(), None).unwrap()
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn clients_are_built_once_per_config_and_variant() {
|
||||
let pool = HttpClientPool::new();
|
||||
let builds = Cell::new(0);
|
||||
let build = |config: &HttpClientConfig, variant| {
|
||||
pool.client_with(config, variant, |builder| {
|
||||
builds.set(builds.get() + 1);
|
||||
builder
|
||||
})
|
||||
.unwrap()
|
||||
};
|
||||
build(&config("a"), ClientVariant::Provider);
|
||||
build(&config("a"), ClientVariant::Provider);
|
||||
assert_eq!(builds.get(), 1);
|
||||
build(&config("a"), ClientVariant::NoRedirect);
|
||||
assert_eq!(builds.get(), 2);
|
||||
build(&config("b"), ClientVariant::Provider);
|
||||
assert_eq!(builds.get(), 3);
|
||||
build(&config("b"), ClientVariant::Provider);
|
||||
build(&config("a"), ClientVariant::NoRedirect);
|
||||
assert_eq!(builds.get(), 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_failures_are_not_cached() {
|
||||
let pool = HttpClientPool::new();
|
||||
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());
|
||||
}
|
||||
|
||||
async fn serve_once(status_line: &'static str) -> (String, tokio::task::JoinHandle<String>) {
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let base = format!("http://{}", listener.local_addr().unwrap());
|
||||
let server = tokio::spawn(async move {
|
||||
let (mut socket, _) = listener.accept().await.unwrap();
|
||||
let mut request = vec![0u8; 4096];
|
||||
let read = socket.read(&mut request).await.unwrap();
|
||||
socket
|
||||
.write_all(
|
||||
format!("{status_line}\r\nLocation: /elsewhere\r\nContent-Length: 0\r\nConnection: close\r\n\r\n")
|
||||
.as_bytes(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
String::from_utf8_lossy(&request[..read]).into_owned()
|
||||
});
|
||||
(base, server)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn provider_client_sends_the_configured_user_agent_over_http1() {
|
||||
let (base, server) = serve_once("HTTP/1.1 204 No Content").await;
|
||||
let config = HttpClientConfig {
|
||||
connect_timeout: Duration::from_secs(2),
|
||||
..config("litellm-test/9")
|
||||
};
|
||||
let response = HttpClientPool::new()
|
||||
.client(&config, ClientVariant::Provider)
|
||||
.unwrap()
|
||||
.get(&base)
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), 204);
|
||||
assert_eq!(response.version(), reqwest::Version::HTTP_11);
|
||||
let request = server.await.unwrap();
|
||||
assert!(request.contains("user-agent: litellm-test/9"), "{request}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn no_redirect_variant_returns_the_redirect_instead_of_following_it() {
|
||||
let (base, server) = serve_once("HTTP/1.1 302 Found").await;
|
||||
let response = HttpClientPool::new()
|
||||
.client(&config("a"), ClientVariant::NoRedirect)
|
||||
.unwrap()
|
||||
.get(&base)
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), 302);
|
||||
assert_eq!(response.headers()["location"], "/elsewhere");
|
||||
server.await.unwrap();
|
||||
}
|
||||
}
|
||||
171
litellm-rust/crates/http/src/settings.rs
Normal file
171
litellm-rust/crates/http/src/settings.rs
Normal file
|
|
@ -0,0 +1,171 @@
|
|||
use std::{path::PathBuf, time::Duration};
|
||||
|
||||
/// `litellm.ssl_verify` / `SSL_VERIFY`: a bool or a CA bundle path.
|
||||
#[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)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The plain inputs `http_handler.py` reads from `litellm.*` globals and the environment.
|
||||
#[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 request_timeout: Option<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: false,
|
||||
connect_timeout: Duration::from_secs(5),
|
||||
request_timeout: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl HttpSettings {
|
||||
/// Overlay the environment variables `http_handler.py` consults, with the same precedence:
|
||||
/// `SSL_VERIFY`, `SSL_CERTIFICATE`, `SSL_SECURITY_LEVEL`, `SSL_ECDH_CURVE` and
|
||||
/// `LITELLM_USER_AGENT` win over the configured value; `SSL_CERT_FILE` only applies when
|
||||
/// verification is on without an explicit bundle; `LITELLM_HTTP2` and `AIOHTTP_TRUST_ENV`
|
||||
/// can only turn their switch on.
|
||||
pub fn with_environment(self, 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"));
|
||||
Self {
|
||||
ssl_verify: env("SSL_VERIFY")
|
||||
.map(|value| SslVerify::parse(&value))
|
||||
.or(self.ssl_verify),
|
||||
ssl_cert_file: env("SSL_CERT_FILE")
|
||||
.map(PathBuf::from)
|
||||
.or(self.ssl_cert_file),
|
||||
ssl_certificate: env("SSL_CERTIFICATE")
|
||||
.map(PathBuf::from)
|
||||
.or(self.ssl_certificate),
|
||||
ssl_security_level: env("SSL_SECURITY_LEVEL").or(self.ssl_security_level),
|
||||
ssl_ecdh_curve: env("SSL_ECDH_CURVE").or(self.ssl_ecdh_curve),
|
||||
http2: self.http2 || enabled("LITELLM_HTTP2"),
|
||||
user_agent: env("LITELLM_USER_AGENT").or(self.user_agent),
|
||||
trust_proxy_env: self.trust_proxy_env || enabled("AIOHTTP_TRUST_ENV"),
|
||||
..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 environment_overrides_configured_ssl_values() {
|
||||
let settings = HttpSettings {
|
||||
ssl_verify: Some(SslVerify::Enabled),
|
||||
ssl_certificate: Some("/configured/client.pem".into()),
|
||||
ssl_security_level: Some("configured".into()),
|
||||
user_agent: Some("configured/1".into()),
|
||||
..HttpSettings::default()
|
||||
}
|
||||
.with_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"),
|
||||
]));
|
||||
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 missing_environment_keeps_configured_values() {
|
||||
let configured = HttpSettings {
|
||||
ssl_verify: Some(SslVerify::CaBundle("/configured/roots.pem".into())),
|
||||
http2: true,
|
||||
trust_proxy_env: true,
|
||||
user_agent: Some("configured/1".into()),
|
||||
..HttpSettings::default()
|
||||
};
|
||||
assert_eq!(configured.clone().with_environment(&no_env), configured);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case("true", true)]
|
||||
#[case("True", true)]
|
||||
#[case("false", false)]
|
||||
#[case("1", false)]
|
||||
fn boolean_switches_only_turn_on_for_true(#[case] value: &'static str, #[case] expected: bool) {
|
||||
let env = move |name: &str| match name {
|
||||
"LITELLM_HTTP2" | "AIOHTTP_TRUST_ENV" => Some(value.to_string()),
|
||||
_ => None,
|
||||
};
|
||||
let settings = HttpSettings::default().with_environment(&env);
|
||||
assert_eq!(settings.http2, expected);
|
||||
assert_eq!(settings.trust_proxy_env, expected);
|
||||
}
|
||||
}
|
||||
|
|
@ -17,6 +17,7 @@ litellm-auth-azure.workspace = true
|
|||
litellm-auth-gcp.workspace = true
|
||||
litellm-callbacks.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_callbacks::event::{Passthrough, WireRequest};
|
||||
use litellm_http::{ClientVariant, HttpClientConfig, HttpClientPool};
|
||||
use serde::{Serialize, de::DeserializeOwned};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
|
|
@ -11,9 +10,8 @@ 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::{
|
||||
|
|
@ -44,30 +42,15 @@ 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) -> Result<Self, transport::Error> {
|
||||
Ok(Self {
|
||||
provider_http,
|
||||
polling_http: no_redirect_http()?,
|
||||
document_fetcher,
|
||||
provider_http: pool.client(config, ClientVariant::Provider)?,
|
||||
polling_http: pool.client(config, ClientVariant::NoRedirect)?,
|
||||
document_fetcher: MediaFetcher::new(pool, config)?,
|
||||
vertex_auth: VertexAuth::default(),
|
||||
})
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
|
|
@ -88,21 +71,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>(
|
||||
|
|
@ -318,6 +296,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, 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")]
|
||||
|
|
@ -63,23 +62,30 @@ 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,
|
||||
) -> Result<Self, litellm_http::Error> {
|
||||
Self::with_resolvers(
|
||||
pool,
|
||||
config,
|
||||
Arc::new(PublicDnsResolver),
|
||||
Arc::new(SystemAddressResolver),
|
||||
)
|
||||
}
|
||||
|
||||
fn with_resolvers<R>(
|
||||
pool: &HttpClientPool,
|
||||
config: &HttpClientConfig,
|
||||
transport_resolver: Arc<R>,
|
||||
address_resolver: Arc<dyn AddressResolver>,
|
||||
) -> Result<Self, reqwest::Error>
|
||||
) -> Result<Self, litellm_http::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()?;
|
||||
let client = pool.client_with(config, ClientVariant::Media, |builder| {
|
||||
builder.dns_resolver(transport_resolver)
|
||||
})?;
|
||||
Ok(Self {
|
||||
client,
|
||||
address_resolver,
|
||||
|
|
@ -281,6 +287,7 @@ impl Resolve for PublicDnsResolver {
|
|||
mod tests {
|
||||
use std::collections::HashSet;
|
||||
|
||||
use litellm_http::HttpSettings;
|
||||
use tokio::{
|
||||
io::{AsyncReadExt, AsyncWriteExt},
|
||||
net::TcpListener,
|
||||
|
|
@ -365,6 +372,8 @@ mod tests {
|
|||
blocked_hosts: HashSet<&'static str>,
|
||||
) -> MediaFetcher {
|
||||
MediaFetcher::with_resolvers(
|
||||
&HttpClientPool::new(),
|
||||
&HttpClientConfig::resolve(&HttpSettings::default(), None).unwrap(),
|
||||
Arc::new(LoopbackDnsResolver(address)),
|
||||
Arc::new(TestAddressResolver { blocked_hosts }),
|
||||
)
|
||||
|
|
@ -542,7 +551,12 @@ 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(),
|
||||
&HttpClientConfig::resolve(&litellm_http::HttpSettings::default(), None)
|
||||
.expect("default settings resolve"),
|
||||
)
|
||||
.expect("media fetcher builds");
|
||||
let url =
|
||||
Url::parse("https://user:password@8.8.8.8/document").expect("credentialed URL parses");
|
||||
assert!(matches!(
|
||||
|
|
|
|||
|
|
@ -26,6 +26,12 @@ impl From<reqwest::Error> for Error {
|
|||
}
|
||||
}
|
||||
|
||||
impl From<litellm_http::Error> for Error {
|
||||
fn from(error: litellm_http::Error) -> Self {
|
||||
Self::Connect(error.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
#[tokio::test]
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ bytes.workspace = true
|
|||
litellm-auth.workspace = true
|
||||
litellm-callbacks-legacy.workspace = true
|
||||
litellm-core.workspace = true
|
||||
litellm-http.workspace = true
|
||||
litellm-llms.workspace = true
|
||||
litellm-types.workspace = true
|
||||
litellm-host-python.workspace = true
|
||||
|
|
|
|||
234
litellm-rust/crates/python-bridge/src/http.rs
Normal file
234
litellm-rust/crates/python-bridge/src/http.rs
Normal file
|
|
@ -0,0 +1,234 @@
|
|||
use std::{path::PathBuf, sync::LazyLock};
|
||||
|
||||
use litellm_http::{HttpClientConfig, HttpClientPool, HttpSettings, SslVerify};
|
||||
use pyo3::{prelude::*, types::PyDict};
|
||||
|
||||
use crate::errors::RustBridgeDeclined;
|
||||
|
||||
static POOL: LazyLock<HttpClientPool> = LazyLock::new(HttpClientPool::new);
|
||||
|
||||
/// Keyword arguments that carry a live Python HTTP client or session. They cannot cross into
|
||||
/// Rust, so a call that supplies one stays on the Python path.
|
||||
const LIVE_CLIENT_ARGUMENTS: [&str; 3] = ["client", "shared_session", "aclient_session"];
|
||||
|
||||
pub(crate) fn pool() -> &'static HttpClientPool {
|
||||
&POOL
|
||||
}
|
||||
|
||||
/// The client configuration for one call: the process settings from the `litellm` module and
|
||||
/// the environment, narrowed by the call's own `ssl_verify`.
|
||||
pub(crate) fn call_config(
|
||||
py: Python<'_>,
|
||||
kwargs: &Bound<'_, PyDict>,
|
||||
) -> PyResult<HttpClientConfig> {
|
||||
decline_live_clients(kwargs)?;
|
||||
let settings = settings(py.import("litellm")?.as_any())?
|
||||
.with_environment(&|name| std::env::var(name).ok());
|
||||
let per_call = kwargs
|
||||
.get_item("ssl_verify")?
|
||||
.map(|value| ssl_verify(&value, "ssl_verify"))
|
||||
.transpose()?
|
||||
.flatten();
|
||||
HttpClientConfig::resolve(&settings, per_call.as_ref())
|
||||
.map_err(|error| RustBridgeDeclined::new_err(error.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) fn decline_live_clients(kwargs: &Bound<'_, PyDict>) -> PyResult<()> {
|
||||
for name in LIVE_CLIENT_ARGUMENTS {
|
||||
if kwargs.get_item(name)?.is_some_and(|value| !value.is_none()) {
|
||||
return Err(RustBridgeDeclined::new_err(format!(
|
||||
"{name} is a live Python HTTP client and cannot be used by the Rust route"
|
||||
)));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Read the `litellm.*` globals `http_handler.py` consults. `globals` is the `litellm` module in
|
||||
/// production and any attribute holder in tests.
|
||||
pub(crate) fn settings(globals: &Bound<'_, PyAny>) -> PyResult<HttpSettings> {
|
||||
Ok(HttpSettings {
|
||||
ssl_verify: ssl_verify(&globals.getattr("ssl_verify")?, "litellm.ssl_verify")?,
|
||||
ssl_certificate: optional_path(globals, "ssl_certificate")?,
|
||||
ssl_security_level: globals.getattr("ssl_security_level")?.extract()?,
|
||||
ssl_ecdh_curve: globals.getattr("ssl_ecdh_curve")?.extract()?,
|
||||
force_ipv4: globals.getattr("force_ipv4")?.extract()?,
|
||||
http2: globals.getattr("http2")?.extract()?,
|
||||
trust_proxy_env: globals.getattr("aiohttp_trust_env")?.extract()?,
|
||||
..HttpSettings::default()
|
||||
})
|
||||
}
|
||||
|
||||
fn optional_path(globals: &Bound<'_, PyAny>, name: &str) -> PyResult<Option<PathBuf>> {
|
||||
Ok(globals
|
||||
.getattr(name)?
|
||||
.extract::<Option<String>>()?
|
||||
.map(PathBuf::from))
|
||||
}
|
||||
|
||||
fn ssl_verify(value: &Bound<'_, PyAny>, name: &str) -> PyResult<Option<SslVerify>> {
|
||||
if value.is_none() {
|
||||
return Ok(None);
|
||||
}
|
||||
if let Ok(enabled) = value.extract::<bool>() {
|
||||
return Ok(Some(if enabled {
|
||||
SslVerify::Enabled
|
||||
} else {
|
||||
SslVerify::Disabled
|
||||
}));
|
||||
}
|
||||
if let Ok(path) = value.extract::<String>() {
|
||||
return Ok(Some(SslVerify::CaBundle(PathBuf::from(path))));
|
||||
}
|
||||
Err(RustBridgeDeclined::new_err(format!(
|
||||
"{name} is a live Python object and cannot be used by the Rust route"
|
||||
)))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use litellm_http::Verify;
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn eval<'py>(py: Python<'py>, source: &std::ffi::CStr) -> Bound<'py, PyDict> {
|
||||
let locals = PyDict::new(py);
|
||||
py.run(source, Some(&locals), Some(&locals)).unwrap();
|
||||
locals
|
||||
}
|
||||
|
||||
fn globals<'py>(py: Python<'py>, overrides: &str) -> Bound<'py, PyAny> {
|
||||
let source = format!(
|
||||
"
|
||||
import types
|
||||
globals = types.SimpleNamespace(
|
||||
ssl_verify=True,
|
||||
ssl_certificate=None,
|
||||
ssl_security_level=None,
|
||||
ssl_ecdh_curve=None,
|
||||
force_ipv4=False,
|
||||
http2=False,
|
||||
aiohttp_trust_env=False,
|
||||
)
|
||||
{overrides}
|
||||
"
|
||||
);
|
||||
let source = std::ffi::CString::new(source).unwrap();
|
||||
eval(py, &source).get_item("globals").unwrap().unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn default_globals_produce_default_settings_with_verification_on() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let settings = settings(&globals(py, "")).unwrap();
|
||||
assert_eq!(
|
||||
settings,
|
||||
HttpSettings {
|
||||
ssl_verify: Some(SslVerify::Enabled),
|
||||
..HttpSettings::default()
|
||||
}
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn globals_flow_into_settings() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let settings = settings(&globals(
|
||||
py,
|
||||
"
|
||||
globals.ssl_verify = '/etc/ssl/corp.pem'
|
||||
globals.ssl_certificate = '/etc/ssl/client.pem'
|
||||
globals.ssl_security_level = '2'
|
||||
globals.ssl_ecdh_curve = 'X25519'
|
||||
globals.force_ipv4 = True
|
||||
globals.http2 = True
|
||||
globals.aiohttp_trust_env = True
|
||||
",
|
||||
))
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
settings,
|
||||
HttpSettings {
|
||||
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: true,
|
||||
http2: true,
|
||||
trust_proxy_env: true,
|
||||
..HttpSettings::default()
|
||||
}
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn disabled_verification_global_resolves_to_disabled() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let settings = settings(&globals(py, "globals.ssl_verify = False")).unwrap();
|
||||
assert_eq!(settings.ssl_verify, Some(SslVerify::Disabled));
|
||||
let config = HttpClientConfig::resolve(&settings, None).unwrap();
|
||||
assert_eq!(config.verify, Verify::Disabled);
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ssl_context_global_declines_instead_of_being_dropped() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let error = settings(&globals(py, "globals.ssl_verify = object()")).unwrap_err();
|
||||
assert!(error.is_instance_of::<RustBridgeDeclined>(py));
|
||||
assert!(error.value(py).to_string().contains("litellm.ssl_verify"));
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::client("client")]
|
||||
#[case::shared_session("shared_session")]
|
||||
#[case::aclient_session("aclient_session")]
|
||||
fn live_python_clients_decline_before_dispatch(#[case] name: &str) {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let kwargs = PyDict::new(py);
|
||||
kwargs
|
||||
.set_item(name, py.eval(c"object()", None, None).unwrap())
|
||||
.unwrap();
|
||||
let error = decline_live_clients(&kwargs).unwrap_err();
|
||||
assert!(error.is_instance_of::<RustBridgeDeclined>(py));
|
||||
assert!(error.value(py).to_string().contains(name));
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn none_valued_client_arguments_are_not_live_clients() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let kwargs = PyDict::new(py);
|
||||
for name in LIVE_CLIENT_ARGUMENTS {
|
||||
kwargs.set_item(name, py.None()).unwrap();
|
||||
}
|
||||
decline_live_clients(&kwargs).unwrap();
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::disabled(c"False", Some(SslVerify::Disabled))]
|
||||
#[case::enabled(c"True", Some(SslVerify::Enabled))]
|
||||
#[case::bundle(c"'/tmp/ca.pem'", Some(SslVerify::CaBundle("/tmp/ca.pem".into())))]
|
||||
#[case::unset(c"None", None)]
|
||||
fn per_call_ssl_verify_values_project(
|
||||
#[case] source: &std::ffi::CStr,
|
||||
#[case] expected: Option<SslVerify>,
|
||||
) {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let value = py.eval(source, None, None).unwrap();
|
||||
assert_eq!(ssl_verify(&value, "ssl_verify").unwrap(), expected);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
@ -1,6 +1,7 @@
|
|||
mod credentials;
|
||||
mod diagnostics;
|
||||
mod errors;
|
||||
mod http;
|
||||
mod marshal;
|
||||
mod routes;
|
||||
mod token_counter;
|
||||
|
|
|
|||
|
|
@ -12,6 +12,8 @@ use pyo3::{
|
|||
types::{PyDict, PyTuple},
|
||||
};
|
||||
|
||||
use crate::{errors::RustBridgeDeclined, http};
|
||||
|
||||
const SURFACE: LegacySurface = LegacySurface {
|
||||
call_type: "ocr",
|
||||
input_description: "OCR document processing",
|
||||
|
|
@ -29,7 +31,9 @@ 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)?;
|
||||
let client = OcrClient::new(http::pool(), &config)
|
||||
.map_err(|error| RustBridgeDeclined::new_err(error.to_string()))?;
|
||||
run_legacy_call(
|
||||
py,
|
||||
if asynchronous { ASYNC_SURFACE } else { SURFACE },
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue