mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
feat(rust): add litellm-s3 crate, AwsCredentialsProvider and cost map fetch (#50547)
* feat(rust): add litellm-s3 crate, AwsCredentialsProvider and cost map fetch * fix(model-catalog-refresh): time each read instead of the whole cost map download The 5 s fetch timeout covered the full body, so a slow but steady link failed every attempt where httpx and boto3 would finish. Each read now gets its own 5 s, matching main. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
yujonglee
GitHub
parent
1432e233cf
commit
25adca8510
25 changed files with 1358 additions and 270 deletions
30
litellm-rust/Cargo.lock
generated
30
litellm-rust/Cargo.lock
generated
|
|
@ -894,6 +894,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
|||
checksum = "cffb0e931875b666fc4fcb20fee52e9bbd1ef836fd9e9e04ec21555f9f85f7ef"
|
||||
dependencies = [
|
||||
"fastrand",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -3657,19 +3658,15 @@ dependencies = [
|
|||
name = "litellm-cache-s3"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"aws-credential-types",
|
||||
"aws-sdk-s3",
|
||||
"aws-smithy-runtime-api",
|
||||
"aws-smithy-types",
|
||||
"aws-types",
|
||||
"futures-util",
|
||||
"http 1.4.2",
|
||||
"litellm-auth-aws",
|
||||
"litellm-cache",
|
||||
"litellm-cache-testing",
|
||||
"litellm-error",
|
||||
"litellm-http",
|
||||
"reqwest",
|
||||
"litellm-s3",
|
||||
"rstest",
|
||||
"serde_json",
|
||||
"tokio",
|
||||
|
|
@ -4361,9 +4358,13 @@ dependencies = [
|
|||
name = "litellm-model-catalog-refresh"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"aws-credential-types",
|
||||
"backon",
|
||||
"bytes",
|
||||
"litellm-error",
|
||||
"litellm-http",
|
||||
"litellm-model-catalog",
|
||||
"litellm-s3",
|
||||
"litellm-tracing",
|
||||
"reqwest",
|
||||
"rstest",
|
||||
"thiserror 2.0.19",
|
||||
|
|
@ -4609,6 +4610,23 @@ dependencies = [
|
|||
"serde_with",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-s3"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"aws-config",
|
||||
"aws-credential-types",
|
||||
"aws-sdk-s3",
|
||||
"aws-smithy-runtime-api",
|
||||
"aws-smithy-types",
|
||||
"aws-types",
|
||||
"http 1.4.2",
|
||||
"litellm-http",
|
||||
"reqwest",
|
||||
"rstest",
|
||||
"thiserror 2.0.19",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-secret-value"
|
||||
version = "0.1.0"
|
||||
|
|
|
|||
|
|
@ -60,6 +60,7 @@ litellm-cache = { path = "crates/cache" }
|
|||
litellm-cache-azure-blob = { path = "crates/cache-azure-blob" }
|
||||
litellm-cache-memory = { path = "crates/cache-memory" }
|
||||
litellm-cache-s3 = { path = "crates/cache-s3" }
|
||||
litellm-s3 = { path = "crates/s3" }
|
||||
litellm-cache-gcs = { path = "crates/cache-gcs" }
|
||||
litellm-cache-disk = { path = "crates/cache-disk" }
|
||||
litellm-cache-response = { path = "crates/cache-response" }
|
||||
|
|
@ -90,6 +91,8 @@ axum = { version = "0.8.9", default-features = false, features = ["http1", "toki
|
|||
axum-login = "0.18.0"
|
||||
tower-sessions = { version = "0.14.0", features = ["memory-store"] }
|
||||
bytes = "1"
|
||||
backon = { version = "1.6.0", default-features = false, features = ["std", "tokio-sleep"] }
|
||||
aws-sdk-s3 = { version = "1.146.1", default-features = false, features = ["rt-tokio"] }
|
||||
http = "1"
|
||||
google-cloud-auth = { version = "1.16.0", default-features = false }
|
||||
jsonwebtoken = { version = "11.1.0", default-features = false }
|
||||
|
|
|
|||
|
|
@ -1,9 +1,11 @@
|
|||
mod aws;
|
||||
pub mod constants;
|
||||
mod error;
|
||||
mod provider;
|
||||
mod signer;
|
||||
|
||||
pub use aws::*;
|
||||
pub use aws_credential_types::Credentials;
|
||||
pub use error::{CredentialFlow, Error, SigningStage};
|
||||
pub use provider::AwsCredentialsProvider;
|
||||
pub use signer::SigV4Signer;
|
||||
|
|
|
|||
135
litellm-rust/crates/auth-aws/src/provider.rs
Normal file
135
litellm-rust/crates/auth-aws/src/provider.rs
Normal file
|
|
@ -0,0 +1,135 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use aws_credential_types::Credentials;
|
||||
use aws_credential_types::provider::{ProvideCredentials, error::CredentialsError, future};
|
||||
|
||||
use crate::{AwsAuthConfig, AwsAuthService};
|
||||
|
||||
type Environment = Arc<dyn Fn(&str) -> Option<String> + Send + Sync>;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct AwsCredentialsProvider {
|
||||
auth: AwsAuthService,
|
||||
config: Box<AwsAuthConfig>,
|
||||
environment: Environment,
|
||||
}
|
||||
|
||||
impl AwsCredentialsProvider {
|
||||
/// Explicit keys in `config` as given, otherwise LiteLLM's `AWS_*` settings read through `environment`
|
||||
pub fn configured(
|
||||
auth: AwsAuthService,
|
||||
config: AwsAuthConfig,
|
||||
environment: impl Fn(&str) -> Option<String> + Send + Sync + 'static,
|
||||
) -> Self {
|
||||
Self {
|
||||
auth,
|
||||
config: Box::new(config),
|
||||
environment: Arc::new(environment),
|
||||
}
|
||||
}
|
||||
|
||||
async fn credentials(&self) -> Result<Credentials, crate::Error> {
|
||||
match (&self.config.access_key_id, &self.config.secret_access_key) {
|
||||
(Some(access_key_id), Some(secret_access_key)) => Ok(Credentials::new(
|
||||
access_key_id.expose(),
|
||||
secret_access_key.expose(),
|
||||
self.config
|
||||
.session_token
|
||||
.as_ref()
|
||||
.map(|token| token.expose().to_owned()),
|
||||
None,
|
||||
"litellm-configured",
|
||||
)),
|
||||
_ => {
|
||||
self.auth
|
||||
.resolve_credentials(
|
||||
AwsAuthConfig::clone(&self.config),
|
||||
self.environment.as_ref(),
|
||||
)
|
||||
.await
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ProvideCredentials for AwsCredentialsProvider {
|
||||
fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a>
|
||||
where
|
||||
Self: 'a,
|
||||
{
|
||||
future::ProvideCredentials::new(async {
|
||||
self.credentials()
|
||||
.await
|
||||
.map_err(CredentialsError::provider_error)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for AwsCredentialsProvider {
|
||||
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
formatter
|
||||
.debug_struct("AwsCredentialsProvider")
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn explicit(session_token: Option<&str>) -> AwsAuthConfig {
|
||||
AwsAuthConfig {
|
||||
access_key_id: Some("key".into()),
|
||||
secret_access_key: Some("secret".into()),
|
||||
session_token: session_token.map(Into::into),
|
||||
region_name: Some("us-east-1".to_string()),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn ambient_token(name: &str) -> Option<String> {
|
||||
(name == "AWS_SESSION_TOKEN").then(|| "ambient".to_string())
|
||||
}
|
||||
|
||||
fn environment_keys(name: &str) -> Option<String> {
|
||||
match name {
|
||||
"AWS_ACCESS_KEY_ID" => Some("env-key".to_string()),
|
||||
"AWS_SECRET_ACCESS_KEY" => Some("env-secret".to_string()),
|
||||
"AWS_SESSION_TOKEN" => Some("env-token".to_string()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::explicit_keys_ignore_an_ambient_session_token(
|
||||
explicit(None), ambient_token, ("key", "secret", None)
|
||||
)]
|
||||
#[case::explicit_keys_keep_their_session_token(
|
||||
explicit(Some("t")), ambient_token, ("key", "secret", Some("t"))
|
||||
)]
|
||||
#[case::environment_keys_resolve_with_their_session_token(
|
||||
AwsAuthConfig::default(), environment_keys, ("env-key", "env-secret", Some("env-token"))
|
||||
)]
|
||||
#[tokio::test]
|
||||
async fn configured_credentials_read_the_environment_they_are_given(
|
||||
#[case] config: AwsAuthConfig,
|
||||
#[case] environment: fn(&str) -> Option<String>,
|
||||
#[case] expected: (&str, &str, Option<&str>),
|
||||
) {
|
||||
let credentials =
|
||||
AwsCredentialsProvider::configured(AwsAuthService::default(), config, environment)
|
||||
.provide_credentials()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
(
|
||||
credentials.access_key_id(),
|
||||
credentials.secret_access_key(),
|
||||
credentials.session_token(),
|
||||
),
|
||||
expected
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
@ -9,14 +9,10 @@ repository.workspace = true
|
|||
litellm-http.workspace = true
|
||||
litellm-cache.workspace = true
|
||||
litellm-auth-aws.workspace = true
|
||||
aws-sdk-s3 = { version = "1.146.1", default-features = false, features = ["rustls", "rt-tokio"] }
|
||||
aws-credential-types = "1.3.0"
|
||||
aws-smithy-runtime-api = { version = "1.16.2", features = ["client", "http-1x"] }
|
||||
aws-smithy-types = { version = "1.6.0", features = ["http-body-1-x"] }
|
||||
aws-types = "1.6.0"
|
||||
litellm-s3.workspace = true
|
||||
aws-sdk-s3.workspace = true
|
||||
aws-smithy-types = "1.6.0"
|
||||
futures-util.workspace = true
|
||||
http.workspace = true
|
||||
reqwest.workspace = true
|
||||
tokio.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
|
|
|
|||
|
|
@ -1,61 +0,0 @@
|
|||
use aws_credential_types::{
|
||||
Credentials as AwsCredentials,
|
||||
provider::{ProvideCredentials, error::CredentialsError, future},
|
||||
};
|
||||
use litellm_auth_aws::{AwsAuthConfig, AwsAuthService};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct S3Credentials {
|
||||
auth: AwsAuthService,
|
||||
config: AwsAuthConfig,
|
||||
env: fn(&str) -> Option<String>,
|
||||
}
|
||||
|
||||
impl S3Credentials {
|
||||
pub fn new(config: AwsAuthConfig) -> Self {
|
||||
Self::with_env(config, |name| std::env::var(name).ok())
|
||||
}
|
||||
|
||||
pub fn with_env(config: AwsAuthConfig, env: fn(&str) -> Option<String>) -> Self {
|
||||
Self {
|
||||
auth: AwsAuthService::default(),
|
||||
config,
|
||||
env,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ProvideCredentials for S3Credentials {
|
||||
fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a>
|
||||
where
|
||||
Self: 'a,
|
||||
{
|
||||
future::ProvideCredentials::new(async {
|
||||
if let (Some(access_key_id), Some(secret_access_key)) = (
|
||||
self.config.access_key_id.clone(),
|
||||
self.config.secret_access_key.clone(),
|
||||
) {
|
||||
return Ok(AwsCredentials::new(
|
||||
access_key_id.expose(),
|
||||
secret_access_key.expose(),
|
||||
self.config
|
||||
.session_token
|
||||
.as_ref()
|
||||
.map(|value| value.expose().to_owned()),
|
||||
None,
|
||||
"litellm-s3-cache",
|
||||
));
|
||||
}
|
||||
self.auth
|
||||
.resolve_credentials(self.config.clone(), &self.env)
|
||||
.await
|
||||
.map_err(CredentialsError::provider_error)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for S3Credentials {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("S3Credentials").finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
|
@ -5,20 +5,18 @@ use std::{
|
|||
};
|
||||
|
||||
use aws_sdk_s3::{
|
||||
config::{BehaviorVersion, Region, RequestChecksumCalculation, ResponseChecksumValidation},
|
||||
config::{RequestChecksumCalculation, ResponseChecksumValidation},
|
||||
error::SdkError,
|
||||
primitives::ByteStream,
|
||||
};
|
||||
use aws_smithy_types::{DateTime, date_time::Format};
|
||||
use futures_util::future::try_join_all;
|
||||
use litellm_auth_aws::AwsAuthConfig;
|
||||
use litellm_auth_aws::AwsCredentialsProvider;
|
||||
use litellm_cache::{
|
||||
BaseCache, BatchCache, CacheCodec, DisconnectCache, Error, ExactCacheContext, FlushCache,
|
||||
};
|
||||
use tokio::runtime::Handle;
|
||||
|
||||
use crate::{auth::S3Credentials, transport::ReqwestHttpClient};
|
||||
|
||||
pub struct S3Endpoint {
|
||||
pub url: String,
|
||||
}
|
||||
|
|
@ -28,7 +26,7 @@ pub struct S3CacheConfig {
|
|||
pub key_prefix: String,
|
||||
pub region: String,
|
||||
pub endpoint: Option<S3Endpoint>,
|
||||
pub auth: AwsAuthConfig,
|
||||
pub credentials: AwsCredentialsProvider,
|
||||
}
|
||||
|
||||
pub struct S3Cache<C: CacheCodec> {
|
||||
|
|
@ -49,17 +47,14 @@ impl<C: CacheCodec> S3Cache<C> {
|
|||
runtime: Handle,
|
||||
) -> Self {
|
||||
let endpoint_url: Option<String> = config.endpoint.map(|endpoint| endpoint.url);
|
||||
let base = aws_sdk_s3::Config::builder()
|
||||
.behavior_version(BehaviorVersion::latest())
|
||||
.region(Region::new(config.region.clone()))
|
||||
.http_client(ReqwestHttpClient(http))
|
||||
.credentials_provider(S3Credentials::new(config.auth))
|
||||
.request_checksum_calculation(RequestChecksumCalculation::WhenRequired)
|
||||
.response_checksum_validation(ResponseChecksumValidation::WhenRequired);
|
||||
let builder = match &endpoint_url {
|
||||
Some(url) => base.endpoint_url(url).force_path_style(true),
|
||||
None => base,
|
||||
};
|
||||
let builder = litellm_s3::config(
|
||||
config.region.clone(),
|
||||
endpoint_url.as_deref(),
|
||||
http,
|
||||
config.credentials,
|
||||
)
|
||||
.request_checksum_calculation(RequestChecksumCalculation::WhenRequired)
|
||||
.response_checksum_validation(ResponseChecksumValidation::WhenRequired);
|
||||
Self {
|
||||
client: aws_sdk_s3::Client::from_conf(builder.build()),
|
||||
codec,
|
||||
|
|
|
|||
|
|
@ -1,6 +1,3 @@
|
|||
mod auth;
|
||||
mod cache;
|
||||
mod transport;
|
||||
|
||||
pub use auth::S3Credentials;
|
||||
pub use cache::{S3Cache, S3CacheConfig, S3Endpoint};
|
||||
|
|
|
|||
|
|
@ -1,57 +0,0 @@
|
|||
use aws_credential_types::provider::ProvideCredentials;
|
||||
use litellm_auth_aws::AwsAuthConfig;
|
||||
use litellm_cache_s3::S3Credentials;
|
||||
use rstest::rstest;
|
||||
|
||||
fn explicit(session_token: Option<&str>) -> AwsAuthConfig {
|
||||
AwsAuthConfig {
|
||||
access_key_id: Some("key".into()),
|
||||
secret_access_key: Some("secret".into()),
|
||||
session_token: session_token.map(Into::into),
|
||||
region_name: Some("us-east-1".to_string()),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn ambient_token(name: &str) -> Option<String> {
|
||||
(name == "AWS_SESSION_TOKEN").then(|| "ambient".to_string())
|
||||
}
|
||||
|
||||
fn environment_keys(name: &str) -> Option<String> {
|
||||
match name {
|
||||
"AWS_ACCESS_KEY_ID" => Some("env-key".to_string()),
|
||||
"AWS_SECRET_ACCESS_KEY" => Some("env-secret".to_string()),
|
||||
"AWS_SESSION_TOKEN" => Some("env-token".to_string()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::explicit_keys_ignore_an_ambient_session_token(
|
||||
explicit(None), ambient_token, ("key", "secret", None)
|
||||
)]
|
||||
#[case::explicit_keys_keep_their_session_token(
|
||||
explicit(Some("t")), ambient_token, ("key", "secret", Some("t"))
|
||||
)]
|
||||
#[case::environment_keys_resolve_with_their_session_token(
|
||||
AwsAuthConfig::default(), environment_keys, ("env-key", "env-secret", Some("env-token"))
|
||||
)]
|
||||
#[tokio::test]
|
||||
async fn credentials_resolve(
|
||||
#[case] config: AwsAuthConfig,
|
||||
#[case] env: fn(&str) -> Option<String>,
|
||||
#[case] expected: (&str, &str, Option<&str>),
|
||||
) {
|
||||
let credentials = S3Credentials::with_env(config, env)
|
||||
.provide_credentials()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
(
|
||||
credentials.access_key_id(),
|
||||
credentials.secret_access_key(),
|
||||
credentials.session_token(),
|
||||
),
|
||||
expected
|
||||
);
|
||||
}
|
||||
|
|
@ -5,7 +5,7 @@ use std::{
|
|||
sync::{Arc, Mutex},
|
||||
};
|
||||
|
||||
use litellm_auth_aws::AwsAuthConfig;
|
||||
use litellm_auth_aws::{AwsAuthConfig, AwsAuthService, AwsCredentialsProvider};
|
||||
use litellm_cache::JsonCodec;
|
||||
use litellm_cache_s3::{S3Cache, S3CacheConfig, S3Endpoint};
|
||||
use serde_json::Value;
|
||||
|
|
@ -22,12 +22,16 @@ pub fn config(endpoint: &str) -> S3CacheConfig {
|
|||
endpoint: Some(S3Endpoint {
|
||||
url: endpoint.to_string(),
|
||||
}),
|
||||
auth: AwsAuthConfig {
|
||||
access_key_id: Some("key".into()),
|
||||
secret_access_key: Some("secret".into()),
|
||||
region_name: Some("us-east-1".to_string()),
|
||||
..Default::default()
|
||||
},
|
||||
credentials: AwsCredentialsProvider::configured(
|
||||
AwsAuthService::default(),
|
||||
AwsAuthConfig {
|
||||
access_key_id: Some("key".into()),
|
||||
secret_access_key: Some("secret".into()),
|
||||
region_name: Some("us-east-1".to_string()),
|
||||
..Default::default()
|
||||
},
|
||||
|_| None,
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -6,14 +6,19 @@ license.workspace = true
|
|||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
backon.workspace = true
|
||||
bytes.workspace = true
|
||||
litellm-error.workspace = true
|
||||
litellm-http.workspace = true
|
||||
litellm-model-catalog.workspace = true
|
||||
litellm-s3.workspace = true
|
||||
litellm-tracing.workspace = true
|
||||
reqwest.workspace = true
|
||||
thiserror.workspace = true
|
||||
tokio = { workspace = true, features = ["time"] }
|
||||
url.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
litellm-error.workspace = true
|
||||
aws-credential-types = "1.3.0"
|
||||
litellm-http = { workspace = true, features = ["test-support"] }
|
||||
reqwest.workspace = true
|
||||
rstest.workspace = true
|
||||
tokio.workspace = true
|
||||
tokio = { workspace = true, features = ["macros", "rt-multi-thread", "time"] }
|
||||
|
|
|
|||
153
litellm-rust/crates/model-catalog-refresh/src/attempt.rs
Normal file
153
litellm-rust/crates/model-catalog-refresh/src/attempt.rs
Normal file
|
|
@ -0,0 +1,153 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use bytes::Bytes;
|
||||
|
||||
use litellm_http::Client;
|
||||
use litellm_s3::S3Uri;
|
||||
use litellm_s3::sdk::config::retry::RetryConfig;
|
||||
use litellm_s3::sdk::config::timeout::TimeoutConfig;
|
||||
use reqwest::header::{ETAG, HeaderMap, RETRY_AFTER};
|
||||
use url::Url;
|
||||
|
||||
use crate::{Error, RETRYABLE_STATUSES, S3Access};
|
||||
|
||||
pub(crate) const MAX_BODY_BYTES: u64 = 100 * 1024 * 1024;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct Fetched {
|
||||
pub body: Vec<u8>,
|
||||
pub etag: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub(crate) enum Cause {
|
||||
Status(u16),
|
||||
Transport(Error),
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct Rejected {
|
||||
pub(crate) cause: Cause,
|
||||
pub(crate) retryable: bool,
|
||||
pub(crate) retry_after: Option<Duration>,
|
||||
}
|
||||
|
||||
impl Rejected {
|
||||
pub(crate) fn fatal(error: Error) -> Self {
|
||||
Self {
|
||||
cause: Cause::Transport(error),
|
||||
retryable: false,
|
||||
retry_after: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Any status below 400 is the document
|
||||
pub(crate) async fn http(
|
||||
client: &Client,
|
||||
url: Url,
|
||||
timeout: Duration,
|
||||
) -> Result<Fetched, Rejected> {
|
||||
let transport = |error: litellm_http::Error| Rejected {
|
||||
cause: Cause::Transport(Error::from(error)),
|
||||
retryable: true,
|
||||
retry_after: None,
|
||||
};
|
||||
let mut response = tokio::time::timeout(timeout, client.get(url).send())
|
||||
.await
|
||||
.map_err(litellm_http::Error::read_timeout)
|
||||
.and_then(|sent| sent.map_err(litellm_http::Error::from))
|
||||
.map_err(transport)?;
|
||||
let status = response.status().as_u16();
|
||||
let etag = etag(response.headers());
|
||||
let retry_after = response
|
||||
.headers()
|
||||
.get(RETRY_AFTER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.and_then(|value| value.trim().parse().ok())
|
||||
.map(Duration::from_secs);
|
||||
if status >= 400 {
|
||||
return Err(Rejected {
|
||||
cause: Cause::Status(status),
|
||||
retryable: RETRYABLE_STATUSES.contains(&status),
|
||||
retry_after,
|
||||
});
|
||||
}
|
||||
if response
|
||||
.content_length()
|
||||
.is_some_and(|length| length > MAX_BODY_BYTES)
|
||||
{
|
||||
return Err(Rejected::fatal(Error::body_too_large(MAX_BODY_BYTES)));
|
||||
}
|
||||
let mut body = Vec::new();
|
||||
while let Some(chunk) = tokio::time::timeout(timeout, response.chunk())
|
||||
.await
|
||||
.map_err(litellm_http::Error::body_read_timeout)
|
||||
.and_then(|read| read.map_err(litellm_http::Error::from))
|
||||
.map_err(transport)?
|
||||
{
|
||||
append(&mut body, &chunk).map_err(Rejected::fatal)?;
|
||||
}
|
||||
Ok(Fetched { body, etag })
|
||||
}
|
||||
|
||||
/// S3 failures are final: the SDK already classifies its own errors and a missing or denied object stays so
|
||||
pub(crate) async fn s3(
|
||||
client: &Client,
|
||||
s3: &S3Access,
|
||||
object: &S3Uri,
|
||||
timeout: Duration,
|
||||
) -> Result<Fetched, Rejected> {
|
||||
let fetched = async {
|
||||
let config = match s3 {
|
||||
S3Access::Environment { region } => {
|
||||
litellm_s3::environment_config(client.clone(), region.clone()).await
|
||||
}
|
||||
S3Access::Config(config) => config.to_builder(),
|
||||
}
|
||||
.retry_config(RetryConfig::disabled())
|
||||
.timeout_config(
|
||||
TimeoutConfig::builder()
|
||||
.operation_attempt_timeout(timeout)
|
||||
.build(),
|
||||
)
|
||||
.build();
|
||||
let output = litellm_s3::sdk::Client::from_conf(config)
|
||||
.get_object()
|
||||
.bucket(&object.bucket)
|
||||
.key(&object.key)
|
||||
.send()
|
||||
.await?;
|
||||
if output
|
||||
.content_length()
|
||||
.is_some_and(|length| length.unsigned_abs() > MAX_BODY_BYTES)
|
||||
{
|
||||
return Err(Error::body_too_large(MAX_BODY_BYTES));
|
||||
}
|
||||
let etag = output.e_tag().map(str::to_owned);
|
||||
let mut stream = output.body;
|
||||
let mut body = Vec::new();
|
||||
while let Some(chunk) = tokio::time::timeout(timeout, stream.try_next()).await?? {
|
||||
append(&mut body, &chunk)?;
|
||||
}
|
||||
Ok::<_, Error>(Fetched { body, etag })
|
||||
};
|
||||
fetched.await.map_err(Rejected::fatal)
|
||||
}
|
||||
|
||||
fn append(body: &mut Vec<u8>, chunk: &Bytes) -> Result<(), Error> {
|
||||
if (body.len() + chunk.len()) as u64 > MAX_BODY_BYTES {
|
||||
return Err(Error::body_too_large(MAX_BODY_BYTES));
|
||||
}
|
||||
body.extend_from_slice(chunk);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn etag(headers: &HeaderMap) -> Option<String> {
|
||||
let values: Vec<&str> = headers
|
||||
.get_all(ETAG)
|
||||
.iter()
|
||||
.filter_map(|value| value.to_str().ok())
|
||||
.collect();
|
||||
(!values.is_empty()).then(|| values.join(", "))
|
||||
}
|
||||
|
|
@ -1,6 +1,9 @@
|
|||
use std::io;
|
||||
|
||||
use litellm_s3::sdk::config::http::HttpResponse;
|
||||
use litellm_s3::sdk::error::SdkError;
|
||||
use litellm_s3::sdk::operation::get_object::GetObjectError;
|
||||
use litellm_s3::sdk::primitives::ByteStreamError;
|
||||
use thiserror::Error;
|
||||
use tokio::time::error::Elapsed;
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
#[error(transparent)]
|
||||
|
|
@ -10,18 +13,48 @@ pub struct Error(Box<Repr>);
|
|||
enum Repr {
|
||||
#[error(transparent)]
|
||||
Http(#[from] litellm_http::Error),
|
||||
#[error("could not read the catalog")]
|
||||
Read(#[source] io::Error),
|
||||
#[error("invalid catalog URL")]
|
||||
InvalidUrl(#[source] url::ParseError),
|
||||
#[error("unsupported catalog URL scheme {scheme}")]
|
||||
UnsupportedScheme { scheme: String },
|
||||
#[error("invalid S3 URL")]
|
||||
S3Uri(#[source] litellm_s3::Error),
|
||||
#[error("S3 GetObject failed")]
|
||||
S3GetObject(#[source] Box<SdkError<GetObjectError, HttpResponse>>),
|
||||
#[error("S3 GetObject body read failed")]
|
||||
S3Body(#[source] ByteStreamError),
|
||||
#[error("S3 GetObject timed out")]
|
||||
S3Timeout(#[source] Elapsed),
|
||||
#[error("response body is over the {limit} byte limit")]
|
||||
BodyTooLarge { limit: u64 },
|
||||
}
|
||||
|
||||
impl Error {
|
||||
pub fn unsupported_scheme(scheme: String) -> Self {
|
||||
Self(Box::new(Repr::UnsupportedScheme { scheme }))
|
||||
}
|
||||
|
||||
pub(crate) fn body_too_large(limit: u64) -> Self {
|
||||
Self(Box::new(Repr::BodyTooLarge { limit }))
|
||||
}
|
||||
}
|
||||
|
||||
impl From<SdkError<GetObjectError, HttpResponse>> for Error {
|
||||
fn from(cause: SdkError<GetObjectError, HttpResponse>) -> Self {
|
||||
Self(Box::new(Repr::S3GetObject(Box::new(cause))))
|
||||
}
|
||||
}
|
||||
|
||||
impl From<ByteStreamError> for Error {
|
||||
fn from(cause: ByteStreamError) -> Self {
|
||||
Self(Box::new(Repr::S3Body(cause)))
|
||||
}
|
||||
}
|
||||
|
||||
impl From<Elapsed> for Error {
|
||||
fn from(cause: Elapsed) -> Self {
|
||||
Self(Box::new(Repr::S3Timeout(cause)))
|
||||
}
|
||||
}
|
||||
|
||||
impl From<litellm_http::Error> for Error {
|
||||
|
|
@ -30,9 +63,9 @@ impl From<litellm_http::Error> for Error {
|
|||
}
|
||||
}
|
||||
|
||||
impl From<io::Error> for Error {
|
||||
fn from(cause: io::Error) -> Self {
|
||||
Self(Box::new(Repr::Read(cause)))
|
||||
impl From<litellm_s3::Error> for Error {
|
||||
fn from(cause: litellm_s3::Error) -> Self {
|
||||
Self(Box::new(Repr::S3Uri(cause)))
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -44,29 +77,18 @@ impl From<url::ParseError> for Error {
|
|||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::io;
|
||||
|
||||
use litellm_error::{Report, find};
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
use crate::{DoneReason, Reply, Step};
|
||||
|
||||
#[rstest]
|
||||
fn exhausted_transport_keeps_its_typed_cause() {
|
||||
let reply = Reply::Transport {
|
||||
error: io::Error::from(io::ErrorKind::ConnectionReset).into(),
|
||||
};
|
||||
let Reply::Transport { error } = reply else {
|
||||
panic!("transport failure")
|
||||
};
|
||||
let step = Step::Done {
|
||||
reason: DoneReason::Exhausted(error),
|
||||
};
|
||||
let Step::Done {
|
||||
reason: DoneReason::Exhausted(error),
|
||||
} = step
|
||||
else {
|
||||
panic!("exhausted attempts")
|
||||
};
|
||||
fn an_s3_failure_keeps_its_typed_cause_once() {
|
||||
let error = Error::from(SdkError::<GetObjectError, HttpResponse>::timeout_error(
|
||||
io::Error::from(io::ErrorKind::ConnectionReset),
|
||||
));
|
||||
let cause = find::<io::Error>(&error).expect("typed I/O cause");
|
||||
assert_eq!(cause.kind(), io::ErrorKind::ConnectionReset);
|
||||
assert_eq!(
|
||||
|
|
|
|||
122
litellm-rust/crates/model-catalog-refresh/src/fetch.rs
Normal file
122
litellm-rust/crates/model-catalog-refresh/src/fetch.rs
Normal file
|
|
@ -0,0 +1,122 @@
|
|||
use std::borrow::Cow;
|
||||
use std::sync::atomic::{AtomicU32, Ordering};
|
||||
use std::time::Duration;
|
||||
|
||||
use backon::Retryable;
|
||||
use litellm_error::Report;
|
||||
use litellm_http::Client;
|
||||
use url::Url;
|
||||
|
||||
use crate::attempt::{self, Cause, Fetched, Rejected};
|
||||
use crate::{RetryPolicy, Source};
|
||||
|
||||
/// How an `s3://` fetch reaches S3
|
||||
#[derive(Clone, Debug)]
|
||||
pub enum S3Access {
|
||||
/// What an AWS SDK finds in the environment and shared config files, loaded on each `s3://` attempt, with
|
||||
/// `region` winning over the region it finds
|
||||
Environment { region: Option<String> },
|
||||
/// A config the caller built with its region, endpoint and credentials
|
||||
Config(Box<litellm_s3::sdk::Config>),
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct Failed {
|
||||
pub reason: String,
|
||||
/// Another attempt could succeed. Only the attempt limit ended the fetch
|
||||
pub retryable: bool,
|
||||
pub retry_after: Option<Duration>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum Start {
|
||||
Fresh,
|
||||
/// Continues a fetch whose first attempt failed retryably elsewhere, waiting before its next attempt
|
||||
After(Failed),
|
||||
}
|
||||
|
||||
pub async fn fetch(
|
||||
client: &Client,
|
||||
s3: &S3Access,
|
||||
url: &str,
|
||||
policy: RetryPolicy,
|
||||
start: Start,
|
||||
) -> Result<Fetched, Failed> {
|
||||
let shown = redacted(url);
|
||||
let source =
|
||||
Source::parse(url).map_err(|error| Failed::of(&shown, &Rejected::fatal(error), 1))?;
|
||||
let mut backoff = policy.backoff();
|
||||
let made = AtomicU32::new(0);
|
||||
if let Start::After(failed) = start {
|
||||
let Some(wait) = backoff.next() else {
|
||||
return Err(failed);
|
||||
};
|
||||
tokio::time::sleep(policy.delay(wait, failed.retry_after)).await;
|
||||
made.store(1, Ordering::Relaxed);
|
||||
}
|
||||
let attempt = || {
|
||||
made.fetch_add(1, Ordering::Relaxed);
|
||||
async {
|
||||
match &source {
|
||||
Source::Http(http) => attempt::http(client, http.clone(), policy.timeout).await,
|
||||
Source::S3(object) => attempt::s3(client, s3, object, policy.timeout).await,
|
||||
}
|
||||
}
|
||||
};
|
||||
attempt
|
||||
.retry(backoff)
|
||||
.sleep(tokio::time::sleep)
|
||||
.when(|rejected| rejected.retryable)
|
||||
.adjust(|rejected, wait| wait.map(|wait| policy.delay(wait, rejected.retry_after)))
|
||||
.notify(|rejected, wait| {
|
||||
litellm_tracing::warn!(
|
||||
"LiteLLM: model cost map fetch attempt {}/{} failed ({}); retrying in {:.1}s",
|
||||
made.load(Ordering::Relaxed),
|
||||
policy.max_attempts,
|
||||
reason(&shown, &rejected.cause),
|
||||
wait.as_secs_f64(),
|
||||
);
|
||||
})
|
||||
.await
|
||||
.map_err(|rejected| Failed::of(&shown, &rejected, made.load(Ordering::Relaxed)))
|
||||
}
|
||||
|
||||
impl Failed {
|
||||
fn of(url: &str, rejected: &Rejected, attempts: u32) -> Self {
|
||||
let reason = reason(url, &rejected.cause);
|
||||
Self {
|
||||
reason: match rejected.retryable && attempts > 1 {
|
||||
true => format!("{reason} (after {attempts} attempts)"),
|
||||
false => reason,
|
||||
},
|
||||
retryable: rejected.retryable,
|
||||
retry_after: rejected.retry_after,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn reason(url: &str, cause: &Cause) -> String {
|
||||
match cause {
|
||||
Cause::Status(status) => format!("HTTP {status} from {url}"),
|
||||
Cause::Transport(error) => format!("error fetching {url}: {}", Report::new(error)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Drops userinfo, query and fragment, where presigned URLs and tokens carry credentials
|
||||
fn redacted(url: &str) -> Cow<'_, str> {
|
||||
let Ok(mut parsed) = Url::parse(url) else {
|
||||
return Cow::Borrowed(url);
|
||||
};
|
||||
if parsed.username().is_empty()
|
||||
&& parsed.password().is_none()
|
||||
&& parsed.query().is_none()
|
||||
&& parsed.fragment().is_none()
|
||||
{
|
||||
return Cow::Borrowed(url);
|
||||
}
|
||||
let _ = parsed.set_username("");
|
||||
let _ = parsed.set_password(None);
|
||||
parsed.set_query(None);
|
||||
parsed.set_fragment(None);
|
||||
Cow::Owned(parsed.into())
|
||||
}
|
||||
|
|
@ -1,9 +1,11 @@
|
|||
mod attempt;
|
||||
mod error;
|
||||
mod fetch;
|
||||
mod policy;
|
||||
mod source;
|
||||
mod step;
|
||||
|
||||
pub use attempt::Fetched;
|
||||
pub use error::Error;
|
||||
pub use fetch::*;
|
||||
pub use policy::*;
|
||||
pub use source::*;
|
||||
pub use step::*;
|
||||
|
|
|
|||
|
|
@ -1,20 +1,76 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_model_catalog::IntegrityLimits;
|
||||
use backon::{BackoffBuilder, ExponentialBackoff, ExponentialBuilder};
|
||||
|
||||
/// Main's defaults, 3 attempts, 30 s max wait, 5 s timeout, backoff base 2.0, jitter 0.0 to 1.0 s
|
||||
#[derive(Clone, Copy, Debug, PartialEq)]
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub struct RetryPolicy {
|
||||
pub max_attempts: u32,
|
||||
pub max_wait: Duration,
|
||||
pub timeout: Duration,
|
||||
pub backoff_base_seconds: f64,
|
||||
pub jitter_min_seconds: f64,
|
||||
pub jitter_max_seconds: f64,
|
||||
/// Main's 50 model floor and 0.5 shrink ratio
|
||||
pub limits: IntegrityLimits,
|
||||
pub min_delay: Duration,
|
||||
pub max_delay: Duration,
|
||||
}
|
||||
|
||||
pub const RETRYABLE_STATUSES: [u16; 5] = [429, 500, 502, 503, 504];
|
||||
|
||||
/// Main's default `LITELLM_MODEL_COST_MAP_URL`
|
||||
pub const DEFAULT_REGISTRY_URL: &str =
|
||||
"https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json";
|
||||
|
||||
impl RetryPolicy {
|
||||
pub(crate) fn backoff(&self) -> ExponentialBackoff {
|
||||
ExponentialBuilder::new()
|
||||
.with_min_delay(self.min_delay)
|
||||
.with_max_delay(self.max_delay)
|
||||
.with_max_times(self.max_attempts.saturating_sub(1) as usize)
|
||||
.with_jitter()
|
||||
.build()
|
||||
}
|
||||
|
||||
pub(crate) fn delay(&self, backoff: Duration, retry_after: Option<Duration>) -> Duration {
|
||||
retry_after.map_or(backoff, |server| server.min(self.max_delay))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
|
||||
const POLICY: RetryPolicy = RetryPolicy {
|
||||
max_attempts: 3,
|
||||
timeout: Duration::from_secs(5),
|
||||
min_delay: Duration::from_secs(1),
|
||||
max_delay: Duration::from_secs(30),
|
||||
};
|
||||
|
||||
#[rstest]
|
||||
#[case::backoff_without_retry_after(None, 4)]
|
||||
#[case::retry_after_wins(Some(7), 7)]
|
||||
#[case::retry_after_zero(Some(0), 0)]
|
||||
#[case::retry_after_capped(Some(120), 30)]
|
||||
fn the_server_retry_after_overrides_the_backoff_up_to_max_delay(
|
||||
#[case] retry_after: Option<u64>,
|
||||
#[case] seconds: u64,
|
||||
) {
|
||||
assert_eq!(
|
||||
POLICY.delay(Duration::from_secs(4), retry_after.map(Duration::from_secs)),
|
||||
Duration::from_secs(seconds)
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::one_attempt(1, 0)]
|
||||
#[case::three_attempts(3, 2)]
|
||||
#[case::zero_attempts(0, 0)]
|
||||
fn the_backoff_allows_one_retry_fewer_than_the_attempts(
|
||||
#[case] max_attempts: u32,
|
||||
#[case] retries: usize,
|
||||
) {
|
||||
let policy = RetryPolicy {
|
||||
max_attempts,
|
||||
..POLICY
|
||||
};
|
||||
assert_eq!(policy.backoff().count(), retries);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,17 +1,59 @@
|
|||
use std::path::PathBuf;
|
||||
|
||||
use litellm_s3::S3Uri;
|
||||
use url::Url;
|
||||
|
||||
use crate::Error;
|
||||
|
||||
/// Where a fetched catalog comes from
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum Source {
|
||||
Http(Url),
|
||||
File(PathBuf),
|
||||
S3(S3Object),
|
||||
S3(S3Uri),
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct S3Object {
|
||||
pub bucket: String,
|
||||
pub key: String,
|
||||
impl Source {
|
||||
pub fn parse(url: &str) -> Result<Self, Error> {
|
||||
let parsed = Url::parse(url)?;
|
||||
match parsed.scheme() {
|
||||
"http" | "https" => Ok(Self::Http(parsed)),
|
||||
"s3" => Ok(Self::S3(S3Uri::try_from(url)?)),
|
||||
scheme => Err(Error::unsupported_scheme(scheme.to_owned())),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[rstest]
|
||||
#[case::https("https://example.com/cost.json")]
|
||||
#[case::http("http://127.0.0.1:8000/cost.json")]
|
||||
fn http_urls_are_fetched_over_http(#[case] url: &str) {
|
||||
assert_eq!(
|
||||
Source::parse(url).unwrap(),
|
||||
Source::Http(Url::parse(url).unwrap())
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn s3_urls_name_a_bucket_and_key() {
|
||||
assert_eq!(
|
||||
Source::parse("s3://bucket/prices/cost.json").unwrap(),
|
||||
Source::S3(S3Uri {
|
||||
bucket: "bucket".into(),
|
||||
key: "prices/cost.json".into(),
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::s3_without_a_key("s3://bucket/")]
|
||||
#[case::unsupported("ftp://example.com/cost.json")]
|
||||
#[case::file("file:///tmp/cost.json")]
|
||||
#[case::not_a_url("not a url")]
|
||||
fn other_urls_are_refused(#[case] url: &str) {
|
||||
assert!(Source::parse(url).is_err());
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,66 +0,0 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_model_catalog::Catalog;
|
||||
|
||||
use crate::Error;
|
||||
use crate::source::Source;
|
||||
|
||||
/// One attempt. Where to fetch from is the `Source`, so the request carries no URL
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct FetchRequest {
|
||||
pub if_none_match: Option<String>,
|
||||
pub timeout: Duration,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum Reply {
|
||||
/// An HTTP, S3 or file answer. The body is read only on 200
|
||||
Response {
|
||||
status: u16,
|
||||
etag: Option<String>,
|
||||
retry_after: Option<Duration>,
|
||||
body: Option<Vec<u8>>,
|
||||
},
|
||||
/// No answer reached us
|
||||
Transport { error: Error },
|
||||
}
|
||||
|
||||
/// What the fetch state machine asks its driver to do next
|
||||
#[derive(Debug)]
|
||||
pub enum Step {
|
||||
Fetch(FetchRequest),
|
||||
Sleep(Duration),
|
||||
/// A parsed document to install, with its rejected row errors
|
||||
Adopt {
|
||||
catalog: Box<Catalog>,
|
||||
installed: Installed,
|
||||
errors: Vec<litellm_model_catalog::Error>,
|
||||
},
|
||||
/// 304, a final status, or retries exhausted
|
||||
Done {
|
||||
reason: DoneReason,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum DoneReason {
|
||||
NotModified,
|
||||
Status(u16),
|
||||
Exhausted(Error),
|
||||
}
|
||||
|
||||
/// What is installed now, the input to the next revalidation
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct Installed {
|
||||
pub origin: Origin,
|
||||
pub etag: Option<String>,
|
||||
/// `git_blob_id` of the installed bytes
|
||||
pub revision: String,
|
||||
pub model_count: usize,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum Origin {
|
||||
Bundled,
|
||||
Remote { source: Source },
|
||||
}
|
||||
553
litellm-rust/crates/model-catalog-refresh/tests/fetch.rs
Normal file
553
litellm-rust/crates/model-catalog-refresh/tests/fetch.rs
Normal file
|
|
@ -0,0 +1,553 @@
|
|||
use std::io::{BufRead, BufReader, Write};
|
||||
use std::net::TcpListener;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Duration;
|
||||
|
||||
use aws_credential_types::Credentials;
|
||||
use aws_credential_types::provider::error::CredentialsError;
|
||||
use aws_credential_types::provider::{ProvideCredentials, future};
|
||||
use litellm_http::Client;
|
||||
use litellm_model_catalog_refresh::{Failed, Fetched, RetryPolicy, S3Access, Start, fetch};
|
||||
use rstest::rstest;
|
||||
|
||||
const POLICY: RetryPolicy = RetryPolicy {
|
||||
max_attempts: 3,
|
||||
timeout: Duration::from_secs(5),
|
||||
min_delay: Duration::from_millis(10),
|
||||
max_delay: Duration::from_millis(50),
|
||||
};
|
||||
|
||||
const OK: &str =
|
||||
"HTTP/1.1 200 OK\r\nETag: \"v1\"\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{}";
|
||||
const UNAVAILABLE: &str = "HTTP/1.1 503 Service Unavailable\r\nRetry-After: 0\r\nContent-Length: 0\r\nConnection: close\r\n\r\n";
|
||||
const NOT_FOUND: &str = "HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\nConnection: close\r\n\r\n";
|
||||
|
||||
const SLOW_BODY: &[u8] = b"{ }";
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum Answer {
|
||||
Raw(&'static str),
|
||||
Stall,
|
||||
StallBody,
|
||||
SlowBody,
|
||||
Oversized,
|
||||
}
|
||||
|
||||
/// A scripted upstream: the n-th request gets the n-th answer, repeating the last, and every request head is kept
|
||||
struct Upstream {
|
||||
base: String,
|
||||
heads: Arc<Mutex<Vec<String>>>,
|
||||
}
|
||||
|
||||
impl Upstream {
|
||||
fn serve(answers: Vec<Answer>) -> Self {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
|
||||
let base = format!("http://{}", listener.local_addr().unwrap());
|
||||
let heads = Arc::new(Mutex::new(Vec::new()));
|
||||
let recorded = Arc::clone(&heads);
|
||||
std::thread::spawn(move || {
|
||||
for (index, stream) in listener.incoming().enumerate() {
|
||||
let answer = answers[index.min(answers.len() - 1)];
|
||||
let recorded = Arc::clone(&recorded);
|
||||
std::thread::spawn(move || {
|
||||
let mut stream = stream.unwrap();
|
||||
let mut reader = BufReader::new(stream.try_clone().unwrap());
|
||||
let mut head = String::new();
|
||||
loop {
|
||||
let mut line = String::new();
|
||||
if reader.read_line(&mut line).unwrap() <= 2 {
|
||||
break;
|
||||
}
|
||||
head.push_str(&line);
|
||||
}
|
||||
recorded.lock().unwrap().push(head);
|
||||
match answer {
|
||||
Answer::Raw(raw) => stream.write_all(raw.as_bytes()).unwrap(),
|
||||
Answer::Stall => std::thread::sleep(Duration::from_secs(5)),
|
||||
Answer::Oversized => {
|
||||
let chunk = vec![b' '; 1024 * 1024];
|
||||
let written = stream
|
||||
.write_all(b"HTTP/1.1 200 OK\r\nConnection: close\r\n\r\n")
|
||||
.and_then(|()| {
|
||||
(0..=100).try_for_each(|_| stream.write_all(&chunk))
|
||||
});
|
||||
drop(written);
|
||||
}
|
||||
Answer::StallBody => {
|
||||
stream
|
||||
.write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\n{")
|
||||
.unwrap();
|
||||
std::thread::sleep(Duration::from_secs(5));
|
||||
}
|
||||
Answer::SlowBody => {
|
||||
stream.set_nodelay(true).unwrap();
|
||||
stream
|
||||
.write_all(
|
||||
b"HTTP/1.1 200 OK\r\nETag: \"v1\"\r\nContent-Length: 8\r\nConnection: close\r\n\r\n",
|
||||
)
|
||||
.unwrap();
|
||||
for byte in SLOW_BODY {
|
||||
std::thread::sleep(Duration::from_millis(60));
|
||||
stream.write_all(&[*byte]).unwrap();
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
});
|
||||
Self { base, heads }
|
||||
}
|
||||
|
||||
fn url(&self) -> String {
|
||||
format!("{}/cost.json", self.base)
|
||||
}
|
||||
|
||||
fn heads(&self) -> Vec<String> {
|
||||
self.heads.lock().unwrap().clone()
|
||||
}
|
||||
}
|
||||
|
||||
fn header<'a>(head: &'a str, name: &str) -> Option<&'a str> {
|
||||
head.lines().find_map(|line| {
|
||||
let (key, value) = line.split_once(':')?;
|
||||
key.eq_ignore_ascii_case(name).then(|| value.trim())
|
||||
})
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
struct NoCredentials;
|
||||
|
||||
impl ProvideCredentials for NoCredentials {
|
||||
fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a>
|
||||
where
|
||||
Self: 'a,
|
||||
{
|
||||
future::ProvideCredentials::ready(Err(CredentialsError::not_loaded("no credentials")))
|
||||
}
|
||||
}
|
||||
|
||||
fn keys() -> Credentials {
|
||||
Credentials::new(
|
||||
"AKIDMODELCOST",
|
||||
"model-cost-secret",
|
||||
None,
|
||||
None,
|
||||
"model-cost-test",
|
||||
)
|
||||
}
|
||||
|
||||
async fn fetched(url: &str, policy: RetryPolicy, start: Start) -> Result<Fetched, Failed> {
|
||||
fetch(
|
||||
&Client::for_test(),
|
||||
&S3Access::Environment { region: None },
|
||||
url,
|
||||
policy,
|
||||
start,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::one_etag(
|
||||
"HTTP/1.1 200 OK\r\nETag: \"v1\"\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{}",
|
||||
Some("\"v1\"")
|
||||
)]
|
||||
#[case::every_etag_joined(
|
||||
"HTTP/1.1 200 OK\r\nETag: \"a\"\r\nETag: \"b\"\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{}",
|
||||
Some("\"a\", \"b\"")
|
||||
)]
|
||||
#[case::no_etag(
|
||||
"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{}",
|
||||
None
|
||||
)]
|
||||
#[tokio::test]
|
||||
async fn a_200_returns_the_body_and_its_etags(
|
||||
#[case] response: &'static str,
|
||||
#[case] etag: Option<&str>,
|
||||
) {
|
||||
let upstream = Upstream::serve(vec![Answer::Raw(response)]);
|
||||
assert_eq!(
|
||||
fetched(&upstream.url(), POLICY, Start::Fresh).await,
|
||||
Ok(Fetched {
|
||||
body: b"{}".to_vec(),
|
||||
etag: etag.map(str::to_owned),
|
||||
})
|
||||
);
|
||||
assert_eq!(upstream.heads().len(), 1);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test(flavor = "multi_thread")]
|
||||
async fn a_fetch_runs_as_a_task_on_a_multi_thread_runtime() {
|
||||
let upstream = Upstream::serve(vec![Answer::Raw(OK)]);
|
||||
let url = upstream.url();
|
||||
let fetched = tokio::spawn(async move {
|
||||
fetch(
|
||||
&Client::for_test(),
|
||||
&S3Access::Environment { region: None },
|
||||
&url,
|
||||
POLICY,
|
||||
Start::Fresh,
|
||||
)
|
||||
.await
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(fetched.is_ok());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_redirect_is_not_followed_and_its_body_is_the_document() {
|
||||
let upstream = Upstream::serve(vec![
|
||||
Answer::Raw(
|
||||
"HTTP/1.1 301 Moved Permanently\r\nLocation: /elsewhere.json\r\nContent-Length: 7\r\nConnection: close\r\n\r\n{\"a\":1}",
|
||||
),
|
||||
Answer::Raw(OK),
|
||||
]);
|
||||
assert_eq!(
|
||||
fetched(&upstream.url(), POLICY, Start::Fresh).await,
|
||||
Ok(Fetched {
|
||||
body: b"{\"a\":1}".to_vec(),
|
||||
etag: None,
|
||||
})
|
||||
);
|
||||
assert_eq!(upstream.heads().len(), 1);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_retryable_status_is_retried_until_it_succeeds() {
|
||||
let upstream = Upstream::serve(vec![Answer::Raw(UNAVAILABLE), Answer::Raw(OK)]);
|
||||
assert!(fetched(&upstream.url(), POLICY, Start::Fresh).await.is_ok());
|
||||
assert_eq!(upstream.heads().len(), 2);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_retry_after_shorter_than_the_backoff_is_honored() {
|
||||
let upstream = Upstream::serve(vec![Answer::Raw(UNAVAILABLE), Answer::Raw(OK)]);
|
||||
let policy = RetryPolicy {
|
||||
min_delay: Duration::from_secs(10),
|
||||
max_delay: Duration::from_secs(10),
|
||||
..POLICY
|
||||
};
|
||||
let fetched = tokio::time::timeout(
|
||||
Duration::from_secs(2),
|
||||
fetched(&upstream.url(), policy, Start::Fresh),
|
||||
)
|
||||
.await
|
||||
.expect("Retry-After: 0 replaces the 10s backoff");
|
||||
assert!(fetched.is_ok(), "{fetched:?}");
|
||||
assert_eq!(upstream.heads().len(), 2);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn retryable_failures_stop_at_the_attempt_limit() {
|
||||
let upstream = Upstream::serve(vec![Answer::Raw(UNAVAILABLE)]);
|
||||
let url = upstream.url();
|
||||
assert_eq!(
|
||||
fetched(&url, POLICY, Start::Fresh).await,
|
||||
Err(Failed {
|
||||
reason: format!("HTTP 503 from {url} (after 3 attempts)"),
|
||||
retryable: true,
|
||||
retry_after: Some(Duration::ZERO),
|
||||
})
|
||||
);
|
||||
assert_eq!(upstream.heads().len(), 3);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_404_fails_on_the_first_attempt() {
|
||||
let upstream = Upstream::serve(vec![Answer::Raw(NOT_FOUND)]);
|
||||
let url = upstream.url();
|
||||
assert_eq!(
|
||||
fetched(&url, POLICY, Start::Fresh).await,
|
||||
Err(Failed {
|
||||
reason: format!("HTTP 404 from {url}"),
|
||||
retryable: false,
|
||||
retry_after: None,
|
||||
})
|
||||
);
|
||||
assert_eq!(upstream.heads().len(), 1);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_stalled_response_times_out_and_is_retried() {
|
||||
let upstream = Upstream::serve(vec![Answer::Stall]);
|
||||
let url = upstream.url();
|
||||
let policy = RetryPolicy {
|
||||
max_attempts: 2,
|
||||
timeout: Duration::from_millis(100),
|
||||
..POLICY
|
||||
};
|
||||
let failed = fetched(&url, policy, Start::Fresh).await.unwrap_err();
|
||||
assert!(
|
||||
failed
|
||||
.reason
|
||||
.starts_with(&format!("error fetching {url}: "))
|
||||
&& failed.reason.ends_with(" (after 2 attempts)"),
|
||||
"{}",
|
||||
failed.reason
|
||||
);
|
||||
assert!(failed.retryable);
|
||||
assert_eq!(upstream.heads().len(), 2);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_resumed_fetch_makes_only_the_attempts_left() {
|
||||
let upstream = Upstream::serve(vec![Answer::Raw(UNAVAILABLE)]);
|
||||
let url = upstream.url();
|
||||
let first = Failed {
|
||||
reason: "first".into(),
|
||||
retryable: true,
|
||||
retry_after: Some(Duration::ZERO),
|
||||
};
|
||||
let failed = fetched(&url, POLICY, Start::After(first))
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert_eq!(
|
||||
failed.reason,
|
||||
format!("HTTP 503 from {url} (after 3 attempts)")
|
||||
);
|
||||
assert_eq!(upstream.heads().len(), 2);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_resumed_fetch_with_no_attempts_left_returns_the_failure_it_resumed() {
|
||||
let upstream = Upstream::serve(vec![Answer::Raw(OK)]);
|
||||
let first = Failed {
|
||||
reason: "first".into(),
|
||||
retryable: true,
|
||||
retry_after: None,
|
||||
};
|
||||
let policy = RetryPolicy {
|
||||
max_attempts: 1,
|
||||
..POLICY
|
||||
};
|
||||
assert_eq!(
|
||||
fetched(&upstream.url(), policy, Start::After(first.clone())).await,
|
||||
Err(first)
|
||||
);
|
||||
assert!(upstream.heads().is_empty());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_failure_reason_drops_url_credentials() {
|
||||
let upstream = Upstream::serve(vec![Answer::Raw(NOT_FOUND)]);
|
||||
let address = upstream.base.trim_start_matches("http://");
|
||||
let url = format!("http://user:hunter2@{address}/cost.json?X-Amz-Signature=sig-secret#frag");
|
||||
let failed = fetched(&url, POLICY, Start::Fresh).await.unwrap_err();
|
||||
assert_eq!(failed.reason, format!("HTTP 404 from {}", upstream.url()));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::unsupported_scheme("ftp://example.com/cost.json")]
|
||||
#[case::not_a_url("not a url")]
|
||||
#[case::s3_without_a_key("s3://prices")]
|
||||
#[tokio::test]
|
||||
async fn an_unfetchable_url_fails_without_a_request(#[case] url: &str) {
|
||||
let failed = fetched(url, POLICY, Start::Fresh).await.unwrap_err();
|
||||
assert!(
|
||||
failed
|
||||
.reason
|
||||
.starts_with(&format!("error fetching {url}: ")),
|
||||
"{}",
|
||||
failed.reason
|
||||
);
|
||||
assert!(!failed.retryable);
|
||||
}
|
||||
|
||||
fn s3(upstream: &Upstream, credentials: impl ProvideCredentials + 'static) -> S3Access {
|
||||
S3Access::Config(Box::new(
|
||||
litellm_s3::config(
|
||||
"us-west-2".into(),
|
||||
Some(&upstream.base),
|
||||
Client::for_test(),
|
||||
credentials,
|
||||
)
|
||||
.build(),
|
||||
))
|
||||
}
|
||||
|
||||
async fn fetched_from_s3(
|
||||
url: &str,
|
||||
upstream: &Upstream,
|
||||
credentials: impl ProvideCredentials + 'static,
|
||||
) -> Result<Fetched, Failed> {
|
||||
fetch(
|
||||
&Client::for_test(),
|
||||
&s3(upstream, credentials),
|
||||
url,
|
||||
POLICY,
|
||||
Start::Fresh,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn an_s3_fetch_sends_a_signed_get_object_to_the_endpoint() {
|
||||
let upstream = Upstream::serve(vec![Answer::Raw(OK)]);
|
||||
assert_eq!(
|
||||
fetched_from_s3("s3://prices/maps/cost.json", &upstream, keys()).await,
|
||||
Ok(Fetched {
|
||||
body: b"{}".to_vec(),
|
||||
etag: Some("\"v1\"".into()),
|
||||
})
|
||||
);
|
||||
let heads = upstream.heads();
|
||||
let [head] = heads.as_slice() else {
|
||||
panic!("expected one request, got {heads:?}");
|
||||
};
|
||||
assert!(head.starts_with("GET /prices/maps/cost.json"), "{head}");
|
||||
let authorization = header(head, "authorization").unwrap_or_default();
|
||||
assert!(
|
||||
authorization.starts_with("AWS4-HMAC-SHA256 Credential=AKIDMODELCOST/")
|
||||
&& authorization.contains("/us-west-2/s3/aws4_request"),
|
||||
"{authorization}"
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn an_s3_body_that_stalls_after_the_headers_times_out() {
|
||||
let upstream = Upstream::serve(vec![Answer::StallBody]);
|
||||
let url = "s3://prices/cost.json";
|
||||
let policy = RetryPolicy {
|
||||
timeout: Duration::from_millis(200),
|
||||
..POLICY
|
||||
};
|
||||
let failed = tokio::time::timeout(
|
||||
Duration::from_secs(2),
|
||||
fetch(
|
||||
&Client::for_test(),
|
||||
&s3(&upstream, keys()),
|
||||
url,
|
||||
policy,
|
||||
Start::Fresh,
|
||||
),
|
||||
)
|
||||
.await
|
||||
.expect("the attempt timeout bounds the body read")
|
||||
.unwrap_err();
|
||||
assert!(
|
||||
failed
|
||||
.reason
|
||||
.starts_with(&format!("error fetching {url}: ")),
|
||||
"{}",
|
||||
failed.reason
|
||||
);
|
||||
assert!(!failed.retryable);
|
||||
assert_eq!(upstream.heads().len(), 1);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::http(|upstream: &Upstream| (S3Access::Environment { region: None }, upstream.url()))]
|
||||
#[case::s3(|upstream: &Upstream| (s3(upstream, keys()), "s3://prices/cost.json".to_owned()))]
|
||||
#[tokio::test]
|
||||
async fn a_body_that_keeps_arriving_finishes_after_the_timeout(
|
||||
#[case] source: fn(&Upstream) -> (S3Access, String),
|
||||
) {
|
||||
let upstream = Upstream::serve(vec![Answer::SlowBody]);
|
||||
let (access, url) = source(&upstream);
|
||||
let policy = RetryPolicy {
|
||||
timeout: Duration::from_millis(300),
|
||||
..POLICY
|
||||
};
|
||||
assert_eq!(
|
||||
fetch(&Client::for_test(), &access, &url, policy, Start::Fresh).await,
|
||||
Ok(Fetched {
|
||||
body: SLOW_BODY.to_vec(),
|
||||
etag: Some("\"v1\"".into()),
|
||||
})
|
||||
);
|
||||
assert_eq!(upstream.heads().len(), 1);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::missing_key(NOT_FOUND)]
|
||||
#[case::denied("HTTP/1.1 403 Forbidden\r\nContent-Length: 0\r\nConnection: close\r\n\r\n")]
|
||||
#[tokio::test]
|
||||
async fn an_s3_error_fails_without_a_retry(#[case] response: &'static str) {
|
||||
let upstream = Upstream::serve(vec![Answer::Raw(response), Answer::Raw(OK)]);
|
||||
let url = "s3://prices/cost.json";
|
||||
let failed = fetched_from_s3(url, &upstream, keys()).await.unwrap_err();
|
||||
assert!(
|
||||
failed
|
||||
.reason
|
||||
.starts_with(&format!("error fetching {url}: ")),
|
||||
"{}",
|
||||
failed.reason
|
||||
);
|
||||
assert!(!failed.retryable);
|
||||
assert_eq!(upstream.heads().len(), 1);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn an_s3_fetch_without_credentials_fails_without_a_request() {
|
||||
let upstream = Upstream::serve(vec![Answer::Raw(OK)]);
|
||||
let url = "s3://prices/cost.json";
|
||||
let failed = fetched_from_s3(url, &upstream, NoCredentials)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(
|
||||
failed
|
||||
.reason
|
||||
.starts_with(&format!("error fetching {url}: ")),
|
||||
"{}",
|
||||
failed.reason
|
||||
);
|
||||
assert!(!failed.retryable);
|
||||
assert!(upstream.heads().is_empty());
|
||||
}
|
||||
|
||||
const DECLARED_OVERSIZED: &str =
|
||||
"HTTP/1.1 200 OK\r\nContent-Length: 104857601\r\nConnection: close\r\n\r\n";
|
||||
|
||||
#[rstest]
|
||||
#[case::http_declared(Answer::Raw(DECLARED_OVERSIZED), false)]
|
||||
#[case::http_streamed(Answer::Oversized, false)]
|
||||
#[case::s3_declared(Answer::Raw(DECLARED_OVERSIZED), true)]
|
||||
#[case::s3_streamed(Answer::Oversized, true)]
|
||||
#[tokio::test]
|
||||
async fn a_body_over_100_mib_fails_without_a_retry(#[case] answer: Answer, #[case] s3: bool) {
|
||||
let upstream = Upstream::serve(vec![answer, Answer::Raw(OK)]);
|
||||
let failed = match s3 {
|
||||
true => fetched_from_s3("s3://prices/cost.json", &upstream, keys()).await,
|
||||
false => fetched(&upstream.url(), POLICY, Start::Fresh).await,
|
||||
}
|
||||
.unwrap_err();
|
||||
assert!(
|
||||
failed
|
||||
.reason
|
||||
.contains("response body is over the 104857600 byte limit"),
|
||||
"{}",
|
||||
failed.reason
|
||||
);
|
||||
assert!(!failed.retryable);
|
||||
assert_eq!(upstream.heads().len(), 1);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_body_of_exactly_100_mib_is_fetched() {
|
||||
let upstream = Upstream::serve(vec![Answer::Raw(Box::leak(
|
||||
format!(
|
||||
"HTTP/1.1 200 OK\r\nContent-Length: 104857600\r\nConnection: close\r\n\r\n{}",
|
||||
" ".repeat(100 * 1024 * 1024)
|
||||
)
|
||||
.into_boxed_str(),
|
||||
))]);
|
||||
let fetched = fetched(&upstream.url(), POLICY, Start::Fresh)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(fetched.body.len(), 100 * 1024 * 1024);
|
||||
}
|
||||
3
litellm-rust/crates/s3/AGENTS.md
Normal file
3
litellm-rust/crates/s3/AGENTS.md
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
`litellm-s3` builds `aws-sdk-s3` clients that send through `litellm_http::Client`, and owns the S3 concepts the SDK lacks, such as the `s3://bucket/key` URI. It reads no environment and picks no credentials: the host passes the region, the endpoint and a credentials provider, usually `litellm_auth_aws::AwsCredentialsProvider`
|
||||
|
||||
Operation policy such as retries, timeouts and checksum settings stays with the caller, which adjusts the returned config builder
|
||||
21
litellm-rust/crates/s3/Cargo.toml
Normal file
21
litellm-rust/crates/s3/Cargo.toml
Normal file
|
|
@ -0,0 +1,21 @@
|
|||
[package]
|
||||
name = "litellm-s3"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
litellm-http.workspace = true
|
||||
aws-sdk-s3.workspace = true
|
||||
aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio", "credentials-process"] }
|
||||
aws-types = "1.4.0"
|
||||
aws-credential-types = "1.3.0"
|
||||
aws-smithy-runtime-api = { version = "1.16.2", features = ["client", "http-1x"] }
|
||||
aws-smithy-types = { version = "1.6.0", features = ["http-body-1-x"] }
|
||||
http.workspace = true
|
||||
reqwest.workspace = true
|
||||
thiserror.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
rstest.workspace = true
|
||||
7
litellm-rust/crates/s3/src/error.rs
Normal file
7
litellm-rust/crates/s3/src/error.rs
Normal file
|
|
@ -0,0 +1,7 @@
|
|||
use thiserror::Error;
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub enum Error {
|
||||
#[error("expected s3://bucket/key")]
|
||||
InvalidUri,
|
||||
}
|
||||
60
litellm-rust/crates/s3/src/lib.rs
Normal file
60
litellm-rust/crates/s3/src/lib.rs
Normal file
|
|
@ -0,0 +1,60 @@
|
|||
mod error;
|
||||
mod transport;
|
||||
mod uri;
|
||||
|
||||
pub use aws_credential_types::provider::ProvideCredentials;
|
||||
pub use aws_sdk_s3 as sdk;
|
||||
pub use error::Error;
|
||||
pub use uri::S3Uri;
|
||||
|
||||
use aws_config::meta::region::RegionProviderChain;
|
||||
use aws_sdk_s3::config::{BehaviorVersion, Builder, Region};
|
||||
use aws_types::SdkConfig;
|
||||
use aws_types::service_config::ServiceConfigKey;
|
||||
|
||||
use crate::transport::ReqwestHttpClient;
|
||||
|
||||
/// A custom endpoint is addressed path style, as boto3 addresses one
|
||||
pub fn config(
|
||||
region: String,
|
||||
endpoint_url: Option<&str>,
|
||||
http: litellm_http::Client,
|
||||
credentials: impl ProvideCredentials + 'static,
|
||||
) -> Builder {
|
||||
let builder = aws_sdk_s3::Config::builder()
|
||||
.behavior_version(BehaviorVersion::latest())
|
||||
.region(Region::new(region))
|
||||
.http_client(ReqwestHttpClient(http))
|
||||
.credentials_provider(credentials);
|
||||
match endpoint_url {
|
||||
Some(url) => builder.endpoint_url(url).force_path_style(true),
|
||||
None => builder,
|
||||
}
|
||||
}
|
||||
|
||||
/// The region, endpoint and credentials an AWS SDK finds in the environment and shared config files, as boto3's
|
||||
/// default `s3` client found them. `region` wins over that chain, and `us-east-1` is used when nothing names one
|
||||
pub async fn environment_config(http: litellm_http::Client, region: Option<String>) -> Builder {
|
||||
let region = RegionProviderChain::first_try(region.map(Region::new))
|
||||
.or_default_provider()
|
||||
.or_else(Region::from_static("us-east-1"));
|
||||
let shared = aws_config::defaults(BehaviorVersion::latest())
|
||||
.region(region)
|
||||
.http_client(ReqwestHttpClient(http))
|
||||
.load()
|
||||
.await;
|
||||
Builder::from(&shared).force_path_style(has_custom_endpoint(&shared))
|
||||
}
|
||||
|
||||
fn has_custom_endpoint(shared: &SdkConfig) -> bool {
|
||||
let s3_endpoint = ServiceConfigKey::builder()
|
||||
.service_id("S3")
|
||||
.env("AWS_ENDPOINT_URL")
|
||||
.profile("endpoint_url")
|
||||
.build()
|
||||
.ok();
|
||||
shared.endpoint_url().is_some()
|
||||
|| s3_endpoint
|
||||
.and_then(|key| shared.service_config()?.load_config(key))
|
||||
.is_some()
|
||||
}
|
||||
76
litellm-rust/crates/s3/src/uri.rs
Normal file
76
litellm-rust/crates/s3/src/uri.rs
Normal file
|
|
@ -0,0 +1,76 @@
|
|||
use crate::Error;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct S3Uri {
|
||||
pub bucket: String,
|
||||
pub key: String,
|
||||
}
|
||||
|
||||
/// Splits the way Python's `urllib.parse.urlsplit` does, so the key keeps dot segments and characters a URL parser
|
||||
/// would normalize or percent-encode. S3 keys are opaque strings
|
||||
impl TryFrom<&str> for S3Uri {
|
||||
type Error = Error;
|
||||
|
||||
fn try_from(uri: &str) -> Result<Self, Error> {
|
||||
let rest = match uri.split_once("://") {
|
||||
Some((scheme, rest)) if scheme.eq_ignore_ascii_case("s3") => rest,
|
||||
_ => return Err(Error::InvalidUri),
|
||||
};
|
||||
let (bucket, path) = rest.split_at(rest.find(['/', '?', '#']).unwrap_or(rest.len()));
|
||||
let key = path
|
||||
.split(['?', '#'])
|
||||
.next()
|
||||
.unwrap_or_default()
|
||||
.trim_start_matches('/');
|
||||
if bucket.is_empty() || key.is_empty() {
|
||||
return Err(Error::InvalidUri);
|
||||
}
|
||||
Ok(Self {
|
||||
bucket: bucket.to_owned(),
|
||||
key: key.to_owned(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[rstest]
|
||||
#[case::key("s3://bucket/cost.json", "bucket", "cost.json")]
|
||||
#[case::nested_key("s3://bucket/prices/cost.json", "bucket", "prices/cost.json")]
|
||||
#[case::dotted_bucket("s3://my.bucket/cost.json", "my.bucket", "cost.json")]
|
||||
#[case::dot_segments("s3://bucket/archives/../cost.json", "bucket", "archives/../cost.json")]
|
||||
#[case::current_dir_segment("s3://bucket/./cost.json", "bucket", "./cost.json")]
|
||||
#[case::space_kept_raw("s3://bucket/my prices.json", "bucket", "my prices.json")]
|
||||
#[case::query_dropped("s3://bucket/cost.json?versionId=1", "bucket", "cost.json")]
|
||||
#[case::fragment_dropped("s3://bucket/cost.json#top", "bucket", "cost.json")]
|
||||
#[case::leading_slashes_trimmed("s3://bucket//cost.json", "bucket", "cost.json")]
|
||||
#[case::uppercase_scheme("S3://bucket/cost.json", "bucket", "cost.json")]
|
||||
fn an_s3_url_names_its_bucket_and_key(
|
||||
#[case] url: &str,
|
||||
#[case] bucket: &str,
|
||||
#[case] key: &str,
|
||||
) {
|
||||
assert_eq!(
|
||||
S3Uri::try_from(url).unwrap(),
|
||||
S3Uri {
|
||||
bucket: bucket.into(),
|
||||
key: key.into(),
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::no_key("s3://bucket/")]
|
||||
#[case::no_path("s3://bucket")]
|
||||
#[case::query_only("s3://bucket?cost.json")]
|
||||
#[case::no_bucket("s3:///cost.json")]
|
||||
#[case::other_scheme("https://bucket/cost.json")]
|
||||
#[case::no_scheme("bucket/cost.json")]
|
||||
fn a_url_without_a_bucket_and_key_is_refused(#[case] url: &str) {
|
||||
assert!(matches!(S3Uri::try_from(url), Err(Error::InvalidUri)));
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue