scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 200

refactor(rust-error): keep causes as sources across http, auth, secrets, cache and utility crates

Browse source 
Messages describe their own layer and every cause is a typed source,
predicates match exhaustively, and boundaries render with
litellm_error::Report. Media downloads no longer count as dispatched,
Error::upstream bounds the provider body itself, and a forwarded
header SigV4 cannot sign is the caller's fault.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Yujong Lee 2026-10-23 12:33:59 -07:00
parent 2de856f685
commit 4432805bf7
228 changed files with 3069 additions and 1742 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm-rust/.agents/skills/rust-error/references/messages-and-causes.md
View file

@ -66,7 +66,7 @@ Store the lower-level error itself, never its string, and never drop it with `ma
| downcasting | the type is gone |
| redaction | `reqwest::Error`'s `Display` includes the request URL, which can carry credentials in the path or query |
If the cause is not `Clone` and the error must be, hold it as `litellm_error::Source`, a cloneable `Arc<dyn Error + Send + Sync>` that is transparent to `Display`, `source()` and `downcast_ref`, or make the crate's error opaque. Do not stringify or drop a cause to get `Clone` or `PartialEq`
If the cause is not `Clone` and the error must be, hold it as `#[source] litellm_error::Source`, a cloneable `Arc<dyn Error + Send + Sync>` that is transparent to `Display` and `source()`, or make the crate's error opaque. `litellm_error::find::<T>(&error)` walks the chain and looks through `Source`, so it is how anyone reaches a typed cause. A third-party `Box<dyn Error + Send + Sync>`, or an `anyhow::Error` a dependency returns, converts with `Source::from`, and a dependency's own boxed alias such as `tokenizers::Error` is held as it is. Never put `#[error(transparent)]` on a `Source`, the erased error would then be no node of the chain and `find` could not reach it. Do not stringify or drop a cause to get `Clone` or `PartialEq`
```rust
// loses the type and leaks the URL

4
litellm-rust/.agents/skills/rust-error/references/testing.md
View file

@ -20,7 +20,7 @@ In order of preference
2. A real cause, such as `tungstenite::Error::ConnectionClosed` or a `serde_json::from_str("{")` error, passed through our `From`
3. A public constructor, such as `Error::upstream(500, "boom".into())`
Inside the owning crate a test may match the private representation through the `#[cfg(test)] fn repr(&self) -> &Repr` accessor in `crates/http/src/error.rs`, as the `media.rs` and `tls.rs` tests do. Anywhere else, walk `source()` and `downcast_ref` to the cause the test expects
Inside the owning crate a test may match the private representation through the `#[cfg(test)] fn repr(&self) -> &Repr` accessor in `crates/http/src/error.rs`, as the `media.rs` and `tls.rs` tests do. Anywhere else, reach the cause the test expects with `litellm_error::find::<T>(&error)`
## What to cover
@ -28,7 +28,7 @@ Inside the owning crate a test may match the private representation through the
|---|---|
| classification | one `#[case]` per input class, asserting the predicate |
| redaction | build from a URL holding `api_key=secret`, assert the message lacks `secret` and the path |
| cause kept | walk `source()` and `downcast_ref` to the expected type |
| cause kept | `litellm_error::find::<T>(&error)` returns the expected cause |
| rendered once | the cause's text appears exactly once in `litellm_error::Report::new(&error).to_string()` and not in `to_string()` |
| boundary mapping | the gateway status and body, or the Python exception type and its `status` attribute |

36
litellm-rust/Cargo.lock generated
View file

@ -3375,6 +3375,7 @@ dependencies = [
"bytes",
"futures-util",
"litellm-auth-types",
"litellm-error",
"litellm-http",
"moka",
"reqwest",
@ -3394,6 +3395,7 @@ dependencies = [
"azure_core",
"azure_identity",
"litellm-auth-types",
"litellm-error",
"moka",
"rstest",
"serde_json",
@ -3412,6 +3414,7 @@ dependencies = [
"google-cloud-auth",
"http 1.4.2",
"litellm-auth-types",
"litellm-error",
"moka",
"rstest",
"serde",
@ -3426,6 +3429,7 @@ dependencies = [
name = "litellm-auth-types"
version = "0.1.0"
dependencies = [
"litellm-error",
"rstest",
"serde",
"serde_json",
@ -3439,6 +3443,7 @@ dependencies = [
name = "litellm-cache"
version = "0.1.0"
dependencies = [
"litellm-error",
"rstest",
"serde",
"serde_json",
@ -3464,6 +3469,7 @@ dependencies = [
"reqwest",
"rstest",
"serde_json",
"thiserror 2.0.19",
"tokio",
"url",
"wiremock",
@ -3482,6 +3488,7 @@ dependencies = [
"serde-pickle",
"serde_json",
"tempfile",
"thiserror 2.0.19",
"tokio",
]
@ -3494,11 +3501,13 @@ dependencies = [
"litellm-auth-types",
"litellm-cache",
"litellm-cache-testing",
"litellm-error",
"litellm-http",
"percent-encoding",
"reqwest",
"rstest",
"serde_json",
"thiserror 2.0.19",
"tokio",
"wiremock",
]
@ -3521,6 +3530,7 @@ dependencies = [
"futures-util",
"litellm-cache",
"litellm-cache-testing",
"litellm-error",
"litellm-http",
"qdrant-client",
"reqwest",
@ -3541,11 +3551,13 @@ version = "0.1.0"
dependencies = [
"litellm-cache",
"litellm-cache-testing",
"litellm-error",
"r2d2",
"redis",
"redis-test",
"rstest",
"serde_json",
"thiserror 2.0.19",
"tokio",
]
@ -3598,6 +3610,7 @@ dependencies = [
"litellm-auth-aws",
"litellm-cache",
"litellm-cache-testing",
"litellm-error",
"litellm-http",
"reqwest",
"rstest",
@ -3620,11 +3633,13 @@ dependencies = [
"litellm-cache",
"litellm-cache-redis",
"litellm-cache-testing",
"litellm-error",
"redis",
"redis-test",
"rstest",
"serde_json",
"sha2 0.10.9",
"thiserror 2.0.19",
"tokio",
"uuid",
]
@ -3648,6 +3663,7 @@ name = "litellm-config"
version = "0.1.0"
dependencies = [
"litellm-auth-types",
"litellm-error",
"litellm-router-types",
"rstest",
"serde",
@ -3663,6 +3679,7 @@ name = "litellm-core-utils"
version = "0.1.0"
dependencies = [
"fancy-regex 0.19.2",
"litellm-error",
"litellm-llms-types",
"litellm-tracing",
"rstest",
@ -3690,6 +3707,7 @@ dependencies = [
"criterion",
"proptest",
"rstest",
"thiserror 2.0.19",
]
[[package]]
@ -3777,6 +3795,7 @@ dependencies = [
"futures-util",
"litellm-auth-types",
"litellm-config",
"litellm-error",
"litellm-secrets",
"rstest",
"serde",
@ -3934,6 +3953,7 @@ dependencies = [
"bytes",
"futures-util",
"litellm-coroutine",
"litellm-error",
"litellm-host",
"parking_lot",
"pyo3",
@ -3954,6 +3974,7 @@ dependencies = [
"http 1.4.2",
"hyper-util",
"litellm-core-utils",
"litellm-error",
"rcgen",
"reqwest",
"rstest",
@ -4068,6 +4089,7 @@ dependencies = [
"litellm-auth-aws",
"litellm-auth-gcp",
"litellm-core-utils",
"litellm-error",
"litellm-host",
"litellm-host-native",
"litellm-http",
@ -4165,6 +4187,7 @@ dependencies = [
"litellm-auth-azure",
"litellm-auth-gcp",
"litellm-core-utils",
"litellm-error",
"litellm-framer",
"litellm-host",
"litellm-http",
@ -4205,6 +4228,7 @@ version = "0.1.0"
dependencies = [
"indexmap 2.14.0",
"jsonschema",
"litellm-error",
"litellm-llms-types",
"rstest",
"schemars 1.2.2",
@ -4231,6 +4255,7 @@ dependencies = [
"litellm-cache-response",
"litellm-callbacks-legacy-python",
"litellm-core-utils",
"litellm-error",
"litellm-host",
"litellm-host-python",
"litellm-http",
@ -4317,6 +4342,7 @@ dependencies = [
"google-cloud-auth",
"google-cloud-kms-v1",
"litellm-core-utils",
"litellm-error",
"litellm-http",
"litellm-python-compat",
"litellm-secrets-aws",
@ -4375,6 +4401,7 @@ dependencies = [
"sha2 0.10.9",
"thiserror 2.0.19",
"tokio",
"url",
"veil",
"wiremock",
]
@ -4398,6 +4425,7 @@ dependencies = [
"tempfile",
"thiserror 2.0.19",
"tokio",
"url",
"veil",
"wiremock",
]
@ -4424,6 +4452,7 @@ dependencies = [
"serde_json",
"thiserror 2.0.19",
"tokio",
"url",
"veil",
"wiremock",
]
@ -4433,6 +4462,7 @@ name = "litellm-secrets-hashicorp"
version = "0.1.0"
dependencies = [
"litellm-core-utils",
"litellm-error",
"litellm-secrets-types",
"rstest",
"rustify",
@ -4467,6 +4497,7 @@ name = "litellm-spend-clickhouse"
version = "0.1.0"
dependencies = [
"flate2",
"litellm-error",
"litellm-http",
"litellm-storage-clickhouse",
"rstest",
@ -4487,6 +4518,7 @@ name = "litellm-storage-clickhouse"
version = "0.1.0"
dependencies = [
"flate2",
"litellm-error",
"litellm-http",
"rstest",
"serde",
@ -4506,6 +4538,7 @@ version = "0.1.0"
dependencies = [
"flate2",
"futures-util",
"litellm-http",
"reqwest",
"rstest",
"semver",
@ -4529,6 +4562,7 @@ dependencies = [
"criterion",
"indexmap 2.14.0",
"itoa",
"litellm-error",
"litellm-token-counter-fast",
"litellm-token-counter-huggingface",
"litellm-token-counter-tiktoken",
@ -4545,6 +4579,7 @@ name = "litellm-token-counter-fast"
version = "0.1.0"
dependencies = [
"base64 0.22.1",
"litellm-error",
"rand 0.8.7",
"rstest",
"rustc-hash",
@ -4570,6 +4605,7 @@ name = "litellm-token-counter-tiktoken"
version = "0.1.0"
dependencies = [
"base64 0.22.1",
"litellm-error",
"once_cell",
"rstest",
"rustc-hash",

1
litellm-rust/crates/auth-aws/Cargo.toml
View file

@ -7,6 +7,7 @@ repository.workspace = true
[dependencies]
litellm-auth-types.workspace = true
litellm-error.workspace = true
litellm-http.workspace = true
reqwest.workspace = true

92
litellm-rust/crates/auth-aws/src/aws.rs
View file

@ -15,13 +15,13 @@ use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
use litellm_auth_types::AwsParams;
use super::Error;
use super::constants::{
AWS_ACCESS_KEY_ID, AWS_EXTERNAL_ID, AWS_PROFILE_NAME, AWS_REGION_NAME, AWS_ROLE_ARN,
AWS_ROLE_NAME, AWS_SECRET_ACCESS_KEY, AWS_SESSION_NAME, AWS_SESSION_TOKEN,
AWS_SIGNED_HEADER_NAMES, AWS_STS_ENDPOINT, AWS_WEB_IDENTITY_TOKEN, AWS_WEB_IDENTITY_TOKEN_FILE,
DEFAULT_BEDROCK_REGION, DEFAULT_SESSION_NAME_PREFIX, SIGV4_COMPUTED_HEADER_NAMES,
};
use crate::error::{CredentialFlow, Error, SigningStage};
const STATIC_CREDENTIALS_TTL: Duration = Duration::from_secs(3600 - 60);
const AMBIENT_CREDENTIALS_TTL: Duration = Duration::from_secs(600);
@ -271,7 +271,7 @@ impl AwsAuthService {
provider
.provide_credentials()
.await
.map_err(|error| Error::AwsProfile(error.to_string()))
.map_err(Error::credentials(CredentialFlow::Profile))
}
AwsAuthFlow::AssumeRole { role, session_name } => {
if is_already_running_as_role(&role, &resolved).await? {
@ -287,7 +287,7 @@ impl AwsAuthService {
let credentials = provider
.provide_credentials()
.await
.map_err(|error| Error::AwsDefaultChain(error.to_string()))?;
.map_err(Error::credentials(CredentialFlow::DefaultChain))?;
self.set_cached_credentials(
key,
credentials.clone(),
@ -327,7 +327,7 @@ impl AwsAuthService {
provider
.provide_credentials()
.await
.map_err(|error| Error::AwsAssumeRole(error.to_string()))
.map_err(Error::credentials(CredentialFlow::AssumeRole))
}
AwsAuthFlow::WebIdentity {
token,
@ -350,12 +350,14 @@ impl AwsAuthService {
.web_identity_token(token)
.send()
.await
.map_err(|error| Error::AwsWebIdentity(error.to_string()))?;
.map_err(Error::credentials(CredentialFlow::WebIdentity))?;
let credentials = response
.credentials()
.ok_or(Error::AwsMissingWebIdentityCredentials)?;
let expiration = SystemTime::try_from(*credentials.expiration())
.map_err(|error| Error::AwsWebIdentityExpiration(error.to_string()))?;
.ok_or(Error::MissingWebIdentityCredentials)?;
let expiration =
SystemTime::try_from(*credentials.expiration()).map_err(|source| {
Error::WebIdentityExpiration(litellm_error::Source::new(source))
})?;
Ok(Credentials::new(
credentials.access_key_id(),
credentials.secret_access_key(),
@ -376,7 +378,7 @@ impl AwsAuthService {
let credentials = provider
.provide_credentials()
.await
.map_err(|error| Error::AwsDefaultChain(error.to_string()))?;
.map_err(Error::credentials(CredentialFlow::DefaultChain))?;
self.set_cached_credentials(
key,
credentials.clone(),
@ -480,14 +482,16 @@ pub fn sign_headers(
.settings(SigningSettings::default())
.build()
.map(SigningParams::from)
.map_err(|error| Error::AwsSigningParameters(error.to_string()))?;
.map_err(Error::signing(SigningStage::Parameters))?;
let header_refs = headers
.iter()
.map(|(name, value)| {
value
.to_str()
.map(|value| (name.as_str(), value))
.map_err(|_| Error::AwsSignableRequest(format!("{name} is not visible ASCII")))
.map_err(|_| Error::NonAsciiHeader {
name: name.to_string(),
})
})
.collect::<Result<Vec<_>, _>>()?;
let request = SignableRequest::new(
@ -496,23 +500,20 @@ pub fn sign_headers(
header_refs.into_iter(),
SignableBody::Bytes(body),
)
.map_err(|error| Error::AwsSignableRequest(error.to_string()))?;
.map_err(Error::signing(SigningStage::Request))?;
let (instructions, _) = sign(request, &params)
.map_err(|error| Error::AwsSigning(error.to_string()))?
.map_err(Error::signing(SigningStage::Signature))?
.into_parts();
instructions
.headers()
.map(|(name, value)| {
HeaderName::from_bytes(name.as_bytes())
.map_err(|error| error.to_string())
.and_then(|name| {
HeaderValue::from_str(value)
.map(|value| (name, value))
.map_err(|error| error.to_string())
})
let name = HeaderName::from_bytes(name.as_bytes())
.map_err(Error::signing(SigningStage::Output))?;
let value =
HeaderValue::from_str(value).map_err(Error::signing(SigningStage::Output))?;
Ok((name, value))
})
.collect::<Result<HeaderMap, _>>()
.map_err(Error::AwsSigning)
.collect()
}
/// Model-id and region parsing shared by every Bedrock route.
@ -1081,6 +1082,53 @@ mod tests {
assert!(!signed.contains_key("X-Amz-Security-Token"));
}
#[test]
fn a_header_value_outside_visible_ascii_is_named_instead_of_signed() {
let (url, body, mut headers) = parity_inputs();
headers.insert("x-amz-meta-note", HeaderValue::from_str("café").unwrap());
let credentials = Credentials::new("AKIDEXAMPLE", "secret", None, None, "test");
let error = sign_headers(
"POST",
&url,
&body,
&headers,
SigningScope {
region: "us-east-1",
service: BEDROCK_SERVICE,
credentials: &credentials,
time: UNIX_EPOCH,
},
)
.unwrap_err();
assert_eq!(
error,
Error::NonAsciiHeader {
name: "x-amz-meta-note".into()
}
);
}
#[test]
fn the_shared_auth_error_nests_the_aws_failure_and_its_cause() {
let cause = std::io::Error::new(std::io::ErrorKind::NotFound, "profile not found");
let error =
litellm_auth_types::Error::from(Error::credentials(CredentialFlow::Profile)(cause));
assert!(matches!(
litellm_error::find::<Error>(&error),
Some(Error::Credentials {
flow: CredentialFlow::Profile,
..
})
));
assert_eq!(
litellm_error::Report::new(&error).to_string(),
"provider authentication failed: could not load AWS profile credentials: profile not found"
);
}
#[ignore]
#[tokio::test]
async fn live_bedrock_invoke_model_returns_200() -> Result<(), Box<dyn std::error::Error>> {

101
litellm-rust/crates/auth-aws/src/error.rs
View file

@ -1,46 +1,73 @@
use thiserror::Error as ThisError;
use litellm_error::Source;
#[derive(Clone, Debug, ThisError, PartialEq, Eq)]
#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)]
pub enum Error {
#[error("AWS profile credentials failed: {0}")]
AwsProfile(String),
#[error("AWS default credentials failed: {0}")]
AwsDefaultChain(String),
#[error("AWS role credentials failed: {0}")]
AwsAssumeRole(String),
#[error("AWS web identity credentials failed: {0}")]
AwsWebIdentity(String),
#[error("AWS web identity expiration was invalid: {0}")]
AwsWebIdentityExpiration(String),
#[error("AWS signing parameters failed: {0}")]
AwsSigningParameters(String),
#[error("AWS signable request failed: {0}")]
AwsSignableRequest(String),
#[error("AWS request signing failed: {0}")]
AwsSigning(String),
#[error("could not load AWS {flow} credentials")]
Credentials {
flow: CredentialFlow,
#[source]
source: Source,
},
#[error("AWS web identity response had no credentials")]
AwsMissingWebIdentityCredentials,
MissingWebIdentityCredentials,
#[error("AWS web identity response had an invalid expiration")]
WebIdentityExpiration(#[source] Source),
#[error("header {name} has a value that is not visible ASCII")]
NonAsciiHeader { name: String },
#[error("{stage}")]
Signing {
stage: SigningStage,
#[source]
source: Source,
},
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)]
pub enum CredentialFlow {
#[error("profile")]
Profile,
#[error("default chain")]
DefaultChain,
#[error("assumed role")]
AssumeRole,
#[error("web identity")]
WebIdentity,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)]
pub enum SigningStage {
#[error("SigV4 signing parameters are invalid")]
Parameters,
#[error("SigV4 cannot sign this request")]
Request,
#[error("SigV4 could not compute the signature")]
Signature,
#[error("SigV4 produced an invalid header")]
Output,
}
impl Error {
pub(crate) fn credentials<E: std::error::Error + Send + Sync + 'static>(
flow: CredentialFlow,
) -> impl FnOnce(E) -> Self {
move |source| Self::Credentials {
flow,
source: Source::new(source),
}
}
pub(crate) fn signing<E: std::error::Error + Send + Sync + 'static>(
stage: SigningStage,
) -> impl FnOnce(E) -> Self {
move |source| Self::Signing {
stage,
source: Source::new(source),
}
}
}
impl From<Error> for litellm_auth_types::Error {
fn from(error: Error) -> Self {
Self::ProviderAuthentication(error.to_string())
}
}
#[cfg(test)]
mod tests {
use super::Error;
#[test]
fn converts_to_shared_auth_error_without_losing_context() {
let error = litellm_auth_types::Error::from(Error::AwsProfile("profile not found".into()));
assert_eq!(
error,
litellm_auth_types::Error::ProviderAuthentication(
"AWS profile credentials failed: profile not found".into()
)
);
Self::provider(error)
}
}

2
litellm-rust/crates/auth-aws/src/lib.rs
View file

@ -5,5 +5,5 @@ mod signer;
pub use aws::*;
pub use aws_credential_types::Credentials;
pub use error::Error;
pub use error::{CredentialFlow, Error, SigningStage};
pub use signer::SigV4Signer;

28
litellm-rust/crates/auth-aws/src/signer.rs
View file

@ -73,7 +73,13 @@ impl RequestSigner for SigV4Signer {
time: (self.clock)(),
},
)
.map_err(litellm_http::Error::signature)
.map_err(|error| match error {
Error::NonAsciiHeader { .. } => litellm_http::Error::unsignable_header(error),
Error::Credentials { .. }
| Error::MissingWebIdentityCredentials
| Error::WebIdentityExpiration(_)
| Error::Signing { .. } => litellm_http::Error::signature(error),
})
}
}
@ -187,6 +193,26 @@ mod tests {
assert!(error.to_string().contains("authorization"), "{error}");
}
#[test]
fn a_forwarded_header_that_cannot_be_signed_is_the_callers_fault() {
let error = signed(
json(
URL,
&[("x-amz-meta-note".into(), "café".into())],
&json_value!({}),
None,
)
.unwrap(),
&signer("textract"),
)
.unwrap_err();
let report = litellm_error::Report::new(&error).to_string();
assert!(report.contains("x-amz-meta-note"), "{report}");
assert!(error.is_request());
assert!(!error.dispatched());
}
#[test]
fn a_streaming_body_cannot_be_signed() {
let mut request = Request::new(Method::POST, Url::parse(URL).unwrap());

1
litellm-rust/crates/auth-azure/Cargo.toml
View file

@ -18,5 +18,6 @@ azure_core = "1.0.0"
azure_identity = { version = "1.0.0", features = ["tokio"] }
[dev-dependencies]
litellm-error.workspace = true
rstest.workspace = true
tokio.workspace = true

43
litellm-rust/crates/auth-azure/src/native.rs
View file

@ -12,7 +12,7 @@ use azure_identity::{
};
use sha2::{Digest, Sha256};
use litellm_auth_types::Error;
use litellm_auth_types::{Error, ErrorDetail};
use litellm_auth_types::{InputSource, ResolvedCredential, SecretValue, Sourced};
use super::credential_provider_cache::{
@ -133,7 +133,9 @@ impl NativeAzureTokenAcquirer {
let token = credential
.get_token(&[scope.as_str()], None)
.await
.map_err(|error| Error::CredentialAcquisition(error.to_string().into()))?;
.map_err(|error| {
Error::CredentialAcquisition(ErrorDetail::failed("Azure token request", error))
})?;
let expires_on = u64::try_from(token.expires_on.unix_timestamp())
.ok()
.map(|seconds| UNIX_EPOCH + Duration::from_secs(seconds));
@ -250,11 +252,8 @@ fn validate_authority(request: &NativeAzureRequest) -> Result<(), Error> {
let Some(authority) = authority else {
return Ok(());
};
let url = url::Url::parse(authority.value()).map_err(|_| {
Error::InvalidConfiguration(
"Azure authority must be an HTTPS origin without credentials, query, or fragment"
.into(),
)
let url = url::Url::parse(authority.value()).map_err(|error| {
Error::InvalidConfiguration(ErrorDetail::failed("Azure authority parsing", error))
})?;
if url.scheme() != "https"
|| url.host_str().is_none()
@ -264,10 +263,10 @@ fn validate_authority(request: &NativeAzureRequest) -> Result<(), Error> {
|| url.fragment().is_some()
|| !matches!(url.path(), "" | "/")
{
return Err(Error::InvalidConfiguration(
"Azure authority must be an HTTPS origin without credentials, query, or fragment"
.into(),
));
return Err(Error::InvalidConfiguration(ErrorDetail::InvalidType {
field: "Azure authority".into(),
expected: "an HTTPS origin without credentials, query, or fragment",
}));
}
Ok(())
}
@ -377,7 +376,7 @@ fn trusted_source(sources: &[InputSource]) -> InputSource {
fn mixed_sources<T>() -> Result<T, Error> {
Err(Error::InvalidConfiguration(
"request-controlled Azure auth inputs cannot be combined with host credentials".into(),
ErrorDetail::MixedInputSources { provider: "Azure" },
))
}
@ -444,7 +443,7 @@ fn build_credential(
.map(|credential| credential as Arc<dyn TokenCredential>),
}
.map_err(|error| {
Error::InvalidConfiguration(litellm_auth_types::ErrorDetail::failed(
Error::InvalidConfiguration(ErrorDetail::failed(
"Azure credential initialization",
error,
))
@ -662,12 +661,11 @@ mod tests {
))
.unwrap_err();
assert_eq!(
error,
litellm_auth_types::Error::InvalidConfiguration(
assert!(
litellm_error::Report::new(&error).to_string().contains(
"request-controlled Azure auth inputs cannot be combined with host credentials"
.into()
)
),
"{error:?}"
);
}
@ -695,12 +693,11 @@ mod tests {
authority,
))
.unwrap_err();
assert_eq!(
error,
litellm_auth_types::Error::InvalidConfiguration(
assert!(
litellm_error::Report::new(&error).to_string().contains(
"Azure authority must be an HTTPS origin without credentials, query, or fragment"
.into()
)
),
"{error:?}"
);
}
}

111
litellm-rust/crates/auth-azure/src/resolve.rs
View file

@ -1,8 +1,8 @@
use litellm_auth_types::Error;
use litellm_auth_types::{
CredentialFileRef, CredentialLookup, CredentialRef, InputSource, ResolvedCredential,
SecretValue, Sourced, TokenProviderHandle,
};
use litellm_auth_types::{Error, ErrorDetail};
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
@ -91,9 +91,9 @@ impl AzureAuthService {
AzureCredentialPlan::Caller(caller) => {
let credential = caller.acquire().await?;
if credential.secret().expose().is_empty() {
return Err(Error::EmptyCallerCredential(
"Azure AD token provider returned an empty token",
));
return Err(Error::EmptyCallerCredential {
caller: "Azure AD token provider",
});
}
Ok(Some(Sourced::new(credential, InputSource::Deployment)))
}
@ -106,11 +106,9 @@ impl AzureAuthService {
} => {
let assertion = resolve_reference(inputs, env_lookup, reference.value())
.await?
.ok_or_else(|| {
Error::CredentialAcquisition(
"Azure OIDC reference did not resolve to a value".into(),
)
})?;
.ok_or(Error::CredentialAcquisition(ErrorDetail::Unresolved {
subject: "Azure OIDC reference",
}))?;
let request = ValidatedAzureRequest::new(NativeAzureRequest::ClientAssertion {
tenant_id,
client_id,
@ -173,7 +171,12 @@ pub(crate) fn select_auth_plan(
.map(|selector| Sourced::new(selector, value.source()))
})
.transpose()
.map_err(|_| Error::InvalidConfiguration("invalid Azure credential selector".into()))?;
.map_err(|error| {
Error::InvalidConfiguration(ErrorDetail::failed(
"Azure credential selector parsing",
error,
))
})?;
let federated_token_file = configured_string(
&inputs.federated_token_file,
AZURE_FEDERATED_TOKEN_FILE_ENV,
@ -263,9 +266,12 @@ fn select_native_plan(
let selection_source = selected.source();
match selected.into_value() {
AzureCredentialType::ClientSecretCredential => Err(Error::InvalidConfiguration(
"ClientSecretCredential requires tenant_id, client_id, and client_secret".into(),
)),
AzureCredentialType::ClientSecretCredential => {
Err(Error::InvalidConfiguration(ErrorDetail::MissingSettings {
credential: "ClientSecretCredential",
settings: "tenant_id, client_id, and client_secret",
}))
}
AzureCredentialType::WorkloadIdentityCredential => {
Ok(AzureCredentialPlan::Native(ValidatedAzureRequest::new(
workload_request(tenant_id, client_id, federated_token_file, scope, authority)?,
@ -349,22 +355,21 @@ fn workload_request(
authority: Option<Sourced<String>>,
) -> Result<NativeAzureRequest, Error> {
Ok(NativeAzureRequest::WorkloadIdentity {
tenant_id: tenant_id.ok_or_else(|| {
Error::InvalidConfiguration("WorkloadIdentityCredential requires tenant_id".into())
})?,
client_id: client_id.ok_or_else(|| {
Error::InvalidConfiguration("WorkloadIdentityCredential requires client_id".into())
})?,
token_file_path: token_file_path.ok_or_else(|| {
Error::InvalidConfiguration(
"WorkloadIdentityCredential requires azure_federated_token_file".into(),
)
})?,
tenant_id: tenant_id.ok_or(workload_requires("tenant_id"))?,
client_id: client_id.ok_or(workload_requires("client_id"))?,
token_file_path: token_file_path.ok_or(workload_requires("azure_federated_token_file"))?,
scope,
authority,
})
}
fn workload_requires(settings: &'static str) -> Error {
Error::InvalidConfiguration(ErrorDetail::MissingSettings {
credential: "WorkloadIdentityCredential",
settings,
})
}
fn configured_string(
configured: &ConfigValue<String>,
environment_name: &str,
@ -410,11 +415,10 @@ async fn resolve_reference(
.map_or(CredentialLookup::Missing, CredentialLookup::Found),
CredentialRef::None => return Ok(None),
CredentialRef::File(_) | CredentialRef::Request(_) | CredentialRef::Host(_) => {
let resolver = inputs.credential_resolver.as_ref().ok_or_else(|| {
Error::InvalidConfiguration(
"credential reference requires a host credential resolver".into(),
)
})?;
let resolver = inputs
.credential_resolver
.as_ref()
.ok_or(Error::InvalidConfiguration(ErrorDetail::MissingResolver))?;
resolver.resolve(reference).await?
}
};
@ -433,7 +437,7 @@ fn oidc_reference(
let value = token.value().expose();
if token.source() == InputSource::Request && value.starts_with("oidc/") {
return Err(Error::InvalidConfiguration(
"request-controlled Azure credential references are not allowed".into(),
ErrorDetail::RequestControlledReference { provider: "Azure" },
));
}
if let Some(name) = value.strip_prefix("oidc/env/") {
@ -457,19 +461,17 @@ fn oidc_reference(
}
if value.starts_with("oidc/") {
return Err(Error::InvalidConfiguration(
"unsupported OIDC reference".into(),
ErrorDetail::UnsupportedReference,
));
}
Ok(None)
}
fn non_empty_reference(value: &str, kind: &str) -> Result<String, Error> {
fn non_empty_reference(value: &str, kind: &'static str) -> Result<String, Error> {
if value.is_empty() {
return Err(Error::InvalidConfiguration(
litellm_auth_types::ErrorDetail::Empty {
subject: kind.into(),
},
));
return Err(Error::InvalidConfiguration(ErrorDetail::Empty {
subject: kind,
}));
}
Ok(value.to_string())
}
@ -518,7 +520,9 @@ mod tests {
expires_on: None,
})
} else {
Err(Error::CredentialAcquisition(kind.into()))
Err(Error::CredentialAcquisition(
litellm_auth_types::ErrorDetail::Unresolved { subject: kind },
))
}
})
}
@ -624,7 +628,11 @@ mod tests {
)))
.expect_err("unsupported backend must fail validation");
assert!(error.to_string().contains("unsupported OIDC reference"));
assert!(
litellm_error::Report::new(&error)
.to_string()
.contains("unsupported OIDC reference")
);
}
#[rstest::rstest]
@ -649,11 +657,11 @@ mod tests {
})
.unwrap_err();
assert_eq!(
error,
Error::InvalidConfiguration(
"request-controlled Azure credential references are not allowed".into()
)
assert!(
litellm_error::Report::new(&error)
.to_string()
.contains("request-controlled Azure credential references are not allowed"),
"{error:?}"
);
}
@ -715,7 +723,16 @@ mod tests {
.await
.unwrap_err();
assert!(matches!(error, Error::CredentialChain(errors) if errors.len() == 2));
assert!(matches!(&error, Error::CredentialChain(errors) if errors.len() == 2));
let report = litellm_error::Report::new(&error).to_string();
assert!(
report.contains("managed-identity did not resolve"),
"{report}"
);
assert!(
report.contains("developer-tools did not resolve"),
"{report}"
);
}
#[derive(Debug)]
@ -763,7 +780,9 @@ mod tests {
assert_eq!(
error,
Error::EmptyCallerCredential("Azure AD token provider returned an empty token")
Error::EmptyCallerCredential {
caller: "Azure AD token provider"
}
);
}
}

1
litellm-rust/crates/auth-gcp/Cargo.toml
View file

@ -24,5 +24,6 @@ google-cloud-auth = { workspace = true, optional = true }
http = { workspace = true, optional = true }
[dev-dependencies]
litellm-error.workspace = true
rstest.workspace = true
tempfile.workspace = true

16
litellm-rust/crates/auth-gcp/src/auth.rs
View file

@ -234,14 +234,13 @@ enum CredentialCacheKey {
Adc,
}
const BASE64_CREDENTIALS: &str = "The `vertex_credentials` value is base64-encoded JSON, not the JSON itself. \
Set it to a readable file path, or to the credentials JSON itself.";
fn load_trusted_service_account(configured: &str) -> Result<CustomServiceAccount, Error> {
if is_base64_encoded_json(configured) {
return Err(Error::InvalidConfiguration(ErrorDetail::Message(
BASE64_CREDENTIALS.into(),
)));
return Err(Error::InvalidConfiguration(
ErrorDetail::Base64EncodedJson {
field: "vertex_credentials",
},
));
}
if names_a_credentials_file(configured) {
CustomServiceAccount::from_file(configured)
@ -299,7 +298,8 @@ mod tests {
else {
panic!("{configured:?} should not load");
};
assert!(error.to_string().contains(expected), "{error}");
let report = litellm_error::Report::new(&error).to_string();
assert!(report.contains(expected), "{report}");
}
#[rstest::rstest]
@ -321,7 +321,7 @@ mod tests {
.await
{
Ok(_) => panic!("{configured:?} should not load"),
Err(error) => error.to_string(),
Err(error) => litellm_error::Report::new(&error).to_string(),
}
}

6
litellm-rust/crates/auth-gcp/src/sdk.rs
View file

@ -55,7 +55,7 @@ impl CredentialsProvider for GoogleCredentials {
entity_tag: EntityTag::new(),
data,
})
.map_err(|_| CredentialsError::from_msg(false, "Google authentication failed"))
.map_err(|error| CredentialsError::from_source(false, error))
}
async fn universe_domain(&self) -> Option<String> {
@ -102,5 +102,9 @@ mod tests {
);
let error = credentials.headers(Extensions::new()).await.unwrap_err();
assert!(!format!("{error:?}").contains("private"));
assert_eq!(
std::error::Error::source(&error).and_then(|cause| cause.downcast_ref::<Error>()),
Some(&Error::InvalidHeader)
);
}
}

1
litellm-rust/crates/auth-types/Cargo.toml
View file

@ -6,6 +6,7 @@ license.workspace = true
repository.workspace = true
[dependencies]
litellm-error.workspace = true
serde.workspace = true
serde_json.workspace = true
subtle.workspace = true

26
litellm-rust/crates/auth-types/src/credential.rs
View file

@ -86,9 +86,9 @@ impl CredentialPlan {
Self::Caller(caller) => {
let credential = caller.acquire().await?;
if credential.secret().expose().is_empty() {
return Err(Error::EmptyCallerCredential(
"credential caller returned an empty credential",
));
return Err(Error::EmptyCallerCredential {
caller: "credential caller",
});
}
Ok(CredentialPlanResolution::Resolved(credential))
}
@ -104,8 +104,8 @@ mod tests {
CredentialLookup, CredentialLookupFuture, CredentialPlan, CredentialPlanResolution,
CredentialRef, CredentialResolver, CredentialResolverHandle,
};
use crate::Error;
use crate::SecretValue;
use crate::{Error, ErrorDetail};
#[derive(Debug)]
struct HostResolver;
@ -150,9 +150,10 @@ mod tests {
impl CredentialResolver for FailingResolver {
fn resolve<'a>(&'a self, _reference: &'a CredentialRef) -> CredentialLookupFuture<'a> {
Box::pin(async {
Err(Error::CredentialAcquisition(
"host credential lookup failed".into(),
))
Err(Error::CredentialAcquisition(ErrorDetail::failed(
"host credential lookup",
std::io::Error::from(std::io::ErrorKind::PermissionDenied),
)))
})
}
}
@ -168,9 +169,14 @@ mod tests {
.await
.expect_err("acquisition errors cannot become fallback");
assert_eq!(
error,
Error::CredentialAcquisition("host credential lookup failed".into())
assert!(
matches!(error, Error::CredentialAcquisition(_)),
"{error:?}"
);
assert!(
litellm_error::Report::new(&error)
.to_string()
.contains("host credential lookup failed")
);
}
}

287
litellm-rust/crates/auth-types/src/error.rs
View file

@ -1,16 +1,18 @@
use thiserror::Error as ThisError;
use litellm_error::{Report, Source};
#[derive(Clone, Debug, ThisError, PartialEq, Eq)]
use crate::{ParamKind, ParamSpec};
#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)]
pub enum Error {
#[error("invalid authentication configuration: {0}")]
InvalidConfiguration(ErrorDetail),
#[error("credential acquisition failed: {0}")]
CredentialAcquisition(ErrorDetail),
#[error("credential caller failed: {0}")]
EmptyCallerCredential(&'static str),
#[error("{0}")]
ProviderAuthentication(String),
#[error("credential acquisition failed: {}", .0.iter().map(ToString::to_string).collect::<Vec<_>>().join("; "))]
#[error("invalid authentication configuration")]
InvalidConfiguration(#[source] ErrorDetail),
#[error("credential acquisition failed")]
CredentialAcquisition(#[source] ErrorDetail),
#[error("{caller} returned an empty credential")]
EmptyCallerCredential { caller: &'static str },
#[error("provider authentication failed")]
Provider(#[source] Source),
#[error("no credential in the chain succeeded: {}", reports(.0))]
CredentialChain(Vec<Error>),
#[error(
"Missing {provider} API Key - Set `api_key` or the {environment_variable} environment variable"
@ -22,7 +24,7 @@ pub enum Error {
#[error("Missing {provider} {} - {}", .spec.setting, .spec.guidance())]
MissingParam {
provider: &'static str,
spec: &'static crate::ParamSpec,
spec: &'static ParamSpec,
},
#[error("Missing {provider} API Base - {guidance}")]
MissingApiBase {
@ -33,31 +35,97 @@ pub enum Error {
InvalidHeader,
}
fn reports(failures: &[Error]) -> String {
failures
.iter()
.map(|failure| Report::new(failure).to_string())
.collect::<Vec<_>>()
.join("; ")
}
impl Error {
pub fn provider(cause: impl std::error::Error + Send + Sync + 'static) -> Self {
Self::Provider(Source::new(cause))
}
/// The environment lacks the credential the provider needs, as opposed to the caller
/// having sent a bad one.
pub fn is_missing_credential(&self) -> bool {
matches!(self, Self::MissingApiKey { .. })
match self {
Self::MissingApiKey { .. } => true,
Self::MissingParam { spec, .. } => match spec.kind {
ParamKind::Credential => true,
ParamKind::Setting => false,
},
Self::InvalidConfiguration(_)
| Self::CredentialAcquisition(_)
| Self::EmptyCallerCredential { .. }
| Self::Provider(_)
| Self::CredentialChain(_)
| Self::MissingApiBase { .. }
| Self::InvalidHeader => false,
}
}
/// A connection setting the caller can pass on the call, such as a region, a project or an
/// API base, is unset everywhere it is read from.
pub fn is_missing_setting(&self) -> bool {
match self {
Self::MissingApiBase { .. } => true,
Self::MissingParam { spec, .. } => match spec.kind {
ParamKind::Setting => true,
ParamKind::Credential => false,
},
Self::InvalidConfiguration(_)
| Self::CredentialAcquisition(_)
| Self::EmptyCallerCredential { .. }
| Self::Provider(_)
| Self::CredentialChain(_)
| Self::MissingApiKey { .. }
| Self::InvalidHeader => false,
}
}
}
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
pub enum ErrorDetail {
#[error("{0}")]
Message(String),
#[error("{field} must be {expected}")]
InvalidType {
field: String,
expected: &'static str,
},
#[error("{subject} cannot be empty")]
Empty { subject: String },
Empty { subject: &'static str },
#[error("credential header {0} already exists")]
DuplicateHeader(&'static str),
#[error("{operation} failed: {cause}")]
#[error("credential plan is not allowed by the provider auth policy")]
PlanNotAllowed,
#[error("{credential} requires {settings}")]
MissingSettings {
credential: &'static str,
settings: &'static str,
},
#[error("credentials cannot be sent to a request-controlled {provider} endpoint")]
RequestControlledEndpoint { provider: &'static str },
#[error("request-controlled {provider} credential references are not allowed")]
RequestControlledReference { provider: &'static str },
#[error("request-controlled {provider} auth inputs cannot be combined with host credentials")]
MixedInputSources { provider: &'static str },
#[error("unsupported OIDC reference")]
UnsupportedReference,
#[error("credential reference requires a host credential resolver")]
MissingResolver,
#[error("{subject} did not resolve to a value")]
Unresolved { subject: &'static str },
#[error(
"The `{field}` value is base64-encoded JSON, not the JSON itself. Set it to a readable file path, or to the credentials JSON itself."
)]
Base64EncodedJson { field: &'static str },
#[error("{operation} failed")]
Failed {
operation: &'static str,
cause: ErrorSource,
#[source]
cause: Source,
},
}
@ -68,50 +136,151 @@ impl ErrorDetail {
) -> Self {
Self::Failed {
operation,
cause: ErrorSource::new(cause),
cause: Source::new(cause),
}
}
}
impl From<String> for ErrorDetail {
fn from(message: String) -> Self {
Self::Message(message)
#[cfg(test)]
mod tests {
use std::io;
use litellm_error::{Report, Source};
use rstest::rstest;
use super::{Error, ErrorDetail};
use crate::{AwsParams, VertexParams};
#[rstest]
#[case::api_key(
Error::MissingApiKey { provider: "Example", environment_variable: "EXAMPLE_API_KEY" },
"Missing Example API Key - Set `api_key` or the EXAMPLE_API_KEY environment variable"
)]
#[case::another_api_key(
Error::MissingApiKey { provider: "Custom", environment_variable: "CUSTOM_KEY" },
"Missing Custom API Key - Set `api_key` or the CUSTOM_KEY environment variable"
)]
#[case::api_base(
Error::MissingApiBase { provider: "Example", guidance: "Pass api_base" },
"Missing Example API Base - Pass api_base"
)]
#[case::another_api_base(
Error::MissingApiBase { provider: "Custom", guidance: "Set CUSTOM_ENDPOINT" },
"Missing Custom API Base - Set CUSTOM_ENDPOINT"
)]
fn missing_inputs_keep_the_python_text(#[case] error: Error, #[case] expected: &str) {
assert_eq!(error.to_string(), expected);
}
#[rstest]
#[case::configuration(
Error::InvalidConfiguration(ErrorDetail::Empty { subject: "credential selector" }),
"invalid authentication configuration: credential selector cannot be empty"
)]
#[case::acquisition(
Error::CredentialAcquisition(ErrorDetail::Unresolved { subject: "OIDC reference" }),
"credential acquisition failed: OIDC reference did not resolve to a value"
)]
#[case::empty_caller(
Error::EmptyCallerCredential { caller: "token provider" },
"token provider returned an empty credential"
)]
#[case::chain(
Error::CredentialChain(vec![
Error::CredentialAcquisition(ErrorDetail::Unresolved { subject: "OIDC reference" }),
Error::EmptyCallerCredential { caller: "token provider" },
]),
"no credential in the chain succeeded: credential acquisition failed: OIDC reference did not resolve \
to a value; token provider returned an empty credential"
)]
fn the_report_names_every_layer_in_order(#[case] error: Error, #[case] expected: &str) {
assert!(
Report::new(&error).to_string().contains(expected),
"{}",
Report::new(&error)
);
}
#[rstest]
#[case::configuration(Error::InvalidConfiguration(ErrorDetail::failed(
"test credential",
denied()
)))]
#[case::acquisition(Error::CredentialAcquisition(ErrorDetail::failed(
"test credential",
denied()
)))]
fn a_failure_keeps_its_cause_as_the_source_and_renders_it_once(#[case] error: Error) {
let cause =
std::iter::successors(std::error::Error::source(&error), |cause| cause.source())
.find_map(|cause| cause.downcast_ref::<Source>()?.downcast_ref::<io::Error>())
.expect("the io cause is in the source chain");
assert_eq!(cause.kind(), io::ErrorKind::PermissionDenied);
let report = Report::new(&error).to_string();
assert_eq!(report.matches("access denied").count(), 1, "{report}");
}
#[rstest]
#[case::configuration(Error::InvalidConfiguration(ErrorDetail::failed(
"test credential",
denied()
)))]
#[case::acquisition(Error::CredentialAcquisition(ErrorDetail::failed(
"test credential",
denied()
)))]
fn a_contextual_failure_leaves_its_cause_out_of_its_own_message(#[case] error: Error) {
assert!(!error.to_string().contains("access denied"), "{error}");
assert!(
Report::new(&error)
.to_string()
.contains("test credential failed: access denied")
);
}
#[test]
fn a_provider_failure_keeps_the_provider_error_as_its_source() {
let error = Error::provider(denied());
assert_eq!(
litellm_error::find::<io::Error>(&error).map(io::Error::kind),
Some(io::ErrorKind::PermissionDenied)
);
assert_eq!(
Report::new(&error).to_string(),
"provider authentication failed: access denied"
);
}
fn denied() -> io::Error {
io::Error::new(io::ErrorKind::PermissionDenied, "access denied")
}
#[rstest]
#[case::api_key(Error::MissingApiKey { provider: "Example", environment_variable: "EXAMPLE_API_KEY" }, true)]
#[case::vertex_credentials(VertexParams::CREDENTIALS.missing("Vertex AI"), true)]
#[case::aws_secret_access_key(AwsParams::SECRET_ACCESS_KEY.missing("AWS"), true)]
#[case::vertex_project(VertexParams::PROJECT.missing("Vertex AI"), false)]
#[case::aws_region(AwsParams::REGION.missing("AWS"), false)]
#[case::api_base(Error::MissingApiBase { provider: "Example", guidance: "Pass api_base" }, false)]
#[case::invalid_configuration(Error::InvalidConfiguration(ErrorDetail::PlanNotAllowed), false)]
fn only_an_absent_credential_is_a_missing_credential(
#[case] error: Error,
#[case] expected: bool,
) {
assert_eq!(error.is_missing_credential(), expected);
}
#[rstest]
#[case::vertex_project(VertexParams::PROJECT.missing("Vertex AI"), true)]
#[case::aws_region(AwsParams::REGION.missing("AWS"), true)]
#[case::api_base(Error::MissingApiBase { provider: "Example", guidance: "Pass api_base" }, true)]
#[case::vertex_credentials(VertexParams::CREDENTIALS.missing("Vertex AI"), false)]
#[case::api_key(Error::MissingApiKey { provider: "Example", environment_variable: "EXAMPLE_API_KEY" }, false)]
#[case::invalid_configuration(Error::InvalidConfiguration(ErrorDetail::PlanNotAllowed), false)]
fn an_absent_connection_setting_is_a_missing_setting(
#[case] error: Error,
#[case] expected: bool,
) {
assert_eq!(error.is_missing_setting(), expected);
}
}
impl From<&str> for ErrorDetail {
fn from(message: &str) -> Self {
Self::Message(message.into())
}
}
#[derive(Clone, Debug)]
pub struct ErrorSource(std::sync::Arc<dyn std::error::Error + Send + Sync>);
impl std::ops::Deref for ErrorSource {
type Target = dyn std::error::Error + Send + Sync;
fn deref(&self) -> &Self::Target {
self.0.as_ref()
}
}
impl std::fmt::Display for ErrorSource {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
std::fmt::Display::fmt(&self.0, formatter)
}
}
impl ErrorSource {
pub fn new(error: impl std::error::Error + Send + Sync + 'static) -> Self {
Self(std::sync::Arc::new(error))
}
}
impl PartialEq for ErrorSource {
fn eq(&self, other: &Self) -> bool {
std::sync::Arc::ptr_eq(&self.0, &other.0)
}
}
impl Eq for ErrorSource {}

12
litellm-rust/crates/auth-types/src/http.rs
View file

@ -21,9 +21,9 @@ pub fn apply_credential(
placement: CredentialPlacement,
) -> Result<Vec<(String, String)>, Error> {
if credential.trim().is_empty() {
return Err(Error::InvalidConfiguration(
"credential cannot be empty".into(),
));
return Err(Error::InvalidConfiguration(crate::ErrorDetail::Empty {
subject: "credential",
}));
}
if headers
.iter()
@ -71,6 +71,10 @@ mod tests {
)
.expect_err("provider policy must handle existing credentials");
assert!(error.to_string().contains("already exists"));
assert!(
litellm_error::Report::new(&error)
.to_string()
.contains("already exists")
);
}
}

4
litellm-rust/crates/auth-types/src/lib.rs
View file

@ -51,9 +51,9 @@ pub use credential::{
CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialPlan,
CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle,
};
pub use error::{Error, ErrorDetail, ErrorSource};
pub use error::{Error, ErrorDetail};
pub use http::CredentialPlacement;
pub use params::{AwsParams, ParamSpec, VertexParams};
pub use params::{AwsParams, ParamKind, ParamSpec, VertexParams};
pub use policy::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy};
pub use secret::SecretValue;
pub use token::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle};

49
litellm-rust/crates/auth-types/src/params/aws.rs
View file

@ -4,15 +4,17 @@ use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use veil::Redact;
use super::ParamSpec;
use super::{ParamKind, ParamSpec};
const fn spec(
setting: &'static str,
kind: ParamKind,
wire: &'static [&'static str],
env: &'static [&'static str],
) -> ParamSpec {
ParamSpec {
setting,
kind,
wire,
module_global: None,
env,
@ -53,40 +55,67 @@ pub struct AwsParams {
impl AwsParams {
pub const ACCESS_KEY_ID: ParamSpec = spec(
"access key id",
ParamKind::Credential,
&["aws_access_key_id"],
&["AWS_ACCESS_KEY_ID"],
);
pub const SECRET_ACCESS_KEY: ParamSpec = spec(
"secret access key",
ParamKind::Credential,
&["aws_secret_access_key"],
&["AWS_SECRET_ACCESS_KEY"],
);
pub const SESSION_TOKEN: ParamSpec = spec(
"session token",
ParamKind::Credential,
&["aws_session_token"],
&["AWS_SESSION_TOKEN"],
);
pub const REGION: ParamSpec = spec(
"region",
ParamKind::Setting,
&["aws_region_name"],
&["AWS_REGION_NAME", "AWS_REGION"],
);
pub const SESSION_NAME: ParamSpec =
spec("session name", &["aws_session_name"], &["AWS_SESSION_NAME"]);
pub const PROFILE_NAME: ParamSpec =
spec("profile name", &["aws_profile_name"], &["AWS_PROFILE_NAME"]);
pub const ROLE_NAME: ParamSpec = spec("role name", &["aws_role_name"], &["AWS_ROLE_NAME"]);
pub const SESSION_NAME: ParamSpec = spec(
"session name",
ParamKind::Setting,
&["aws_session_name"],
&["AWS_SESSION_NAME"],
);
pub const PROFILE_NAME: ParamSpec = spec(
"profile name",
ParamKind::Setting,
&["aws_profile_name"],
&["AWS_PROFILE_NAME"],
);
pub const ROLE_NAME: ParamSpec = spec(
"role name",
ParamKind::Setting,
&["aws_role_name"],
&["AWS_ROLE_NAME"],
);
pub const WEB_IDENTITY_TOKEN: ParamSpec = spec(
"web identity token",
ParamKind::Credential,
&["aws_web_identity_token"],
&["AWS_WEB_IDENTITY_TOKEN"],
);
pub const STS_ENDPOINT: ParamSpec =
spec("STS endpoint", &["aws_sts_endpoint"], &["AWS_STS_ENDPOINT"]);
pub const EXTERNAL_ID: ParamSpec =
spec("external id", &["aws_external_id"], &["AWS_EXTERNAL_ID"]);
pub const STS_ENDPOINT: ParamSpec = spec(
"STS endpoint",
ParamKind::Setting,
&["aws_sts_endpoint"],
&["AWS_STS_ENDPOINT"],
);
pub const EXTERNAL_ID: ParamSpec = spec(
"external id",
ParamKind::Setting,
&["aws_external_id"],
&["AWS_EXTERNAL_ID"],
);
pub const BEDROCK_RUNTIME_ENDPOINT: ParamSpec = spec(
"Bedrock runtime endpoint",
ParamKind::Setting,
&["aws_bedrock_runtime_endpoint"],
&["AWS_BEDROCK_RUNTIME_ENDPOINT"],
);

7
litellm-rust/crates/auth-types/src/params/mod.rs
View file

@ -10,11 +10,18 @@ pub use vertex::VertexParams;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ParamSpec {
pub setting: &'static str,
pub kind: ParamKind,
pub wire: &'static [&'static str],
pub module_global: Option<&'static str>,
pub env: &'static [&'static str],
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ParamKind {
Credential,
Setting,
}
impl ParamSpec {
/// The first non-blank value: each wire name through `params`, then each environment
/// name through `env`.

5
litellm-rust/crates/auth-types/src/params/vertex.rs
View file

@ -2,7 +2,7 @@ use serde::{Deserialize, Deserializer, Serialize, de::Unexpected};
use serde_json::Value;
use veil::Redact;
use super::ParamSpec;
use super::{ParamKind, ParamSpec};
/// The Vertex AI fields of Python's `GenericLiteLLMParams`, both the current names and the
/// `vertex_ai_*` spellings `VertexBase.safe_get_vertex_ai_*` still read.
@ -27,18 +27,21 @@ pub struct VertexParams {
impl VertexParams {
pub const CREDENTIALS: ParamSpec = ParamSpec {
setting: "credentials",
kind: ParamKind::Credential,
wire: &["vertex_credentials", "vertex_ai_credentials"],
module_global: None,
env: &["VERTEXAI_CREDENTIALS"],
};
pub const PROJECT: ParamSpec = ParamSpec {
setting: "project",
kind: ParamKind::Setting,
wire: &["vertex_project", "vertex_ai_project"],
module_global: Some("vertex_project"),
env: &["VERTEXAI_PROJECT"],
};
pub const LOCATION: ParamSpec = ParamSpec {
setting: "location",
kind: ParamKind::Setting,
wire: &["vertex_location", "vertex_ai_location"],
module_global: Some("vertex_location"),
env: &["VERTEXAI_LOCATION", "VERTEX_LOCATION"],

35
litellm-rust/crates/auth-types/src/policy.rs
View file

@ -1,5 +1,5 @@
use crate::http::apply_credential;
use crate::{CredentialPlacement, Error, ResolvedCredential};
use crate::{CredentialPlacement, Error, ErrorDetail, ResolvedCredential};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CredentialPlanKind {
@ -31,11 +31,18 @@ pub struct ProviderAuthPolicy {
impl ProviderAuthPolicy {
pub fn has_existing_credential(&self, headers: &[(String, String)]) -> bool {
headers.iter().any(|(name, _)| {
self.accepted_existing_headers
.iter()
.any(|accepted| name.eq_ignore_ascii_case(accepted))
})
self.existing_credential_header(headers).is_some()
}
fn existing_credential_header(&self, headers: &[(String, String)]) -> Option<&'static str> {
self.accepted_existing_headers
.iter()
.copied()
.find(|accepted| {
headers
.iter()
.any(|(name, _)| name.eq_ignore_ascii_case(accepted))
})
}
pub fn apply(
@ -44,11 +51,11 @@ impl ProviderAuthPolicy {
kind: CredentialPlanKind,
credential: &ResolvedCredential,
) -> Result<Vec<(String, String)>, Error> {
if self.has_existing_credential(&headers) {
if let Some(existing) = self.existing_credential_header(&headers) {
return match self.existing_header_behavior {
ExistingHeaderBehavior::Preserve => Ok(headers),
ExistingHeaderBehavior::Reject => Err(Error::InvalidConfiguration(
"credential header already exists".into(),
ErrorDetail::DuplicateHeader(existing),
)),
};
}
@ -56,11 +63,7 @@ impl ProviderAuthPolicy {
.rules
.iter()
.find(|rule| rule.kind == kind)
.ok_or_else(|| {
Error::InvalidConfiguration(
"credential plan is not allowed by the provider auth policy".into(),
)
})?;
.ok_or(Error::InvalidConfiguration(ErrorDetail::PlanNotAllowed))?;
apply_credential(headers, credential.secret().expose(), rule.placement)
}
}
@ -108,6 +111,10 @@ mod tests {
)
.unwrap_err();
assert!(error.to_string().contains("not allowed"));
assert!(
litellm_error::Report::new(&error)
.to_string()
.contains("not allowed")
);
}
}

81
litellm-rust/crates/auth-types/tests/error.rs
View file

@ -1,81 +0,0 @@
use litellm_auth_types::Error;
use rstest::rstest;
#[rstest]
#[case::api_key(
Error::MissingApiKey { provider: "Example", environment_variable: "EXAMPLE_API_KEY" },
"Missing Example API Key - Set `api_key` or the EXAMPLE_API_KEY environment variable"
)]
#[case::another_api_key(
Error::MissingApiKey { provider: "Custom", environment_variable: "CUSTOM_KEY" },
"Missing Custom API Key - Set `api_key` or the CUSTOM_KEY environment variable"
)]
#[case::api_base(
Error::MissingApiBase { provider: "Example", guidance: "Pass api_base" },
"Missing Example API Base - Pass api_base"
)]
#[case::another_api_base(
Error::MissingApiBase { provider: "Custom", guidance: "Set CUSTOM_ENDPOINT" },
"Missing Custom API Base - Set CUSTOM_ENDPOINT"
)]
#[case::configuration(
Error::InvalidConfiguration("credential selector is invalid".into()),
"invalid authentication configuration: credential selector is invalid"
)]
#[case::acquisition(
Error::CredentialAcquisition("token expired".into()),
"credential acquisition failed: token expired"
)]
#[case::caller(
Error::EmptyCallerCredential("empty token"),
"credential caller failed: empty token"
)]
#[case::provider(
Error::ProviderAuthentication("provider rejected credentials".into()),
"provider rejected credentials"
)]
#[case::chain(
Error::CredentialChain(vec![
Error::CredentialAcquisition("token expired".into()),
Error::EmptyCallerCredential("empty token"),
]),
"credential acquisition failed: credential acquisition failed: token expired; credential caller failed: empty token"
)]
fn display_preserves_failure_phase_and_caller_context(
#[case] error: Error,
#[case] expected: &str,
) {
assert_eq!(error.to_string(), expected);
}
#[rstest]
#[case::configuration(true)]
#[case::acquisition(false)]
fn contextual_failures_render_their_cause_once(#[case] configuration: bool) {
use litellm_auth_types::ErrorDetail;
let cause = std::io::Error::from(std::io::ErrorKind::PermissionDenied);
let cause_message = cause.to_string();
let detail = ErrorDetail::failed("test credential", cause);
let error = if configuration {
Error::InvalidConfiguration(detail)
} else {
Error::CredentialAcquisition(detail)
};
let (Error::InvalidConfiguration(detail) | Error::CredentialAcquisition(detail)) = &error
else {
unreachable!("built above");
};
let ErrorDetail::Failed { cause, .. } = detail else {
panic!("a contextual failure keeps its typed cause: {detail:?}");
};
assert_eq!(
cause
.downcast_ref::<std::io::Error>()
.map(std::io::Error::kind),
Some(std::io::ErrorKind::PermissionDenied)
);
assert!(error.to_string().contains("test credential failed:"));
assert!(error.to_string().ends_with(&cause_message));
assert!(std::error::Error::source(&error).is_none());
}

3
litellm-rust/crates/auth-types/tests/params.rs
View file

@ -1,8 +1,9 @@
use litellm_auth_types::{AwsParams, ParamSpec};
use litellm_auth_types::{AwsParams, ParamKind, ParamSpec};
use rstest::rstest;
const LOCATION: ParamSpec = ParamSpec {
setting: "location",
kind: ParamKind::Setting,
wire: &["vertex_location", "vertex_ai_location"],
module_global: Some("vertex_location"),
env: &["VERTEXAI_LOCATION", "VERTEX_LOCATION"],

1
litellm-rust/crates/cache-azure-blob/Cargo.toml
View file

@ -16,6 +16,7 @@ azure_core = "1.1.0"
azure_storage_blob = "1.1.0"
futures-util.workspace = true
reqwest.workspace = true
thiserror.workspace = true
tokio.workspace = true
url.workspace = true

30
litellm-rust/crates/cache-azure-blob/src/cache.rs
View file

@ -18,6 +18,10 @@ use url::Url;
use crate::{credential::AzureBlobCredential, transport::ReqwestTransport};
#[derive(Debug, thiserror::Error)]
#[error("account URL cannot hold a container path")]
struct NotABaseUrl;
pub struct AzureBlobCache<C> {
container: BlobContainerClient,
codec: C,
@ -57,12 +61,12 @@ impl<C: CacheCodec> AzureBlobCache<C> {
codec: C,
runtime: Handle,
) -> Result<Self, Error> {
let parsed = Url::parse(account_url).map_err(|_| Error::Unavailable)?;
let parsed = Url::parse(account_url).map_err(Error::unavailable)?;
let account_url = parsed.as_str().trim_end_matches('/').to_string();
let container_url = {
let mut url = parsed;
url.path_segments_mut()
.map_err(|()| Error::Unavailable)?
.map_err(|()| Error::unavailable(NotABaseUrl))?
.pop_if_empty()
.push(container);
url
@ -75,7 +79,7 @@ impl<C: CacheCodec> AzureBlobCache<C> {
..BlobContainerClientOptions::default()
}),
)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
let cache = Self {
container: client,
codec,
@ -101,7 +105,7 @@ impl<C: CacheCodec> AzureBlobCache<C> {
Err(error) if is_storage_error(&error, StorageErrorCode::ContainerAlreadyExists) => {
Ok(())
}
Err(_) => Err(Error::Unavailable),
Err(error) => Err(Error::unavailable(error)),
}
}
@ -116,7 +120,7 @@ impl<C: CacheCodec> AzureBlobCache<C> {
{
Ok(_) => Ok(()),
Err(error) if !overwrite && is_already_present(&error) => Ok(()),
Err(_) => Err(Error::Unavailable),
Err(error) => Err(Error::unavailable(error)),
}
}
@ -126,13 +130,9 @@ impl<C: CacheCodec> AzureBlobCache<C> {
Err(error) if is_storage_error(&error, StorageErrorCode::BlobNotFound) => {
return Ok(None);
}
Err(_) => return Err(Error::Unavailable),
Err(error) => return Err(Error::unavailable(error)),
};
let bytes = response
.body
.collect()
.await
.map_err(|_| Error::Unavailable)?;
let bytes = response.body.collect().await.map_err(Error::unavailable)?;
self.codec.decode(&bytes).map(Some)
}
@ -140,16 +140,16 @@ impl<C: CacheCodec> AzureBlobCache<C> {
let mut pages = self
.container
.list_blobs(None)
.map_err(|_| Error::Unavailable)?
.map_err(Error::unavailable)?
.into_pages();
while let Some(page) = pages.try_next().await.map_err(|_| Error::Unavailable)? {
let page = page.into_model().map_err(|_| Error::Unavailable)?;
while let Some(page) = pages.try_next().await.map_err(Error::unavailable)? {
let page = page.into_model().map_err(Error::unavailable)?;
for name in page.blob_items.into_iter().filter_map(|item| item.name) {
self.container
.blob_client(&name)
.delete(None)
.await
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
}
}
Ok(())

15
litellm-rust/crates/cache-azure-blob/src/credential.rs
View file

@ -17,6 +17,10 @@ const LLM_TOKEN_ENV: &str = "AZURE_AD_TOKEN";
type EnvLookup = Arc<dyn Fn(&str) -> Option<String> + Send + Sync>;
#[derive(Debug, thiserror::Error)]
#[error("no Azure credential is available for blob storage")]
struct NoCredential;
pub struct AzureBlobCredential {
service: AzureAuthService,
env_lookup: EnvLookup,
@ -62,15 +66,8 @@ impl TokenCredential for AzureBlobCredential {
&lookup,
)
.await
.map_err(|error| {
azure_core::Error::with_message(ErrorKind::Credential, error.to_string())
})?
.ok_or_else(|| {
azure_core::Error::with_message(
ErrorKind::Credential,
"no Azure credential is available for blob storage",
)
})?;
.map_err(|error| azure_core::Error::new(ErrorKind::Credential, error))?
.ok_or_else(|| azure_core::Error::new(ErrorKind::Credential, NoCredential))?;
let (token, expires_on) = match credential.into_value() {
ResolvedCredential::AccessToken { token, expires_on } => (token, expires_on),
ResolvedCredential::Static(token) => (token, None),

8
litellm-rust/crates/cache-azure-blob/src/transport.rs
View file

@ -29,7 +29,7 @@ impl HttpClient for ReqwestTransport {
} else {
ErrorKind::Io
};
azure_core::Error::new(kind, error)
azure_core::Error::new(kind, litellm_http::Error::from(error))
})?;
let status = response.status().as_u16().into();
let mut headers = Headers::new();
@ -41,9 +41,9 @@ impl HttpClient for ReqwestTransport {
);
}
}
let body = response
.bytes_stream()
.map_err(|error| azure_core::Error::new(ErrorKind::Io, error));
let body = response.bytes_stream().map_err(|error| {
azure_core::Error::new(ErrorKind::Io, litellm_http::Error::from(error))
});
Ok(AsyncRawResponse::new(status, headers, Box::pin(body)))
}
}

12
litellm-rust/crates/cache-azure-blob/tests/cache.rs
View file

@ -128,7 +128,7 @@ fn connect_surfaces_service_failures() {
ResponseCacheCodec,
runtime.handle().clone(),
));
assert!(matches!(result, Err(Error::Unavailable)));
assert!(matches!(result, Err(Error::Unavailable(_))));
}
#[rstest]
@ -263,7 +263,7 @@ fn malformed_blobs_are_invalid_entries(fixture: Fixture, #[case] key: &str, #[ca
fixture.service.seed_blob(key, bytes);
assert!(matches!(
fixture.cache.get_cache(key, &no_ttl()),
Err(Error::InvalidEntry)
Err(Error::InvalidEntry(_))
));
}
@ -378,15 +378,15 @@ fn service_failures_map_to_unavailable(fixture: Fixture) {
fixture.service.set_failing(true);
assert!(matches!(
fixture.cache.get_cache("key", &no_ttl()),
Err(Error::Unavailable)
Err(Error::Unavailable(_))
));
assert!(matches!(
fixture.cache.set_cache("key", entry(json!(1)), &no_ttl()),
Err(Error::Unavailable)
Err(Error::Unavailable(_))
));
assert!(matches!(
fixture.cache.flush_cache(),
Err(Error::Unavailable)
Err(Error::Unavailable(_))
));
assert!(matches!(
fixture.runtime.block_on(
@ -394,7 +394,7 @@ fn service_failures_map_to_unavailable(fixture: Fixture) {
.cache
.async_set_cache_pipeline(vec![("k".into(), entry(json!(1)))], no_ttl())
),
Err(Error::Unavailable)
Err(Error::Unavailable(_))
));
}

1
litellm-rust/crates/cache-disk/Cargo.toml
View file

@ -12,6 +12,7 @@ rand.workspace = true
rusqlite = { version = "0.39", features = ["bundled"] }
serde-pickle = "1.2"
serde_json.workspace = true
thiserror.workspace = true
tokio.workspace = true
[dev-dependencies]

8
litellm-rust/crates/cache-disk/src/cache.rs
View file

@ -64,7 +64,7 @@ impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> DiskCache<S, D, A> {
{
tokio::task::spawn_blocking(move || operation(&store))
.await
.map_err(|_| Error::Unavailable)?
.map_err(Error::unavailable)?
}
}
@ -162,7 +162,7 @@ impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> BatchCache for DiskCache<S, D
.map(|key| match self.get_cache(key, context) {
Ok(Some(value)) => Ok(BatchEntry::Hit(value)),
Ok(None) => Ok(BatchEntry::Miss),
Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid),
Err(Error::InvalidEntry(_)) => Ok(BatchEntry::Invalid),
Err(error) => Err(error),
})
.collect()
@ -186,7 +186,7 @@ impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> BatchCache for DiskCache<S, D
Some(value) => match self.decode_stored(value) {
Ok(Some(value)) => Ok(BatchEntry::Hit(value)),
Ok(None) => Ok(BatchEntry::Miss),
Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid),
Err(Error::InvalidEntry(_)) => Ok(BatchEntry::Invalid),
Err(error) => Err(error),
},
})
@ -272,7 +272,7 @@ fn increment<A: ValueAdapter, D: DiskStore>(
Ok((stored, ttl.map(|ttl| unix_now() + ttl.as_secs_f64())))
};
store.update(key, unix_now(), &mut apply)?;
result.ok_or(Error::InvalidEntry)
result.ok_or(Error::InvalidEntry(None))
}
fn unix_now() -> f64 {

18
litellm-rust/crates/cache-disk/src/python/value.rs
View file

@ -3,8 +3,8 @@ use py_literal::Value;
use serde_json::{Map, Number};
pub(crate) fn from_pickle(bytes: &[u8]) -> Result<Value, Error> {
let value = serde_pickle::value_from_slice(bytes, Default::default())
.map_err(|_| Error::InvalidEntry)?;
let value =
serde_pickle::value_from_slice(bytes, Default::default()).map_err(Error::invalid_entry)?;
from_pickle_value(value)
}
@ -70,7 +70,7 @@ fn from_pickle_hashable(value: serde_pickle::HashableValue) -> Result<Value, Err
}
fn integer(value: String) -> Result<Value, Error> {
value.parse().map_err(|_| Error::InvalidEntry)
value.parse().map_err(Error::invalid_entry)
}
pub(crate) fn from_json(value: serde_json::Value) -> Value {
@ -101,7 +101,7 @@ pub(crate) fn from_json(value: serde_json::Value) -> Value {
pub(crate) fn from_json_text(value: &str) -> Result<Value, Error> {
serde_json::from_str(value)
.map(from_json)
.map_err(|_| Error::InvalidEntry)
.map_err(Error::invalid_entry)
}
pub(crate) fn is_truthy(value: &Value) -> bool {
@ -131,7 +131,7 @@ pub(crate) fn to_f64(value: &Value) -> Option<f64> {
}
pub(crate) fn to_json(value: &Value) -> Result<Vec<u8>, Error> {
serde_json::to_vec(&to_json_value(value)?).map_err(|_| Error::InvalidEntry)
serde_json::to_vec(&to_json_value(value)?).map_err(Error::invalid_entry)
}
fn to_json_value(value: &Value) -> Result<serde_json::Value, Error> {
@ -142,12 +142,12 @@ fn to_json_value(value: &Value) -> Result<serde_json::Value, Error> {
value
.to_string()
.parse::<Number>()
.map_err(|_| Error::InvalidEntry)?,
.map_err(Error::invalid_entry)?,
),
Value::Float(value) => {
serde_json::Value::Number(Number::from_f64(*value).ok_or(Error::InvalidEntry)?)
serde_json::Value::Number(Number::from_f64(*value).ok_or(Error::InvalidEntry(None))?)
}
Value::Complex(_) | Value::Bytes(_) => return Err(Error::InvalidEntry),
Value::Complex(_) | Value::Bytes(_) => return Err(Error::InvalidEntry(None)),
Value::String(value) => serde_json::Value::String(value.clone()),
Value::Tuple(values) | Value::List(values) | Value::Set(values) => {
serde_json::Value::Array(
@ -162,7 +162,7 @@ fn to_json_value(value: &Value) -> Result<serde_json::Value, Error> {
.iter()
.map(|(key, value)| {
let Value::String(key) = key else {
return Err(Error::InvalidEntry);
return Err(Error::InvalidEntry(None));
};
Ok((key.clone(), to_json_value(value)?))
})

168
litellm-rust/crates/cache-disk/src/sqlite.rs
View file

@ -48,13 +48,13 @@ struct Row {
impl DiskcacheSqliteStore {
pub fn open(directory: impl AsRef<Path>) -> Result<Self, Error> {
let directory = directory.as_ref().to_path_buf();
fs::create_dir_all(&directory).map_err(|_| Error::Unavailable)?;
let directory = std::path::absolute(&directory).map_err(|_| Error::Unavailable)?;
fs::create_dir_all(&directory).map_err(Error::unavailable)?;
let directory = std::path::absolute(&directory).map_err(Error::unavailable)?;
let database = directory.join("cache.db");
let connection = Connection::open(database).map_err(|_| Error::Unavailable)?;
let connection = Connection::open(database).map_err(Error::unavailable)?;
connection
.busy_timeout(std::time::Duration::from_secs(60))
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
let mut settings = read_settings(&connection)?;
for (key, value) in default_settings() {
@ -74,7 +74,7 @@ impl DiskcacheSqliteStore {
value
)",
)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
for (key, value) in &settings {
if !matches!(key.as_str(), "count" | "size" | "hits" | "misses") {
connection
@ -82,7 +82,7 @@ impl DiskcacheSqliteStore {
"INSERT OR REPLACE INTO Settings VALUES (?, ?)",
params![key, value],
)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
}
}
for (key, value) in [
@ -96,7 +96,7 @@ impl DiskcacheSqliteStore {
"INSERT OR IGNORE INTO Settings VALUES (?, ?)",
params![key, value],
)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
}
connection
.execute_batch(
@ -117,7 +117,7 @@ impl DiskcacheSqliteStore {
CREATE UNIQUE INDEX IF NOT EXISTS Cache_key_raw ON Cache(key, raw);
CREATE INDEX IF NOT EXISTS Cache_expire_time ON Cache(expire_time);",
)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
let eviction_policy = setting_string(&settings, "eviction_policy")
.unwrap_or_else(|| "least-recently-stored".to_string());
@ -128,23 +128,27 @@ impl DiskcacheSqliteStore {
.execute_batch(
"CREATE INDEX IF NOT EXISTS Cache_store_time ON Cache(store_time)",
)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
}
"least-recently-used" => {
connection
.execute_batch(
"CREATE INDEX IF NOT EXISTS Cache_access_time ON Cache(access_time)",
)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
}
"least-frequently-used" => {
connection
.execute_batch(
"CREATE INDEX IF NOT EXISTS Cache_access_count ON Cache(access_count)",
)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
}
_ => {
return Err(Error::unavailable(SettingFault::EvictionPolicy(
eviction_policy,
)));
}
_ => return Err(Error::Unavailable),
}
connection
.execute_batch(
@ -170,12 +174,12 @@ impl DiskcacheSqliteStore {
UPDATE Settings SET value = value - OLD.size
WHERE key = \"size\"; END;",
)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
let min_file_size = setting_i64(&settings, "disk_min_file_size")
.unwrap_or(DEFAULT_DISK_MIN_FILE_SIZE)
.try_into()
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
let size_limit = setting_i64(&settings, "size_limit").unwrap_or(DEFAULT_SIZE_LIMIT);
let cull_limit = setting_i64(&settings, "cull_limit").unwrap_or(DEFAULT_CULL_LIMIT);
let statistics = setting_i64(&settings, "statistics").unwrap_or_default() != 0;
@ -207,7 +211,7 @@ impl DiskcacheSqliteStore {
|row| row.get::<_, Option<String>>(0),
)
.optional()
.map_err(|_| Error::Unavailable)?
.map_err(Error::unavailable)?
.flatten()
{
cleanup.push(old_filename);
@ -221,7 +225,7 @@ impl DiskcacheSqliteStore {
|row| row.get::<_, i64>(0),
)
.optional()
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
if let Some(rowid) = rowid {
connection
.execute(
@ -230,7 +234,7 @@ impl DiskcacheSqliteStore {
WHERE rowid = ?",
params![now, expire_time, now, size, mode, filename, value, rowid],
)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
} else {
connection
.execute(
@ -240,7 +244,7 @@ impl DiskcacheSqliteStore {
) VALUES (?, 1, ?, ?, ?, 0, NULL, ?, ?, ?, ?)",
params![key, now, expire_time, now, size, mode, filename, value],
)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
}
cleanup.extend(self.cull(connection, now)?);
Ok(cleanup)
@ -257,13 +261,13 @@ impl DiskcacheSqliteStore {
WHERE expire_time IS NOT NULL AND expire_time < ?
ORDER BY expire_time LIMIT ?",
)
.map_err(|_| Error::Unavailable)?
.map_err(Error::unavailable)?
.query_map(params![now, self.cull_limit], |row| {
Ok((row.get::<_, i64>(0)?, row.get::<_, Option<String>>(1)?))
})
.map_err(|_| Error::Unavailable)?
.map_err(Error::unavailable)?
.collect::<Result<Vec<_>, _>>()
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
for (_, filename) in &expired {
if let Some(filename) = filename {
cleanup.push(filename.clone());
@ -272,7 +276,7 @@ impl DiskcacheSqliteStore {
for (rowid, _) in &expired {
connection
.execute("DELETE FROM Cache WHERE rowid = ?", params![rowid])
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
}
let remaining = self.cull_limit - i64::try_from(expired.len()).unwrap_or(self.cull_limit);
if remaining <= 0 || self.volume(connection)? < self.size_limit {
@ -283,19 +287,23 @@ impl DiskcacheSqliteStore {
"least-recently-stored" => "store_time",
"least-recently-used" => "access_time",
"least-frequently-used" => "access_count",
_ => return Err(Error::Unavailable),
policy => {
return Err(Error::unavailable(SettingFault::EvictionPolicy(
policy.to_owned(),
)));
}
};
let rows = connection
.prepare(&format!(
"SELECT rowid, filename FROM Cache ORDER BY {order} LIMIT ?"
))
.map_err(|_| Error::Unavailable)?
.map_err(Error::unavailable)?
.query_map(params![remaining], |row| {
Ok((row.get::<_, i64>(0)?, row.get::<_, Option<String>>(1)?))
})
.map_err(|_| Error::Unavailable)?
.map_err(Error::unavailable)?
.collect::<Result<Vec<_>, _>>()
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
for (_, filename) in &rows {
if let Some(filename) = filename {
cleanup.push(filename.clone());
@ -304,7 +312,7 @@ impl DiskcacheSqliteStore {
for (rowid, _) in rows {
connection
.execute("DELETE FROM Cache WHERE rowid = ?", params![rowid])
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
}
Ok(cleanup)
}
@ -312,15 +320,15 @@ impl DiskcacheSqliteStore {
fn volume(&self, connection: &Connection) -> Result<i64, Error> {
let page_count: i64 = connection
.query_row("PRAGMA page_count", [], |row| row.get(0))
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
let page_size: i64 = connection
.query_row("PRAGMA page_size", [], |row| row.get(0))
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
let size: i64 = connection
.query_row("SELECT value FROM Settings WHERE key = 'size'", [], |row| {
row.get(0)
})
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
Ok(page_count.saturating_mul(page_size).saturating_add(size))
}
}
@ -331,13 +339,13 @@ impl DiskStore for DiskcacheSqliteStore {
}
fn get(&self, key: &str, now: f64) -> Result<Option<StoredValue>, Error> {
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
let connection = self.connection.lock()?;
let select = "SELECT rowid, expire_time, mode, filename, value FROM Cache
WHERE key = ? AND raw = 1 AND (expire_time IS NULL OR expire_time > ?)";
let row = connection
.query_row(select, params![key, now], row_from_query)
.optional()
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
if !self.statistics && !has_get_update(&self.eviction_policy) {
return row
.map(|row| fetch_row(&self.directory, row))
@ -348,7 +356,7 @@ impl DiskStore for DiskcacheSqliteStore {
let row = connection
.query_row(select, params![key, now], row_from_query)
.optional()
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
let Some(row) = row else {
if self.statistics {
connection
@ -356,7 +364,7 @@ impl DiskStore for DiskcacheSqliteStore {
"UPDATE Settings SET value = value + 1 WHERE key = 'misses'",
[],
)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
}
return Ok(None);
};
@ -369,14 +377,14 @@ impl DiskStore for DiskcacheSqliteStore {
"UPDATE Settings SET value = value + 1 WHERE key = 'hits'",
[],
)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
} else if !hit && self.statistics {
connection
.execute(
"UPDATE Settings SET value = value + 1 WHERE key = 'misses'",
[],
)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
}
if has_get_update(&self.eviction_policy) && hit {
let update = match self.eviction_policy.as_str() {
@ -384,16 +392,20 @@ impl DiskStore for DiskcacheSqliteStore {
"least-frequently-used" => {
"UPDATE Cache SET access_count = access_count + 1 WHERE rowid = ?"
}
_ => return Err(Error::Unavailable),
policy => {
return Err(Error::unavailable(SettingFault::EvictionPolicy(
policy.to_owned(),
)));
}
};
if self.eviction_policy == "least-recently-used" {
connection
.execute(update, params![now, rowid])
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
} else {
connection
.execute(update, params![rowid])
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
}
}
value
@ -409,7 +421,7 @@ impl DiskStore for DiskcacheSqliteStore {
) -> Result<(), Error> {
let columns = store_value(&self.directory, self.min_file_size, value)?;
let new_filename = columns.filename.clone();
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
let connection = self.connection.lock()?;
let result = transactional(&connection, |connection| {
self.set_locked(connection, key, columns, expire_time, now)
});
@ -428,7 +440,7 @@ impl DiskStore for DiskcacheSqliteStore {
}
fn pop(&self, key: &str, now: f64) -> Result<Option<StoredValue>, Error> {
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
let connection = self.connection.lock()?;
let selected = transactional(&connection, |connection| {
let row = connection
.query_row(
@ -439,13 +451,13 @@ impl DiskStore for DiskcacheSqliteStore {
row_from_query,
)
.optional()
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
let Some(row) = row else {
return Ok(None);
};
connection
.execute("DELETE FROM Cache WHERE rowid = ?", params![row.rowid])
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
Ok(Some(row))
})?;
let Some(row) = selected else {
@ -460,7 +472,7 @@ impl DiskStore for DiskcacheSqliteStore {
}
fn clear(&self) -> Result<(), Error> {
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
let connection = self.connection.lock()?;
let mut last_rowid = 0_i64;
loop {
let batch = transactional(&connection, |connection| {
@ -469,13 +481,13 @@ impl DiskStore for DiskcacheSqliteStore {
"SELECT rowid, filename FROM Cache
WHERE rowid > ? ORDER BY rowid LIMIT 100",
)
.map_err(|_| Error::Unavailable)?
.map_err(Error::unavailable)?
.query_map(params![last_rowid], |row| {
Ok((row.get::<_, i64>(0)?, row.get::<_, Option<String>>(1)?))
})
.map_err(|_| Error::Unavailable)?
.map_err(Error::unavailable)?
.collect::<Result<Vec<_>, _>>()
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
if rows.is_empty() {
return Ok(rows);
}
@ -486,7 +498,7 @@ impl DiskStore for DiskcacheSqliteStore {
.join(",");
connection
.execute(&format!("DELETE FROM Cache WHERE rowid IN ({ids})"), [])
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
Ok(rows)
})?;
if batch.is_empty() {
@ -509,7 +521,7 @@ impl DiskStore for DiskcacheSqliteStore {
now: f64,
apply: &mut dyn FnMut(Option<StoredValue>) -> Result<(StoredValue, Option<f64>), Error>,
) -> Result<(), Error> {
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
let connection = self.connection.lock()?;
let mut created_filename = None;
let result = transactional(&connection, |connection| {
let current = connection
@ -521,7 +533,7 @@ impl DiskStore for DiskcacheSqliteStore {
row_from_query,
)
.optional()
.map_err(|_| Error::Unavailable)?
.map_err(Error::unavailable)?
.map(|row| fetch_row(&self.directory, row))
.transpose()?
.flatten();
@ -582,21 +594,35 @@ fn read_settings(connection: &Connection) -> Result<HashMap<String, Value>, Erro
};
statement
.query_map([], |row| Ok((row.get(0)?, row.get(1)?)))
.map_err(|_| Error::Unavailable)?
.map_err(Error::unavailable)?
.collect::<Result<HashMap<_, _>, _>>()
.map_err(|_| Error::Unavailable)
.map_err(Error::unavailable)
}
#[derive(Debug, thiserror::Error)]
enum SettingFault {
#[error("unknown eviction policy {0:?}")]
EvictionPolicy(String),
#[error("setting {0:?} is not an SQLite pragma")]
NotPragma(String),
#[error("setting {0:?} has an unsupported value type")]
PragmaValue(String),
}
fn apply_pragma(connection: &Connection, key: &str, value: &Value) -> Result<(), Error> {
let pragma = key.strip_prefix("sqlite_").ok_or(Error::Unavailable)?;
let pragma = key
.strip_prefix("sqlite_")
.ok_or_else(|| Error::unavailable(SettingFault::NotPragma(key.to_owned())))?;
match value {
Value::Integer(value) => connection
.pragma_update(None, pragma, value)
.map_err(|_| Error::Unavailable),
.map_err(Error::unavailable),
Value::Text(value) => connection
.pragma_update(None, pragma, value)
.map_err(|_| Error::Unavailable),
_ => Err(Error::Unavailable),
.map_err(Error::unavailable),
_ => Err(Error::unavailable(SettingFault::PragmaValue(
key.to_owned(),
))),
}
}
@ -634,7 +660,7 @@ fn fetch_row(directory: &Path, row: Row) -> Result<Option<StoredValue>, Error> {
Value::Text(value) => Ok(Some(StoredValue::Text(value))),
Value::Integer(value) => Ok(Some(StoredValue::Integer(value))),
Value::Real(value) => Ok(Some(StoredValue::Float(value))),
Value::Null => Err(Error::InvalidEntry),
Value::Null => Err(Error::InvalidEntry(None)),
},
MODE_BINARY | MODE_PICKLE => {
let bytes = match row.value {
@ -645,7 +671,7 @@ fn fetch_row(directory: &Path, row: Row) -> Result<Option<StoredValue>, Error> {
};
value
}
_ => return Err(Error::InvalidEntry),
_ => return Err(Error::InvalidEntry(None)),
};
Ok(Some(if row.mode == MODE_BINARY {
StoredValue::Bytes(bytes)
@ -663,24 +689,24 @@ fn fetch_row(directory: &Path, row: Row) -> Result<Option<StoredValue>, Error> {
}
Value::Blob(value) => value,
Value::Text(value) => value.into_bytes(),
_ => return Err(Error::InvalidEntry),
_ => return Err(Error::InvalidEntry(None)),
};
Ok(Some(StoredValue::Text(
String::from_utf8(bytes).map_err(|_| Error::InvalidEntry)?,
String::from_utf8(bytes).map_err(Error::invalid_entry)?,
)))
}
_ => Err(Error::InvalidEntry),
_ => Err(Error::InvalidEntry(None)),
}
}
fn read_file(directory: &Path, filename: Option<&str>) -> Result<Option<Vec<u8>>, Error> {
let Some(filename) = filename else {
return Err(Error::InvalidEntry);
return Err(Error::InvalidEntry(None));
};
match fs::read(directory.join(filename)) {
Ok(value) => Ok(Some(value)),
Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None),
Err(_) => Err(Error::Unavailable),
Err(error) => Err(Error::unavailable(error)),
}
}
@ -712,7 +738,7 @@ fn store_value(
let bytes = value.into_bytes();
let filename = write_file(directory, &bytes)?;
Ok(StoredColumns {
size: i64::try_from(bytes.len()).map_err(|_| Error::Unavailable)?,
size: i64::try_from(bytes.len()).map_err(Error::unavailable)?,
mode: MODE_TEXT,
filename: Some(filename),
value: None,
@ -727,7 +753,7 @@ fn store_value(
StoredValue::Bytes(value) => {
let filename = write_file(directory, &value)?;
Ok(StoredColumns {
size: i64::try_from(value.len()).map_err(|_| Error::Unavailable)?,
size: i64::try_from(value.len()).map_err(Error::unavailable)?,
mode: MODE_BINARY,
filename: Some(filename),
value: None,
@ -742,7 +768,7 @@ fn store_value(
StoredValue::Pickle(value) => {
let filename = write_file(directory, &value)?;
Ok(StoredColumns {
size: i64::try_from(value.len()).map_err(|_| Error::Unavailable)?,
size: i64::try_from(value.len()).map_err(Error::unavailable)?,
mode: MODE_PICKLE,
filename: Some(filename),
value: None,
@ -761,14 +787,14 @@ fn write_file(directory: &Path, bytes: &[u8]) -> Result<String, Error> {
let filename = format!("{}/{}/{}.val", &hex[..2], &hex[2..4], &hex[4..]);
let path = directory.join(&filename);
if let Some(parent) = path.parent() {
fs::create_dir_all(parent).map_err(|_| Error::Unavailable)?;
fs::create_dir_all(parent).map_err(Error::unavailable)?;
}
let mut file = OpenOptions::new()
.write(true)
.create_new(true)
.open(path)
.map_err(|_| Error::Unavailable)?;
file.write_all(bytes).map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
file.write_all(bytes).map_err(Error::unavailable)?;
Ok(filename)
}
@ -789,12 +815,12 @@ fn transactional<T>(
) -> Result<T, Error> {
connection
.execute_batch("BEGIN IMMEDIATE")
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
match operation(connection) {
Ok(value) => {
connection
.execute_batch("COMMIT")
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
Ok(value)
}
Err(error) => {

2
litellm-rust/crates/cache-disk/tests/python_compat.rs
View file

@ -84,7 +84,7 @@ fn python_read_cases(#[case] row: StoredValue, #[case] expected: ReadExpectation
match expected {
ReadExpectation::Bytes(expected) => assert_eq!(result.unwrap().unwrap(), expected),
ReadExpectation::Miss => assert_eq!(result.unwrap(), None),
ReadExpectation::Invalid => assert!(matches!(result, Err(Error::InvalidEntry))),
ReadExpectation::Invalid => assert!(matches!(result, Err(Error::InvalidEntry(_)))),
}
}

2
litellm-rust/crates/cache-gcs/Cargo.toml
View file

@ -13,9 +13,11 @@ litellm-auth-types.workspace = true
litellm-cache.workspace = true
percent-encoding.workspace = true
reqwest.workspace = true
thiserror.workspace = true
tokio.workspace = true
[dev-dependencies]
litellm-error.workspace = true
litellm-http = { workspace = true, features = ["test-support"] }
litellm-cache-testing.workspace = true
rstest.workspace = true

32
litellm-rust/crates/cache-gcs/src/cache.rs
View file

@ -25,6 +25,14 @@ pub fn key_prefix(gcs_path: Option<&str>) -> String {
}
}
#[derive(Debug, thiserror::Error)]
#[error("GCS answered {0}")]
struct UnexpectedStatus(reqwest::StatusCode);
#[derive(Debug, thiserror::Error)]
#[error("GCS cache worker thread panicked")]
struct WorkerPanicked;
#[derive(Clone, Debug)]
pub struct GcsConfig {
pub bucket_name: String,
@ -114,9 +122,9 @@ impl<S: CacheCodec> GcsCache<S> {
.body(payload)
.send()
.await
.map_err(|_| Error::Unavailable)?;
.map_err(|error| Error::unavailable(litellm_http::Error::from(error)))?;
if !response.status().is_success() {
return Err(Error::Unavailable);
return Err(Error::unavailable(UnexpectedStatus(response.status())));
}
Ok(())
}
@ -135,18 +143,18 @@ impl<S: CacheCodec> GcsCache<S> {
.header(reqwest::header::CONTENT_TYPE, "application/json")
.send()
.await
.map_err(|_| Error::Unavailable)?;
.map_err(|error| Error::unavailable(litellm_http::Error::from(error)))?;
if response.status() == reqwest::StatusCode::NOT_FOUND {
return Ok(None);
}
if !response.status().is_success() {
return Err(Error::Unavailable);
return Err(Error::unavailable(UnexpectedStatus(response.status())));
}
let body = response.bytes().await.map_err(|_| Error::Unavailable)?;
self.codec
.decode(&body)
.map(Some)
.map_err(|_| Error::InvalidEntry)
let body = response
.bytes()
.await
.map_err(|error| Error::unavailable(litellm_http::Error::from(error)))?;
self.codec.decode(&body).map(Some)
}
fn run_sync<T, F>(future: F) -> Result<T, Error>
@ -158,7 +166,7 @@ impl<S: CacheCodec> GcsCache<S> {
tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.map_err(|_| Error::Unavailable)
.map_err(Error::unavailable)
.and_then(|runtime| runtime.block_on(future))
};
match tokio::runtime::Handle::try_current() {
@ -169,7 +177,7 @@ impl<S: CacheCodec> GcsCache<S> {
scope
.spawn(|| run(future))
.join()
.map_err(|_| Error::Unavailable)
.map_err(|_| Error::unavailable(WorkerPanicked))
.and_then(|result| result)
}),
Err(_) => run(future),
@ -242,7 +250,7 @@ impl<S: CacheCodec> BatchCache for GcsCache<S> {
match self.async_get_cache(&key, &context).await {
Ok(Some(value)) => Ok(BatchEntry::Hit(value)),
Ok(None) => Ok(BatchEntry::Miss),
Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid),
Err(Error::InvalidEntry(_)) => Ok(BatchEntry::Invalid),
Err(error) => Err(error),
}
}

2
litellm-rust/crates/cache-gcs/src/token.rs
View file

@ -30,7 +30,7 @@ impl TokenSource for GcpTokenSource {
self.auth
.access_token(&self.config, &|name| std::env::var(name).ok())
.await
.map_err(|_| Error::Unavailable)
.map_err(Error::unavailable)
})
}
}

39
litellm-rust/crates/cache-gcs/tests/cache.rs
View file

@ -55,19 +55,27 @@ async fn set_writes_encoded_object_and_headers(#[future(awt)] server: MockServer
Ok(Some(json!({"value": "entry"})))
)]
#[case::missing("missing", ResponseTemplate::new(404), Ok(None))]
#[case::server_error("server-error", ResponseTemplate::new(500), Err(Error::Unavailable))]
#[case::unauthorized("unauthorized", ResponseTemplate::new(401), Err(Error::Unavailable))]
#[case::server_error(
"server-error",
ResponseTemplate::new(500),
Err("cache is unavailable: GCS answered 500 Internal Server Error")
)]
#[case::unauthorized(
"unauthorized",
ResponseTemplate::new(401),
Err("cache is unavailable: GCS answered 401 Unauthorized")
)]
#[case::invalid(
"invalid",
ResponseTemplate::new(200).set_body_string("not json"),
Err(Error::InvalidEntry)
Err("invalid cache entry: ")
)]
#[tokio::test]
async fn get_maps_statuses_and_decode_failures(
#[future(awt)] server: MockServer,
#[case] key: &str,
#[case] response: ResponseTemplate,
#[case] expected: Result<Option<Value>, Error>,
#[case] expected: Result<Option<Value>, &str>,
) {
Mock::given(method("GET"))
.and(path(format!("/storage/v1/b/bucket/o/{key}")))
@ -76,8 +84,21 @@ async fn get_maps_statuses_and_decode_failures(
.mount(&server)
.await;
let cache = support::cache(&server, None);
assert_eq!(cache.get_cache(key, &context()), expected);
assert_eq!(cache.async_get_cache(key, &context()).await, expected);
let outcome = |result: Result<Option<Value>, Error>| {
result.map_err(|error| litellm_error::Report::new(&error).to_string())
};
let sync = outcome(cache.get_cache(key, &context()));
let asynchronous = outcome(cache.async_get_cache(key, &context()).await);
match expected {
Ok(value) => {
assert_eq!(sync, Ok(value.clone()));
assert_eq!(asynchronous, Ok(value));
}
Err(prefix) => {
assert!(sync.unwrap_err().starts_with(prefix));
assert!(asynchronous.unwrap_err().starts_with(prefix));
}
}
}
#[rstest]
@ -257,7 +278,7 @@ struct FailingTokenSource;
impl TokenSource for FailingTokenSource {
fn bearer_token(&self) -> Pin<Box<dyn Future<Output = Result<String, Error>> + Send + '_>> {
Box::pin(async { Err(Error::Unavailable) })
Box::pin(async { Err(Error::Unavailable(None)) })
}
}
@ -267,14 +288,14 @@ async fn token_source_failure_skips_http(#[future(awt)] server: MockServer) {
let cache = support::cache_with_token(&server, None, Arc::new(FailingTokenSource));
assert_eq!(
cache.get_cache("key", &context()).unwrap_err(),
Error::Unavailable
Error::Unavailable(None)
);
assert_eq!(
cache
.async_set_cache("key", json!(1), context())
.await
.unwrap_err(),
Error::Unavailable
Error::Unavailable(None)
);
assert_eq!(server.received_requests().await.unwrap().len(), 0);
}

48
litellm-rust/crates/cache-memory/src/cache.rs
View file

@ -2,7 +2,7 @@ use std::{
cmp::Reverse,
collections::{BinaryHeap, HashMap, HashSet},
hash::Hash,
sync::{Arc, Mutex},
sync::{Arc, Mutex, MutexGuard},
time::{Duration, SystemTime, UNIX_EPOCH},
};
@ -84,6 +84,10 @@ impl<V: Clone> InMemoryCache<V> {
}
}
fn lock(&self) -> Result<MutexGuard<'_, CacheState<V>>, Error> {
Ok(self.state.lock()?)
}
pub fn set_cache(
&self,
key: impl Into<String>,
@ -94,13 +98,13 @@ impl<V: Clone> InMemoryCache<V> {
return Ok(CacheWrite::Disabled);
}
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
let mut state = self.lock()?;
self.store(&mut state, key.into(), value, ttl, now)
}
pub fn get_cache(&self, key: &str) -> Result<Option<V>, Error> {
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
let mut state = self.lock()?;
if state
.expirations
.get(key)
@ -124,7 +128,7 @@ impl<V: Clone> InMemoryCache<V> {
/// fits.
pub fn evict_cache(&self) -> Result<(), Error> {
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
let mut state = self.lock()?;
Self::evict(&mut state, self.max_size_in_memory, now, None);
Ok(())
}
@ -132,7 +136,7 @@ impl<V: Clone> InMemoryCache<V> {
/// `evict_element_if_expired`: `true` when `key` had expired and was removed.
pub fn evict_element_if_expired(&self, key: &str) -> Result<bool, Error> {
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
let mut state = self.lock()?;
let expired = state
.expirations
.get(key)
@ -153,12 +157,7 @@ impl<V: Clone> InMemoryCache<V> {
/// The number of stored entries, expired ones included until they are evicted.
pub fn len(&self) -> Result<usize, Error> {
Ok(self
.state
.lock()
.map_err(|_| Error::Unavailable)?
.values
.len())
Ok(self.lock()?.values.len())
}
pub fn is_empty(&self) -> Result<bool, Error> {
@ -167,12 +166,7 @@ impl<V: Clone> InMemoryCache<V> {
/// Entries in the expiration heap, stale ones included; bounded by eviction.
pub fn expiration_heap_len(&self) -> Result<usize, Error> {
Ok(self
.state
.lock()
.map_err(|_| Error::Unavailable)?
.expiration_heap
.len())
Ok(self.lock()?.expiration_heap.len())
}
pub fn max_size_in_memory(&self) -> usize {
@ -184,13 +178,7 @@ impl<V: Clone> InMemoryCache<V> {
}
pub fn expires_at(&self, key: &str) -> Result<Option<Duration>, Error> {
Ok(self
.state
.lock()
.map_err(|_| Error::Unavailable)?
.expirations
.get(key)
.copied())
Ok(self.lock()?.expirations.get(key).copied())
}
pub async fn async_get_ttl(&self, key: &str) -> Result<Option<Duration>, Error> {
@ -198,7 +186,7 @@ impl<V: Clone> InMemoryCache<V> {
}
pub async fn async_get_oldest_n_keys(&self, count: usize) -> Result<Vec<String>, Error> {
let state = self.state.lock().map_err(|_| Error::Unavailable)?;
let state = self.lock()?;
let mut expirations = state
.expirations
.iter()
@ -213,13 +201,13 @@ impl<V: Clone> InMemoryCache<V> {
}
pub fn delete_cache(&self, key: &str) -> Result<(), Error> {
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
let mut state = self.lock()?;
Self::remove(&mut state, key);
Ok(())
}
pub fn flush_cache(&self) -> Result<(), Error> {
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
let mut state = self.lock()?;
state.values.clear();
state.expirations.clear();
state.expiration_heap.clear();
@ -316,7 +304,7 @@ where
return Ok(candidate);
}
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
let mut state = self.lock()?;
Self::evict(&mut state, self.max_size_in_memory, now, Some(key));
let existing = state
.values
@ -351,7 +339,7 @@ impl CounterCache for InMemoryCache<f64> {
return Ok(amount);
}
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
let mut state = self.lock()?;
let value = Self::live(&mut state, key, now).unwrap_or_default() + amount;
self.store(&mut state, key.into(), value, self.get_ttl(&context), now)?;
Ok(value)
@ -424,7 +412,7 @@ where
return Ok(values);
}
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
let mut state = self.lock()?;
let mut stored = Self::live(&mut state, key, now).unwrap_or_default();
stored.extend(values.iter().cloned());
self.store(&mut state, key.into(), stored, ttl, now)?;

6
litellm-rust/crates/cache-memory/tests/cache.rs
View file

@ -52,7 +52,7 @@ fn measured(capacity: usize) -> InMemoryCache<String> {
Some(4),
Some(Arc::new(|value: &String| {
if value.is_empty() {
return Err(Error::InvalidEntry);
return Err(Error::InvalidEntry(None));
}
Ok(value.len())
})),
@ -275,7 +275,7 @@ fn evict_element_if_expired_reports_removal(clock: Clock) {
#[case::fits("ok", Ok(true))]
#[case::at_limit("four", Ok(true))]
#[case::too_large("oversized", Ok(false))]
#[case::measure_error("", Err(Error::InvalidEntry))]
#[case::measure_error("", Err(Error::InvalidEntry(None)))]
fn check_value_size_applies_the_entry_limit(
#[case] value: &str,
#[case] expected: Result<bool, Error>,
@ -310,7 +310,7 @@ fn disabled_size_limited_and_validated_writes_are_observable() {
assert_eq!(cache.get_cache("small").unwrap(), Some("ok".into()));
assert_eq!(
cache.set_cache("invalid", String::new(), None),
Err(Error::InvalidEntry)
Err(Error::InvalidEntry(None))
);
assert_eq!(cache.get_cache("invalid").unwrap(), None);
cache.delete_cache("small").unwrap();

1
litellm-rust/crates/cache-qdrant-semantic/Cargo.toml
View file

@ -18,6 +18,7 @@ tokio.workspace = true
uuid.workspace = true
[dev-dependencies]
litellm-error.workspace = true
litellm-http = { workspace = true, features = ["test-support"] }
futures-executor = "0.3"
litellm-cache-testing.workspace = true

14
litellm-rust/crates/cache-qdrant-semantic/src/cache.rs
View file

@ -36,7 +36,7 @@ impl<E: Embedder, C: CacheCodec> QdrantSemanticCache<E, C> {
let exists = client
.collection_exists(config.collection_name.clone())
.await
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
if !exists {
client
.create_collection(
@ -48,7 +48,7 @@ impl<E: Embedder, C: CacheCodec> QdrantSemanticCache<E, C> {
.quantization_config(quantization(&config.quantization)),
)
.await
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
}
let _ = client
.create_field_index(CreateFieldIndexCollectionBuilder::new(
@ -99,13 +99,13 @@ impl<E: Embedder, C: CacheCodec> QdrantSemanticCache<E, C> {
.async_embed(&prompt, context.metadata.as_ref())
.await?;
let response =
String::from_utf8(self.codec.encode(&value)?).map_err(|_| Error::InvalidEntry)?;
String::from_utf8(self.codec.encode(&value)?).map_err(Error::invalid_entry)?;
let payload = Payload::try_from(json!({
"litellm_cache_key": key,
"text": prompt,
"response": response,
}))
.map_err(|_| Error::InvalidEntry)?;
.map_err(Error::invalid_entry)?;
self.client
.upsert_points(
UpsertPointsBuilder::new(
@ -119,7 +119,7 @@ impl<E: Embedder, C: CacheCodec> QdrantSemanticCache<E, C> {
.wait(true),
)
.await
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
Ok(())
}
@ -152,7 +152,7 @@ impl<E: Embedder, C: CacheCodec> QdrantSemanticCache<E, C> {
),
)
.await
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
let Some(point) = result.result.into_iter().next() else {
return Ok(SemanticLookup::miss(Some(0.0)));
};
@ -170,7 +170,7 @@ impl<E: Embedder, C: CacheCodec> QdrantSemanticCache<E, C> {
let response = payload
.get("response")
.and_then(Value::as_str)
.ok_or(Error::InvalidEntry)?;
.ok_or(Error::InvalidEntry(None))?;
Ok(SemanticLookup {
value: Some(self.codec.decode(response.as_bytes())?),
similarity: Some(similarity),

17
litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs
View file

@ -55,10 +55,9 @@ impl Embedder for OpenAiEmbedder {
}
.send()
.await
.map_err(|_| Error::Unavailable)?
.error_for_status()
.map_err(|_| Error::Unavailable)?;
let body: Value = response.json().await.map_err(|_| Error::Unavailable)?;
.and_then(reqwest::Response::error_for_status)
.map_err(http_error)?;
let body: Value = response.json().await.map_err(http_error)?;
body.get("data")
.and_then(Value::as_array)
.and_then(|data| data.first())
@ -70,6 +69,14 @@ impl Embedder for OpenAiEmbedder {
.map(|value| value.as_f64().map(|value| value as f32))
.collect::<Option<Vec<_>>>()
})
.ok_or(Error::Unavailable)
.ok_or(Error::unavailable(MissingEmbedding))
}
}
fn http_error(error: reqwest::Error) -> Error {
Error::unavailable(litellm_http::Error::from(error))
}
#[derive(Debug, thiserror::Error)]
#[error("embeddings response has no embedding")]
struct MissingEmbedding;

21
litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs
View file

@ -134,13 +134,19 @@ async fn posts_embeddings_request_and_parses_vector() {
}
#[rstest]
#[case::error_status("500 Internal Server Error", "{}", 0, None, Err(Error::Unavailable))]
#[case::error_status(
"500 Internal Server Error",
"{}",
0,
None,
Err(Error::Unavailable(None))
)]
#[case::timed_out(
"200 OK",
r#"{"data":[{"embedding":[0.1,0.2]}]}"#,
500,
Some(Duration::from_millis(200)),
Err(Error::Unavailable)
Err(Error::Unavailable(None))
)]
#[case::within_timeout(
"200 OK",
@ -149,7 +155,7 @@ async fn posts_embeddings_request_and_parses_vector() {
Some(Duration::from_secs(1)),
Ok(vec![0.1, 0.2])
)]
#[case::missing_embedding("200 OK", r#"{"data":[]}"#, 0, None, Err(Error::Unavailable))]
#[case::missing_embedding("200 OK", r#"{"data":[]}"#, 0, None, Err(Error::Unavailable(None)))]
#[tokio::test]
async fn status_timeout_and_body_errors_are_unavailable(
#[case] status: &str,
@ -164,7 +170,14 @@ async fn status_timeout_and_body_errors_are_unavailable(
litellm_http::Client::for_test(),
config(server.base_url(), timeout),
);
assert_eq!(embedder.async_embed("hello", None).await, expected);
let result = embedder.async_embed("hello", None).await;
if let Err(error) = &result {
let report = litellm_error::Report::new(error).to_string();
assert!(!report.contains(&server.base_url()), "{report}");
}
let kind =
|result: Result<Vec<f32>, Error>| result.map_err(|error| std::mem::discriminant(&error));
assert_eq!(kind(result), kind(expected));
}
#[rstest]

26
litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs
View file

@ -53,7 +53,10 @@ impl Embedder for FixedEmbedder {
.lock()
.unwrap()
.push((input.to_owned(), metadata.cloned()));
self.vectors.get(input).cloned().ok_or(Error::Unavailable)
self.vectors
.get(input)
.cloned()
.ok_or(Error::Unavailable(None))
}
}
@ -396,9 +399,9 @@ async fn lookup_reports_python_semantic_similarity(
#[rstest]
#[case::codec_decodes_the_payload(Some(json!("{\"a\":1}")), Ok(Some(json!({"a": 1}))))]
#[case::undecodable_response(Some(json!("not json")), Err(Error::InvalidEntry))]
#[case::non_string_response(Some(json!(1)), Err(Error::InvalidEntry))]
#[case::missing_response(None, Err(Error::InvalidEntry))]
#[case::undecodable_response(Some(json!("not json")), Err(Error::InvalidEntry(None)))]
#[case::non_string_response(Some(json!(1)), Err(Error::InvalidEntry(None)))]
#[case::missing_response(None, Err(Error::InvalidEntry(None)))]
#[tokio::test(flavor = "multi_thread")]
async fn stored_responses_go_through_the_codec(
#[case] response: Option<JsonValue>,
@ -419,9 +422,12 @@ async fn stored_responses_go_through_the_codec(
.into(),
});
let kind = |result: Result<Option<JsonValue>, Error>| {
result.map_err(|error| std::mem::discriminant(&error))
};
assert_eq!(
cache.async_get_cache("key", &context("hello")).await,
expected
kind(cache.async_get_cache("key", &context("hello")).await),
kind(expected)
);
server.stop();
}
@ -434,7 +440,7 @@ async fn embedding_failures_propagate() {
assert_eq!(
cache.async_get_cache("key", &context("unknown")).await,
Err(Error::Unavailable)
Err(Error::Unavailable(None))
);
server.stop();
}
@ -493,10 +499,10 @@ async fn stopped_qdrant_server_maps_to_unavailable() {
let cache = connect(&server, [("hello", vec![1.0, 0.0])]).await;
server.stop();
tokio::time::sleep(Duration::from_millis(50)).await;
assert_eq!(
assert!(matches!(
cache.async_get_cache("key", &context("hello")).await,
Err(Error::Unavailable)
);
Err(Error::Unavailable(Some(_)))
));
}
/// `_payload_matches_cache_key` compares `str(cached_key) == str(key)`, so a point whose stored

6
litellm-rust/crates/cache-redis-semantic/src/cache.rs
View file

@ -63,13 +63,13 @@ impl Inner {
.arg(CACHE_KEY_FIELD)
.arg(tag)
.query::<()>(connection)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
if let Some(ttl) = ttl {
redis::cmd("EXPIRE")
.arg(&hash_key)
.arg(ttl_seconds(ttl))
.query::<()>(connection)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
}
Ok(())
}
@ -111,7 +111,7 @@ impl Inner {
.arg("vector")
.arg(vector_buffer(vector))
.query::<redis::Value>(connection)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
let Some(fields) = first_document(&result) else {
return Ok(SemanticLookup::miss(Some(0.0)));
};

18
litellm-rust/crates/cache-redis-semantic/src/index.rs
View file

@ -40,10 +40,10 @@ impl Index {
Some(false) => self.isolated(connection, dims)?,
None => match create_index(connection, &self.name, dims) {
Ok(()) => self.name.clone(),
Err(_) => match index_compatible(connection, &self.name, dims)? {
Err(error) => match index_compatible(connection, &self.name, dims)? {
Some(true) => self.name.clone(),
Some(false) => self.isolated(connection, dims)?,
None => return Err(Error::Unavailable),
None => return Err(error),
},
},
};
@ -59,7 +59,7 @@ impl Index {
redis::cmd("FT.DROPINDEX")
.arg(&name)
.query::<()>(connection)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
create_index(connection, &name, dims)?;
Ok(name)
}
@ -109,7 +109,7 @@ fn create_index(connection: &mut ConnectionRef<'_>, name: &str, dims: usize) ->
.arg("SEPARATOR")
.arg(",")
.query::<()>(connection)
.map_err(|_| Error::Unavailable)
.map_err(Error::unavailable)
}
fn index_compatible(
@ -123,14 +123,18 @@ fn index_compatible(
{
Ok(info) => info,
Err(error) if unknown_index(&error) => return Ok(None),
Err(_) => return Err(Error::Unavailable),
Err(error) => return Err(Error::unavailable(error)),
};
Ok(Some(schema_compatible(&info, dims)))
}
fn unknown_index(error: &redis::RedisError) -> bool {
let message = error.to_string().to_lowercase();
message.contains("unknown") && message.contains("index")
error
.code()
.is_some_and(|code| code.eq_ignore_ascii_case("unknown"))
&& error
.detail()
.is_some_and(|detail| detail.to_ascii_lowercase().starts_with("index"))
}
struct Attribute {

30
litellm-rust/crates/cache-redis-semantic/tests/cache.rs
View file

@ -79,7 +79,7 @@ fn s(value: &str) -> redis::Value {
}
fn unknown_index_error() -> redis::RedisError {
redis::RedisError::from((redis::ErrorKind::Extension, "Unknown index name"))
support::server_error("Unknown index name")
}
fn info_missing(index: &str) -> MockCmd {
@ -320,7 +320,7 @@ fn store_writes_the_redisvl_hash(
#[case::below_distance_threshold("key1", "0.05", None, Ok(Some(entry())))]
#[case::above_distance_threshold("key1", "0.5", None, Ok(None))]
#[case::other_cache_key("other", "0.05", None, Ok(None))]
#[case::malformed_response("key1", "0.05", Some(b"not json!".as_slice()), Err(Error::InvalidEntry))]
#[case::malformed_response("key1", "0.05", Some(b"not json!".as_slice()), Err(Error::InvalidEntry(None)))]
fn lookup_applies_threshold_scope_and_codec(
#[case] stored_tag: &str,
#[case] distance: &str,
@ -343,7 +343,10 @@ fn lookup_applies_threshold_scope_and_codec(
FakeEmbedder::new(&[]),
);
assert_eq!(cache.get_cache("key1", &context), expected);
let kind = |result: Result<Option<Value>, Error>| {
result.map_err(|error| std::mem::discriminant(&error))
};
assert_eq!(kind(cache.get_cache("key1", &context)), kind(expected));
}
#[rstest]
@ -451,10 +454,7 @@ fn create_index_race_rechecks_schema_and_stores(entry: Value, context: SemanticC
info_missing(INDEX),
MockCmd::new(
create_index_command(INDEX, 3),
Err::<&str, _>(redis::RedisError::from((
redis::ErrorKind::Extension,
"Index already exists",
))),
Err::<&str, _>(support::server_error("Index already exists")),
),
info(INDEX, compatible_info(3)),
hset(INDEX, PROMPT, "key1", &VECTOR, &entry),
@ -673,10 +673,10 @@ fn shared_base_index_across_dimensions_replaces_the_isolated_index(entry: Value)
.set_cache("key1", entry.clone(), &context())
.unwrap();
assert_eq!(
worker_b.get_cache("key1", &context()).unwrap_err(),
Error::Unavailable
);
assert!(matches!(
worker_b.get_cache("key1", &context()),
Err(Error::Unavailable(Some(_)))
));
}
#[fixture]
@ -732,10 +732,10 @@ fn live_shared_index_is_replaced_across_dimensions(redis_stack_url: Option<Strin
.set_cache("key1", entry.clone(), &context())
.unwrap();
assert_eq!(
worker_b.get_cache("key1", &context()).unwrap_err(),
Error::Unavailable
);
assert!(matches!(
worker_b.get_cache("key1", &context()),
Err(Error::Unavailable(Some(_)))
));
let mut connection = redis::Client::open(url).unwrap().get_connection().unwrap();
for index in [&base, &isolated] {

15
litellm-rust/crates/cache-redis-semantic/tests/support/mod.rs
View file

@ -74,7 +74,7 @@ impl FakeSearch {
"FT.CREATE" => {
let name = text(1);
if state.indexes.contains_key(&name) {
return Err(error("Index already exists"));
return Err(server_error("Index already exists"));
}
let position = |token: &str| args.iter().position(|arg| arg == token.as_bytes());
let prefix = args[position("PREFIX").unwrap() + 2].clone();
@ -94,7 +94,7 @@ impl FakeSearch {
let index = state
.indexes
.get(&text(1))
.ok_or_else(|| error("Unknown index name"))?;
.ok_or_else(|| server_error("Unknown index name"))?;
Ok(index_info(index))
}
"FT.DROPINDEX" => {
@ -118,7 +118,7 @@ impl FakeSearch {
let index = state
.indexes
.get(&text(1))
.ok_or_else(|| error("no such index"))?;
.ok_or_else(|| server_error("no such index"))?;
let query = text(2);
let tag = query_tag(&query);
let params = args.iter().position(|arg| arg == b"PARAMS").unwrap();
@ -155,7 +155,7 @@ impl FakeSearch {
]))
}
"PING" => Ok(redis::Value::SimpleString("PONG".into())),
_ => Err(error("unsupported command")),
_ => Err(server_error("unsupported command")),
}
}
}
@ -192,8 +192,11 @@ impl redis::ConnectionLike for FakeSearch {
}
}
fn error(message: &'static str) -> redis::RedisError {
redis::RedisError::from((redis::ErrorKind::Extension, message))
pub fn server_error(reply: &str) -> redis::RedisError {
redis::parse_redis_value(format!("-{reply}\r\n").as_bytes())
.expect("error reply")
.extract_error()
.expect_err("server error")
}
fn bulk(bytes: &[u8]) -> redis::Value {

2
litellm-rust/crates/cache-redis/Cargo.toml
View file

@ -7,8 +7,10 @@ repository.workspace = true
[dependencies]
litellm-cache.workspace = true
litellm-error.workspace = true
redis = { version = "1.7.0", features = ["cluster", "tls-rustls"] }
r2d2 = "0.8.10"
thiserror.workspace = true
tokio.workspace = true
[dev-dependencies]

4
litellm-rust/crates/cache-redis/src/cache.rs
View file

@ -141,7 +141,7 @@ where
redis::Value::Nil => Ok(None),
redis::Value::BulkString(bytes) => self.codec.decode(&bytes).map(Some),
redis::Value::SimpleString(text) => self.codec.decode(text.as_bytes()).map(Some),
_ => Err(Error::InvalidEntry),
_ => Err(Error::InvalidEntry(None)),
}
}
@ -152,7 +152,7 @@ where
match self.decode_response(value) {
Ok(Some(value)) => Ok(BatchEntry::Hit(value)),
Ok(None) => Ok(BatchEntry::Miss),
Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid),
Err(Error::InvalidEntry(_)) => Ok(BatchEntry::Invalid),
Err(error) => Err(error),
}
}

14
litellm-rust/crates/cache-redis/src/claim.rs
View file

@ -17,7 +17,7 @@ fn stored_bytes(value: redis::Value) -> Result<Option<Vec<u8>>, Error> {
redis::Value::Nil => Ok(None),
redis::Value::BulkString(bytes) => Ok(Some(bytes)),
redis::Value::SimpleString(text) => Ok(Some(text.into_bytes())),
_ => Err(Error::InvalidEntry),
_ => Err(Error::InvalidEntry(None)),
}
}
@ -37,13 +37,13 @@ where
{
let payload = codec.encode(&candidate)?;
if payload.is_empty() {
return Err(Error::InvalidEntry);
return Err(Error::InvalidEntry(None));
}
for _ in 0..CLAIM_ATTEMPTS {
let current = stored_bytes(
connection
.get::<_, redis::Value>(key)
.map_err(|_| Error::Unavailable)?,
.map_err(Error::unavailable)?,
)?
.filter(|bytes| !bytes.is_empty());
let existing = current
@ -63,14 +63,18 @@ where
.arg(write)
.arg(u8::from(refresh))
.query::<bool>(connection)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
if applied {
return Ok(existing.unwrap_or(candidate));
}
}
Err(Error::Unavailable)
Err(Error::unavailable(ClaimContended))
}
#[derive(Debug, thiserror::Error)]
#[error("claim kept losing the race after {CLAIM_ATTEMPTS} attempts")]
struct ClaimContended;
impl<S, C> ClaimCache for RedisCache<S, C>
where
S: CacheCodec + Clone + 'static,

74
litellm-rust/crates/cache-redis/src/connection.rs
View file

@ -36,19 +36,19 @@ where
) -> Result<T, Error> {
match self {
Self::Pool(pool) => {
let mut pooled = pool.get().map_err(|_| Error::Unavailable)?;
let mut pooled = pool.get().map_err(Error::unavailable)?;
let result = operation(&mut ConnectionRef::Node(&mut pooled.connection));
pooled.failed = matches!(result, Err(Error::Unavailable));
pooled.failed = matches!(result, Err(Error::Unavailable(_)));
result
}
Self::Cluster(pool) => {
let mut pooled = pool.get().map_err(|_| Error::Unavailable)?;
let mut pooled = pool.get().map_err(Error::unavailable)?;
let result = operation(&mut ConnectionRef::Cluster(&mut pooled.connection));
pooled.failed = matches!(result, Err(Error::Unavailable));
pooled.failed = matches!(result, Err(Error::Unavailable(_)));
result
}
Self::Fixed(connection) => {
let mut connection = connection.lock().map_err(|_| Error::Unavailable)?;
let mut connection = connection.lock()?;
operation(&mut ConnectionRef::Node(&mut *connection))
}
}
@ -61,7 +61,7 @@ where
{
tokio::task::spawn_blocking(move || connections.execute(operation))
.await
.map_err(|_| Error::Unavailable)?
.map_err(Error::unavailable)?
}
pub fn fixed(connection: C) -> Self {
@ -95,7 +95,7 @@ fn pool<M: r2d2::ManageConnection>(manager: M) -> Result<r2d2::Pool<M>, Error> {
.connection_timeout(REDIS_TIMEOUT)
.test_on_check_out(false)
.build(manager)
.map_err(|_| Error::Unavailable)
.map_err(Error::unavailable)
}
fn close_idle<M, T>(pool: &r2d2::Pool<M>)
@ -123,7 +123,7 @@ impl ConnectionManager {
fn open(url: &str) -> Result<Self, Error> {
redis::Client::open(url)
.map(Self)
.map_err(|_| Error::Unavailable)
.map_err(Error::unavailable)
}
}
@ -151,14 +151,26 @@ impl r2d2::ManageConnection for ConnectionManager {
}
}
#[derive(Debug, thiserror::Error)]
enum TopologyFault {
#[error("Redis Cluster needs at least one startup node")]
NoStartupNodes,
#[error("Redis Cluster nodes need a TCP address")]
NonTcpAddress,
#[error("Redis Cluster answered a fan-out command without a per-node map")]
FanOutReply,
#[error("scan target does not match the connection topology")]
ScanTarget,
}
pub(crate) struct ClusterConnectionManager(ClusterClient);
impl ClusterConnectionManager {
fn open(url: &str, startup_nodes: &[RedisNode]) -> Result<Self, Error> {
if startup_nodes.is_empty() {
return Err(Error::Unavailable);
return Err(Error::unavailable(TopologyFault::NoStartupNodes));
}
let info = url.into_connection_info().map_err(|_| Error::Unavailable)?;
let info = url.into_connection_info().map_err(Error::unavailable)?;
let nodes = startup_nodes
.iter()
.map(|node| node_info(&info, node))
@ -168,7 +180,7 @@ impl ClusterConnectionManager {
.response_timeout(REDIS_TIMEOUT)
.build()
.map(Self)
.map_err(|_| Error::Unavailable)
.map_err(Error::unavailable)
}
}
@ -185,7 +197,7 @@ fn node_info(info: &ConnectionInfo, node: &RedisNode) -> Result<ConnectionInfo,
insecure: *insecure,
tls_params: tls_params.clone(),
},
_ => return Err(Error::Unavailable),
_ => return Err(Error::unavailable(TopologyFault::NonTcpAddress)),
};
Ok(info.clone().set_addr(addr))
}
@ -298,7 +310,7 @@ impl ConnectionRef<'_> {
.and_then(|replies| pipeline.query(&mut Replies(Some(replies))))
}
}
.map_err(|_| Error::Unavailable)
.map_err(Error::unavailable)
}
pub(crate) fn scan(
@ -311,7 +323,7 @@ impl ConnectionRef<'_> {
Self::Node(connection) => {
let page = scan_command(0, pattern, count)
.query::<ScanPage>(*connection)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
vec![(None, page)]
}
Self::Cluster(connection) => connection
@ -322,7 +334,7 @@ impl ConnectionRef<'_> {
Some(ResponsePolicy::Special),
)),
)
.map_err(|_| Error::Unavailable)
.map_err(Error::unavailable)
.and_then(primary_pages)?
.into_iter()
.map(|(node, page)| (Some(node), page))
@ -362,7 +374,7 @@ impl ConnectionRef<'_> {
pub(crate) fn node_text(&mut self, command: &redis::Cmd) -> Result<String, Error> {
match self {
Self::Node(connection) => command.query(*connection).map_err(|_| Error::Unavailable),
Self::Node(connection) => command.query(*connection).map_err(Error::unavailable),
Self::Cluster(connection) => {
let value = connection
.route_command(
@ -372,18 +384,16 @@ impl ConnectionRef<'_> {
Some(ResponsePolicy::Special),
)),
)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
let redis::Value::Map(entries) = value else {
return Err(Error::Unavailable);
return Err(Error::unavailable(TopologyFault::FanOutReply));
};
let mut replies = entries
.into_iter()
.map(|(node, reply)| {
Ok((
redis::from_redis_value::<String>(node)
.map_err(|_| Error::Unavailable)?,
redis::from_redis_value::<String>(reply)
.map_err(|_| Error::Unavailable)?,
redis::from_redis_value::<String>(node).map_err(Error::unavailable)?,
redis::from_redis_value::<String>(reply).map_err(Error::unavailable)?,
))
})
.collect::<Result<Vec<(String, String)>, Error>>()?;
@ -400,7 +410,7 @@ impl ConnectionRef<'_> {
pub(crate) fn flushall(&mut self) -> Result<(), Error> {
let command = redis::cmd("FLUSHALL");
match self {
Self::Node(connection) => command.query(*connection).map_err(|_| Error::Unavailable),
Self::Node(connection) => command.query(*connection).map_err(Error::unavailable),
Self::Cluster(connection) => connection
.route_command(
&command,
@ -410,7 +420,7 @@ impl ConnectionRef<'_> {
)),
)
.map(|_| ())
.map_err(|_| Error::Unavailable),
.map_err(Error::unavailable),
}
}
@ -424,7 +434,7 @@ impl ConnectionRef<'_> {
let command = scan_command(cursor, pattern, count);
match (self, node) {
(Self::Node(connection), None) => {
command.query(*connection).map_err(|_| Error::Unavailable)
command.query(*connection).map_err(Error::unavailable)
}
(Self::Cluster(connection), Some(node)) => connection
.route_command(
@ -434,9 +444,9 @@ impl ConnectionRef<'_> {
port: node.port(),
}),
)
.map_err(|_| Error::Unavailable)
.and_then(|value| redis::from_redis_value(value).map_err(|_| Error::Unavailable)),
_ => Err(Error::Unavailable),
.map_err(Error::unavailable)
.and_then(|value| redis::from_redis_value(value).map_err(Error::unavailable)),
_ => Err(Error::unavailable(TopologyFault::ScanTarget)),
}
}
}
@ -445,14 +455,14 @@ type ScanPage = (u64, Vec<String>);
fn primary_pages(value: redis::Value) -> Result<Vec<(NodeAddress, ScanPage)>, Error> {
let redis::Value::Map(entries) = value else {
return Err(Error::Unavailable);
return Err(Error::unavailable(TopologyFault::FanOutReply));
};
entries
.into_iter()
.map(|(node, page)| {
let node = redis::from_redis_value::<String>(node).map_err(|_| Error::Unavailable)?;
let node = NodeAddress::try_from(node.as_str()).map_err(|_| Error::Unavailable)?;
let page = redis::from_redis_value::<ScanPage>(page).map_err(|_| Error::Unavailable)?;
let node = redis::from_redis_value::<String>(node).map_err(Error::unavailable)?;
let node = NodeAddress::try_from(node.as_str()).map_err(Error::unavailable)?;
let page = redis::from_redis_value::<ScanPage>(page).map_err(Error::unavailable)?;
Ok((node, page))
})
.collect()

8
litellm-rust/crates/cache-redis/src/counter.rs
View file

@ -100,7 +100,7 @@ fn increment(
.arg(amount)
.arg(ttl)
.query(connection)
.map_err(|_| Error::Unavailable);
.map_err(Error::unavailable);
}
connection
.query_pipeline(
@ -140,7 +140,7 @@ where
}
fn count(value: redis::Value) -> Result<Option<i64>, Error> {
redis::from_redis_value(value).map_err(|_| Error::InvalidEntry)
redis::from_redis_value(value).map_err(Error::invalid_entry)
}
impl<S, C> BoundedCounterCache for RedisCache<S, C>
@ -182,7 +182,7 @@ where
.arg(value)
.arg(ttl)
.query(connection)
.map_err(|_| Error::Unavailable)
.map_err(Error::unavailable)
})
.await
}
@ -201,5 +201,5 @@ fn increment_with_floor(
.arg(amount)
.arg(ttl)
.query(connection)
.map_err(|_| Error::Unavailable)
.map_err(Error::unavailable)
}

4
litellm-rust/crates/cache-redis/src/keys.rs
View file

@ -16,7 +16,7 @@ where
redis::cmd("TTL")
.arg(key)
.query::<i64>(connection)
.map_err(|_| Error::Unavailable)
.map_err(Error::unavailable)
})
.await?;
Ok(u64::try_from(ttl).ok().map(Duration::from_secs))
@ -36,7 +36,7 @@ where
.arg(key)
.arg(ttl)
.query(connection)
.map_err(|_| Error::Unavailable)
.map_err(Error::unavailable)
})
.await
}

9
litellm-rust/crates/cache-redis/src/lifecycle.rs
View file

@ -11,11 +11,11 @@ where
C: redis::ConnectionLike + Send + 'static,
{
fn sync_ping(&self) -> Result<bool, Error> {
self.execute(|connection| connection.ping().map_err(|_| Error::Unavailable))
self.execute(|connection| connection.ping().map_err(Error::unavailable))
}
async fn ping(&self) -> Result<bool, Error> {
self.run(|connection| connection.ping().map_err(|_| Error::Unavailable))
self.run(|connection| connection.ping().map_err(Error::unavailable))
.await
}
}
@ -33,9 +33,10 @@ where
RedisTopology::Cluster { .. } => "Redis Cluster",
};
let ping = self
.run(|connection| Ok(connection.ping().map_err(|error| error.to_string())))
.run(|connection| Ok(connection.ping().map_err(Error::unavailable)))
.await
.unwrap_or_else(|error| Err(error.to_string()));
.and_then(|ping| ping)
.map_err(|error| litellm_error::Report::new(&error).to_string());
Ok(match ping {
Ok(true) => CacheConnectionResult {
status: CacheConnectionStatus::Success,

16
litellm-rust/crates/cache-redis/src/queue.rs
View file

@ -29,7 +29,7 @@ where
ttl: Option<Duration>,
) -> Result<Self::SetResult, Error> {
if values.is_empty() {
return Err(Error::InvalidEntry);
return Err(Error::InvalidEntry(None));
}
let key = self.namespaced_key(key);
let ttl = self.ttl_or_default(ttl);
@ -58,7 +58,7 @@ where
async fn async_rpush(&self, key: &str, values: Vec<Self::QueueValue>) -> Result<usize, Error> {
if values.is_empty() {
return Err(Error::InvalidEntry);
return Err(Error::InvalidEntry(None));
}
let key = self.namespaced_key(key);
self.run(move |connection| {
@ -66,7 +66,7 @@ where
.arg(key)
.arg(values)
.query(connection)
.map_err(|_| Error::Unavailable)
.map_err(Error::unavailable)
})
.await
}
@ -78,7 +78,7 @@ where
max_len: usize,
) -> Result<usize, Error> {
if values.is_empty() {
return Err(Error::InvalidEntry);
return Err(Error::InvalidEntry(None));
}
let key = self.namespaced_key(key);
let start = i64::try_from(max_len).map_or(i64::MIN, |max_len| -max_len);
@ -108,7 +108,7 @@ where
let mut pipeline = redis::pipe();
for operation in operations {
if operation.values.is_empty() {
return Err(Error::InvalidEntry);
return Err(Error::InvalidEntry(None));
}
pipeline
.cmd("RPUSH")
@ -130,7 +130,7 @@ where
.run(move |connection| {
command
.query::<redis::Value>(connection)
.map_err(|_| Error::Unavailable)
.map_err(Error::unavailable)
})
.await?;
lpop_result(value, count.is_some())
@ -182,7 +182,7 @@ fn redis_bytes(value: redis::Value) -> Result<Vec<u8>, Error> {
match value {
redis::Value::BulkString(bytes) => Ok(bytes),
redis::Value::SimpleString(text) => Ok(text.into_bytes()),
_ => Err(Error::InvalidEntry),
_ => Err(Error::InvalidEntry(None)),
}
}
@ -221,6 +221,6 @@ fn lpop_result(value: redis::Value, multiple: bool) -> Result<RedisLpopResult, E
.collect::<Result<Vec<_>, _>>()
.map(RedisLpopResult::Values),
value if !multiple => redis_bytes(value).map(RedisLpopResult::Value),
_ => Err(Error::InvalidEntry),
_ => Err(Error::InvalidEntry(None)),
}
}

2
litellm-rust/crates/cache-redis/src/script.rs
View file

@ -116,7 +116,7 @@ fn eval(
.arg(keys)
.arg(arguments)
.query(connection)
.map_err(|_| Error::Unavailable)
.map_err(Error::unavailable)
}
impl<S, C> ScriptCache for RedisCache<S, C>

18
litellm-rust/crates/cache-redis/src/store.rs
View file

@ -32,7 +32,7 @@ where
self.execute(|connection| {
connection
.set_ex::<_, _, ()>(key, payload, ttl)
.map_err(|_| Error::Unavailable)
.map_err(Error::unavailable)
})
}
@ -41,7 +41,7 @@ where
let value = self.execute(|connection| {
connection
.get::<_, redis::Value>(key)
.map_err(|_| Error::Unavailable)
.map_err(Error::unavailable)
})?;
self.decode_response(value)
}
@ -58,7 +58,7 @@ where
self.run(move |connection| {
connection
.set_ex::<_, _, ()>(key, payload, ttl)
.map_err(|_| Error::Unavailable)
.map_err(Error::unavailable)
})
.await
}
@ -73,7 +73,7 @@ where
.run(move |connection| {
connection
.get::<_, redis::Value>(key)
.map_err(|_| Error::Unavailable)
.map_err(Error::unavailable)
})
.await?;
self.decode_response(value)
@ -158,7 +158,7 @@ pub(crate) fn mget(
redis::cmd("MGET")
.arg(keys)
.query(connection)
.map_err(|_| Error::Unavailable)
.map_err(Error::unavailable)
}
impl<S, C> DeleteCache for RedisCache<S, C>
@ -168,12 +168,12 @@ where
{
fn delete_cache(&self, key: &str) -> Result<(), Error> {
let key = self.namespaced_key(key);
self.execute(|connection| connection.del::<_, ()>(key).map_err(|_| Error::Unavailable))
self.execute(|connection| connection.del::<_, ()>(key).map_err(Error::unavailable))
}
async fn async_delete_cache(&self, key: &str) -> Result<(), Error> {
let key = self.namespaced_key(key);
self.run(move |connection| connection.del::<_, ()>(key).map_err(|_| Error::Unavailable))
self.run(move |connection| connection.del::<_, ()>(key).map_err(Error::unavailable))
.await
}
}
@ -188,7 +188,7 @@ where
return Ok(0);
}
let keys = self.namespaced_keys(&keys);
self.run(move |connection| connection.del(keys).map_err(|_| Error::Unavailable))
self.run(move |connection| connection.del(keys).map_err(Error::unavailable))
.await
}
}
@ -225,7 +225,7 @@ fn flush_matching(connection: &mut ConnectionRef<'_>, pattern: &str) -> Result<(
if !keys.is_empty() {
connection
.del::<_, usize>(keys)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
}
Ok(true)
})

38
litellm-rust/crates/cache-redis/tests/cache.rs
View file

@ -178,11 +178,11 @@ async fn codec_errors_propagate_without_writing_partial_batches(context: ExactCa
]);
assert_eq!(
cache.set_cache("invalid", 255, &context),
Err(Error::InvalidEntry)
Err(Error::InvalidEntry(None))
);
assert_eq!(
cache.async_set_cache("invalid", 255, context.clone()).await,
Err(Error::InvalidEntry)
Err(Error::InvalidEntry(None))
);
assert_eq!(
cache
@ -191,7 +191,7 @@ async fn codec_errors_propagate_without_writing_partial_batches(context: ExactCa
context.clone(),
)
.await,
Err(Error::InvalidEntry)
Err(Error::InvalidEntry(None))
);
assert_eq!(
cache
@ -200,15 +200,15 @@ async fn codec_errors_propagate_without_writing_partial_batches(context: ExactCa
("invalid".into(), 255, None),
])
.await,
Err(Error::InvalidEntry)
Err(Error::InvalidEntry(None))
);
assert_eq!(
cache.get_cache("invalid", &context),
Err(Error::InvalidEntry)
Err(Error::InvalidEntry(None))
);
assert_eq!(
cache.async_get_cache("invalid", &context).await,
Err(Error::InvalidEntry)
Err(Error::InvalidEntry(None))
);
}
@ -516,7 +516,7 @@ async fn sets_add_members_and_arm_the_default_ttl() {
cache
.async_set_cache_sadd("members", Vec::new(), None)
.await,
Err(Error::InvalidEntry)
Err(Error::InvalidEntry(None))
);
}
@ -547,7 +547,7 @@ async fn queues_push_and_pop_namespaced_lists() {
);
assert_eq!(
cache.async_rpush("queue", Vec::new()).await,
Err(Error::InvalidEntry)
Err(Error::InvalidEntry(None))
);
assert_eq!(
cache.async_lpop("queue", Some(2)).await.unwrap(),
@ -673,15 +673,15 @@ async fn rpush_and_trim_raises_when_a_queued_command_fails() {
)])
.with_namespace(Some("ns".into()));
assert_eq!(
assert!(matches!(
cache
.async_rpush_and_trim("buf", vec!["c".into(), "d".into()], 3)
.await,
Err(Error::Unavailable)
);
Err(Error::Unavailable(Some(_)))
));
assert_eq!(
cache.async_rpush_and_trim("buf", Vec::new(), 3).await,
Err(Error::InvalidEntry)
Err(Error::InvalidEntry(None))
);
}
@ -1122,24 +1122,30 @@ async fn test_connection_reports_success_with_the_python_message() {
}
#[rstest]
#[case::unexpected_reply(Ok("NOPE"), "Redis ping returned False", false)]
#[case::unexpected_reply(Ok("NOPE"), "Redis ping returned False", None)]
#[case::connection_refused(
Err(redis::RedisError::from((redis::ErrorKind::Io, "connection refused"))),
"Redis connection failed:",
true
Some("connection refused")
)]
#[tokio::test]
async fn test_connection_failures_use_the_python_result_contract(
#[case] reply: redis::RedisResult<&'static str>,
#[case] message: &str,
#[case] has_error: bool,
#[case] cause: Option<&str>,
) {
let cache = json_cache(vec![MockCmd::new(redis::cmd("PING"), reply)]);
let result = cache.test_connection().await.unwrap();
assert_eq!(result.status, CacheConnectionStatus::Failed);
assert!(result.message.starts_with(message), "{}", result.message);
assert_eq!(result.error.is_some(), has_error);
assert_eq!(result.error.is_some(), cause.is_some());
assert!(cause.is_none_or(|cause| {
result
.error
.as_deref()
.is_some_and(|error| error.contains(cause))
}));
}
#[rstest]

2
litellm-rust/crates/cache-redis/tests/cluster.rs
View file

@ -66,7 +66,7 @@ fn constructor_rejects_unusable_cluster_configs(
JsonCodec::new(),
)
.err();
assert!(matches!(error, Some(Error::Unavailable)));
assert!(matches!(error, Some(Error::Unavailable(_))));
}
#[rstest]

4
litellm-rust/crates/cache-redis/tests/support/mod.rs
View file

@ -16,7 +16,7 @@ impl CacheCodec for TaggedByteCodec {
fn encode(&self, value: &u8) -> Result<Vec<u8>, Error> {
if *value > 127 {
return Err(Error::InvalidEntry);
return Err(Error::InvalidEntry(None));
}
Ok(vec![self.0, *value])
}
@ -24,7 +24,7 @@ impl CacheCodec for TaggedByteCodec {
fn decode(&self, bytes: &[u8]) -> Result<u8, Error> {
match bytes {
[tag, value] if *tag == self.0 => Ok(*value),
_ => Err(Error::InvalidEntry),
_ => Err(Error::InvalidEntry(None)),
}
}
}

4
litellm-rust/crates/cache-response/src/buffer.rs
View file

@ -26,7 +26,7 @@ impl WriteBuffer {
now: Duration,
) -> Result<(), Error> {
let pending = {
let mut entries = self.entries.lock().map_err(|_| Error::Unavailable)?;
let mut entries = self.entries.lock()?;
entries.push((request.clone(), response, now));
(entries.len() >= self.flush_size).then(|| std::mem::take(&mut *entries))
};
@ -39,7 +39,7 @@ impl WriteBuffer {
}
pub fn clear(&self) -> Result<(), Error> {
self.entries.lock().map_err(|_| Error::Unavailable)?.clear();
self.entries.lock()?.clear();
Ok(())
}
}

28
litellm-rust/crates/cache-response/src/codec.rs
View file

@ -14,24 +14,24 @@ impl CacheCodec for ResponseCacheCodec {
.timestamp
.is_some_and(|timestamp| !timestamp.is_finite())
{
return Err(Error::InvalidEntry);
return Err(Error::InvalidEntry(None));
}
// Python reads a `response` that is either a dict or a serialized string, so every
// other shape is written serialized. A string on the wire is therefore always a
// serialized response, which keeps string-valued responses unambiguous.
if value.timestamp.is_none() || value.response.is_object() {
return serde_json::to_vec(value).map_err(|_| Error::InvalidEntry);
return serde_json::to_vec(value).map_err(Error::invalid_entry);
}
let response = serde_json::to_string(&value.response).map_err(|_| Error::InvalidEntry)?;
let response = serde_json::to_string(&value.response).map_err(Error::invalid_entry)?;
serde_json::to_vec(&CacheEntry {
timestamp: value.timestamp,
response: Value::String(response),
})
.map_err(|_| Error::InvalidEntry)
.map_err(Error::invalid_entry)
}
fn decode(&self, bytes: &[u8]) -> Result<CacheEntry, Error> {
let text = std::str::from_utf8(bytes).map_err(|_| Error::InvalidEntry)?;
let text = std::str::from_utf8(bytes).map_err(Error::invalid_entry)?;
let value = decode_value(text)?;
let Some(timestamp) = value.get("timestamp") else {
return Ok(CacheEntry {
@ -40,9 +40,9 @@ impl CacheCodec for ResponseCacheCodec {
});
};
let Some(timestamp) = timestamp.as_f64().filter(|timestamp| timestamp.is_finite()) else {
return Err(Error::InvalidEntry);
return Err(Error::InvalidEntry(None));
};
let response = match value.get("response").ok_or(Error::InvalidEntry)? {
let response = match value.get("response").ok_or(Error::InvalidEntry(None))? {
Value::String(text) => decode_value(text)?,
response => response.clone(),
};
@ -58,25 +58,25 @@ fn decode_value(text: &str) -> Result<Value, Error> {
return Ok(value);
}
check_literal_depth(text)?;
let literal: py_literal::Value = text.parse().map_err(|_| Error::InvalidEntry)?;
let literal: py_literal::Value = text.parse().map_err(Error::invalid_entry)?;
literal_value(literal, 0)
}
fn literal_value(value: py_literal::Value, depth: usize) -> Result<Value, Error> {
use py_literal::Value as Literal;
if depth > 128 {
return Err(Error::InvalidEntry);
return Err(Error::InvalidEntry(None));
}
match value {
Literal::String(text) => Ok(Value::String(text)),
Literal::Boolean(value) => Ok(Value::Bool(value)),
Literal::None => Ok(Value::Null),
Literal::Integer(value) => {
serde_json::from_str(&value.to_string()).map_err(|_| Error::InvalidEntry)
serde_json::from_str(&value.to_string()).map_err(Error::invalid_entry)
}
Literal::Float(value) => serde_json::Number::from_f64(value)
.map(Value::Number)
.ok_or(Error::InvalidEntry),
.ok_or(Error::InvalidEntry(None)),
Literal::List(values) | Literal::Tuple(values) => values
.into_iter()
.map(|value| literal_value(value, depth + 1))
@ -86,13 +86,13 @@ fn literal_value(value: py_literal::Value, depth: usize) -> Result<Value, Error>
.into_iter()
.map(|(key, value)| {
let Literal::String(key) = key else {
return Err(Error::InvalidEntry);
return Err(Error::InvalidEntry(None));
};
Ok((key, literal_value(value, depth + 1)?))
})
.collect::<Result<serde_json::Map<_, _>, _>>()
.map(Value::Object),
_ => Err(Error::InvalidEntry),
_ => Err(Error::InvalidEntry(None)),
}
}
@ -118,7 +118,7 @@ fn check_literal_depth(text: &str) -> Result<(), Error> {
'[' | '{' | '(' => {
depth += 1;
if depth > 128 {
return Err(Error::InvalidEntry);
return Err(Error::InvalidEntry(None));
}
}
']' | '}' | ')' => depth = depth.saturating_sub(1),

13
litellm-rust/crates/cache-response/src/response.rs
View file

@ -2,7 +2,7 @@ use std::{sync::Arc, time::Duration};
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, CacheConnectionResult, CacheContext, ConnectionCache, Error,
FlushCache,
FlushCache, ShortBatch,
semantic::{SemanticCache, SemanticLookup},
};
use serde_json::Value;
@ -114,7 +114,7 @@ where
.get_cache(&cache_key(&request.key), &request.context)
{
Ok(entry) => entry,
Err(Error::InvalidEntry) => None,
Err(Error::InvalidEntry(_)) => None,
Err(error) => return Err(error),
};
Ok(Self::fresh_or_miss(entry, now, request.max_age))
@ -134,7 +134,7 @@ where
.await
{
Ok(entry) => entry,
Err(Error::InvalidEntry) => None,
Err(Error::InvalidEntry(_)) => None,
Err(error) => return Err(error),
};
Ok(Self::fresh_or_miss(entry, now, request.max_age))
@ -338,7 +338,10 @@ where
now: Duration,
) -> Result<PartialHits, Error> {
if readable.len() != entries.len() {
return Err(Error::Unavailable);
return Err(Error::unavailable(ShortBatch {
requested: readable.len(),
returned: entries.len(),
}));
}
let mut values = vec![None; requests.len()];
for ((index, request), entry) in readable.into_iter().zip(entries) {
@ -361,7 +364,7 @@ where
value: Self::fresh_or_miss(lookup.value, now, max_age),
similarity: lookup.similarity,
}),
Err(Error::InvalidEntry) => Ok(SemanticLookup::miss(None)),
Err(Error::InvalidEntry(_)) => Ok(SemanticLookup::miss(None)),
Err(error) => Err(error),
}
}

24
litellm-rust/crates/cache-response/tests/codec.rs
View file

@ -44,25 +44,23 @@ fn decode_reads_every_python_envelope(#[case] bytes: &[u8], #[case] response: Va
#[case::non_utf8(vec![0xff, 0xfe])]
#[case::deep_nesting(format!("{}None{}", "[".repeat(1000), "]".repeat(1000)).into_bytes())]
fn decode_rejects_invalid_entries(#[case] bytes: Vec<u8>) {
assert_eq!(
ResponseCacheCodec.decode(&bytes).unwrap_err(),
Error::InvalidEntry
);
assert!(matches!(
ResponseCacheCodec.decode(&bytes),
Err(Error::InvalidEntry(_))
));
}
#[rstest]
#[case::nan(f64::NAN)]
#[case::infinity(f64::INFINITY)]
fn encode_rejects_non_finite_timestamps(#[case] timestamp: f64) {
assert_eq!(
ResponseCacheCodec
.encode(&CacheEntry {
timestamp: Some(timestamp),
response: json!({}),
})
.unwrap_err(),
Error::InvalidEntry
);
assert!(matches!(
ResponseCacheCodec.encode(&CacheEntry {
timestamp: Some(timestamp),
response: json!({}),
}),
Err(Error::InvalidEntry(_))
));
}
#[rstest]

12
litellm-rust/crates/cache-response/tests/response.rs
View file

@ -158,8 +158,16 @@ fn scored(timestamp: f64, similarity: f64) -> Result<SemanticLookup<CacheEntry>,
)]
#[case::no_search(Ok(SemanticLookup::miss(None)), true, Ok(SemanticLookup::miss(None)))]
#[case::disabled_reads_skip_the_backend(scored(95.0, 0.95), false, Ok(SemanticLookup::miss(None)))]
#[case::invalid_entry_is_a_miss(Err(Error::InvalidEntry), true, Ok(SemanticLookup::miss(None)))]
#[case::backend_errors_propagate(Err(Error::Unavailable), true, Err(Error::Unavailable))]
#[case::invalid_entry_is_a_miss(
Err(Error::InvalidEntry(None)),
true,
Ok(SemanticLookup::miss(None))
)]
#[case::backend_errors_propagate(
Err(Error::Unavailable(None)),
true,
Err(Error::Unavailable(None))
)]
#[tokio::test]
async fn semantic_lookup_applies_freshness_to_the_value_only(
#[case] backend: Result<SemanticLookup<CacheEntry>, Error>,

1
litellm-rust/crates/cache-s3/Cargo.toml
View file

@ -20,6 +20,7 @@ reqwest.workspace = true
tokio.workspace = true
[dev-dependencies]
litellm-error.workspace = true
litellm-http = { workspace = true, features = ["test-support"] }
litellm-cache-testing.workspace = true
rstest.workspace = true

2
litellm-rust/crates/cache-s3/src/auth.rs
View file

@ -46,7 +46,7 @@ impl ProvideCredentials for S3Credentials {
self.auth
.resolve_credentials(self.config.clone(), &self.env)
.await
.map_err(|_| CredentialsError::provider_error("S3 cache authentication failed"))
.map_err(CredentialsError::provider_error)
})
}
}

6
litellm-rust/crates/cache-s3/src/cache.rs
View file

@ -125,7 +125,7 @@ impl<C: CacheCodec> S3Cache<C> {
}
None => request.cache_control("immutable, max-age=31536000, s-maxage=31536000"),
};
request.send().await.map_err(|_| Error::Unavailable)?;
request.send().await.map_err(Error::unavailable)?;
Ok(())
}
@ -152,7 +152,7 @@ impl<C: CacheCodec> S3Cache<C> {
return Ok(None);
}
}
return Err(Error::Unavailable);
return Err(Error::unavailable(error));
}
};
if let Some(expires) = output.expires_string()
@ -165,7 +165,7 @@ impl<C: CacheCodec> S3Cache<C> {
.body
.collect()
.await
.map_err(|_| Error::Unavailable)?
.map_err(Error::unavailable)?
.into_bytes();
self.codec.decode(&bytes).map(Some)
}

4
litellm-rust/crates/cache-s3/src/transport.rs
View file

@ -32,7 +32,9 @@ impl HttpConnector for ReqwestHttpClient {
let request = reqwest::Request::try_from(request)
.map_err(|error| ConnectorError::other(error.into(), None))?;
let response = client.execute(request).await.map_err(|error| {
if error.is_timeout() {
let timeout = error.is_timeout();
let error = litellm_http::Error::from(error);
if timeout {
ConnectorError::timeout(error.into())
} else {
ConnectorError::io(error.into())

24
litellm-rust/crates/cache-s3/tests/cache.rs
View file

@ -147,15 +147,15 @@ async fn async_set_signs_with_the_configured_keys(#[future(awt)] server: MockSer
#[case::malformed(
"malformed",
ResponseTemplate::new(200).set_body_string("not a cache entry"),
Err(Error::InvalidEntry)
Err("invalid cache entry: ")
)]
#[case::server_error("broken", ResponseTemplate::new(500), Err(Error::Unavailable))]
#[case::server_error("broken", ResponseTemplate::new(500), Err("cache is unavailable: "))]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn get_maps_s3_responses(
#[future(awt)] server: MockServer,
#[case] key: &str,
#[case] response: ResponseTemplate,
#[case] expected: Result<Option<Value>, Error>,
#[case] expected: Result<Option<Value>, &str>,
) {
Mock::given(method("GET"))
.and(path(format!("/cache-bucket/team/{key}")))
@ -164,9 +164,21 @@ async fn get_maps_s3_responses(
.await;
let cache = support::cache(&server.uri());
let context = ExactCacheContext::default();
assert_eq!(cache.get_cache(key, &context), expected);
assert_eq!(cache.async_get_cache(key, &context).await, expected);
let outcome = |result: Result<Option<Value>, Error>| {
result.map_err(|error| litellm_error::Report::new(&error).to_string())
};
let sync = outcome(cache.get_cache(key, &context));
let asynchronous = outcome(cache.async_get_cache(key, &context).await);
match expected {
Ok(value) => {
assert_eq!(sync, Ok(value.clone()));
assert_eq!(asynchronous, Ok(value));
}
Err(prefix) => {
assert!(sync.unwrap_err().starts_with(prefix));
assert!(asynchronous.unwrap_err().starts_with(prefix));
}
}
}
#[rstest]

2
litellm-rust/crates/cache-valkey-semantic/Cargo.toml
View file

@ -10,9 +10,11 @@ litellm-cache.workspace = true
litellm-cache-redis.workspace = true
redis = { version = "1.7.0", features = ["tls-rustls"] }
sha2.workspace = true
thiserror.workspace = true
uuid = { version = "1", features = ["v4"] }
[dev-dependencies]
litellm-error.workspace = true
litellm-cache-testing.workspace = true
redis-test = "1.0.4"
rstest.workspace = true

39
litellm-rust/crates/cache-valkey-semantic/src/index.rs
View file

@ -23,8 +23,7 @@ pub(crate) fn ensure_index(
) -> Result<(), Error> {
if index
.dimension
.lock()
.map_err(|_| Error::Unavailable)?
.lock()?
.is_some_and(|existing| existing == dimension)
{
return Ok(());
@ -50,24 +49,42 @@ pub(crate) fn ensure_index(
.arg("DISTANCE_METRIC")
.arg("COSINE")
.query::<String>(connection)
.map(|_| ())
.map_err(|error| error.to_string());
if let Err(message) = create {
if !message.to_ascii_lowercase().contains("already exists") {
return Err(Error::Unavailable);
.map(|_| ());
if let Err(error) = create {
if !index_exists(&error) {
return Err(Error::unavailable(error));
}
let info = redis::cmd("FT.INFO")
.arg(&index.name)
.query::<redis::Value>(connection)
.map_err(|_| Error::Unavailable)?;
if index_dimension_from_info(&info).is_some_and(|existing| existing != dimension) {
return Err(Error::Unavailable);
.map_err(Error::unavailable)?;
if let Some(existing) =
index_dimension_from_info(&info).filter(|existing| *existing != dimension)
{
return Err(Error::unavailable(DimensionMismatch {
existing,
dimension,
}));
}
}
*index.dimension.lock().map_err(|_| Error::Unavailable)? = Some(dimension);
*index.dimension.lock()? = Some(dimension);
Ok(())
}
fn index_exists(error: &redis::RedisError) -> bool {
error.code().is_some()
&& error
.detail()
.is_some_and(|detail| detail.to_ascii_lowercase().contains("already exists"))
}
#[derive(Debug, thiserror::Error)]
#[error("index exists with dimension {existing}, not {dimension}")]
struct DimensionMismatch {
existing: usize,
dimension: usize,
}
/// `_extract_index_dim`: flatten each attribute one level and read the value after
/// `dimensions`.
fn index_dimension_from_info(value: &redis::Value) -> Option<usize> {

28
litellm-rust/crates/cache-valkey-semantic/src/search.rs
View file

@ -53,9 +53,7 @@ pub(crate) fn write_document(
.arg(ttl.as_secs())
.ignore();
}
pipeline
.query::<()>(connection)
.map_err(|_| Error::Unavailable)
pipeline.query::<()>(connection).map_err(Error::unavailable)
}
/// The KNN-1 search within `scope`: the closest document's similarity, and its stored response
@ -83,7 +81,7 @@ pub(crate) fn search_document(
.arg("DIALECT")
.arg(2)
.query::<redis::Value>(connection)
.map_err(|_| Error::Unavailable)?;
.map_err(Error::unavailable)?;
let Some(fields) = search_fields(response)? else {
return Ok(SemanticLookup::miss(Some(0.0)));
};
@ -91,7 +89,7 @@ pub(crate) fn search_document(
fields
.iter()
.find_map(|(field, value)| (field == name).then(|| value.clone()))
.ok_or(Error::InvalidEntry)
.ok_or(Error::InvalidEntry(None))
};
let response = field("response")?;
let similarity = 1.0 - parse_f64(&field("vector_distance")?)?;
@ -105,24 +103,24 @@ type SearchFields = Vec<(String, Vec<u8>)>;
fn search_fields(value: redis::Value) -> Result<Option<SearchFields>, Error> {
let redis::Value::Array(values) = value else {
return Err(Error::InvalidEntry);
return Err(Error::InvalidEntry(None));
};
let total = parse_i64(values.first().ok_or(Error::InvalidEntry)?)?;
let total = parse_i64(values.first().ok_or(Error::InvalidEntry(None))?)?;
if total <= 0 || values.len() < 3 {
return Ok(None);
}
let redis::Value::Array(fields) = &values[2] else {
return Err(Error::InvalidEntry);
return Err(Error::InvalidEntry(None));
};
let (pairs, remainder) = fields.as_chunks::<2>();
if !remainder.is_empty() {
return Err(Error::InvalidEntry);
return Err(Error::InvalidEntry(None));
}
pairs
.iter()
.map(|pair| {
Ok((
value_text(&pair[0]).ok_or(Error::InvalidEntry)?,
value_text(&pair[0]).ok_or(Error::InvalidEntry(None))?,
value_bytes(&pair[1])?,
))
})
@ -132,16 +130,16 @@ fn search_fields(value: redis::Value) -> Result<Option<SearchFields>, Error> {
fn parse_i64(value: &redis::Value) -> Result<i64, Error> {
value_text(value)
.ok_or(Error::InvalidEntry)?
.ok_or(Error::InvalidEntry(None))?
.parse()
.map_err(|_| Error::InvalidEntry)
.map_err(Error::invalid_entry)
}
fn parse_f64(value: &[u8]) -> Result<f64, Error> {
std::str::from_utf8(value)
.map_err(|_| Error::InvalidEntry)?
.map_err(Error::invalid_entry)?
.parse()
.map_err(|_| Error::InvalidEntry)
.map_err(Error::invalid_entry)
}
pub(crate) fn value_text(value: &redis::Value) -> Option<String> {
@ -158,6 +156,6 @@ fn value_bytes(value: &redis::Value) -> Result<Vec<u8>, Error> {
redis::Value::BulkString(bytes) => Ok(bytes.clone()),
redis::Value::SimpleString(value) => Ok(value.as_bytes().to_vec()),
redis::Value::Int(value) => Ok(value.to_string().into_bytes()),
_ => Err(Error::InvalidEntry),
_ => Err(Error::InvalidEntry(None)),
}
}

43
litellm-rust/crates/cache-valkey-semantic/tests/cache.rs
View file

@ -83,10 +83,7 @@ fn ok() -> redis::RedisResult<redis::Value> {
}
fn already_exists() -> redis::RedisResult<redis::Value> {
Err(redis::RedisError::from((
redis::ErrorKind::Io,
"already exists",
)))
Err(support::server_error("Index already exists"))
}
fn bulk(value: &[u8]) -> redis::Value {
@ -250,19 +247,29 @@ fn second_set_skips_create_after_dimension_is_cached(entry: Value, context: Sema
#[rstest]
#[case::nested_matching(nested_dimension_info(3), Ok(()))]
#[case::nested_mismatch(nested_dimension_info(2), Err(Error::Unavailable))]
#[case::nested_mismatch(
nested_dimension_info(2),
Err("cache is unavailable: index exists with dimension 2, not 3")
)]
#[case::flat_matching(flat_dimension_info("3"), Ok(()))]
#[case::flat_mismatch(flat_dimension_info("2"), Err(Error::Unavailable))]
#[case::flat_mismatch(
flat_dimension_info("2"),
Err("cache is unavailable: index exists with dimension 2, not 3")
)]
#[case::unreported_dimension_is_accepted(redis::Value::Array(vec![]), Ok(()))]
fn existing_index_dimension_must_match_embedding(
#[case] info: redis::Value,
#[case] expected: Result<(), Error>,
#[case] expected: Result<(), &str>,
entry: Value,
context: SemanticCacheContext,
) {
let recording = recording([already_exists(), Ok(info)]);
assert_eq!(recording.cache.set_cache("key", entry, &context), expected);
let result = recording.cache.set_cache("key", entry, &context);
assert_eq!(
result.map_err(|error| litellm_error::Report::new(&error).to_string()),
expected.map_err(str::to_owned)
);
}
#[rstest]
@ -274,10 +281,10 @@ fn index_creation_failures_are_unavailable(
) {
let recording = recording([reply]);
assert_eq!(
assert!(matches!(
recording.cache.set_cache("key", entry, &context),
Err(Error::Unavailable)
);
Err(Error::Unavailable(Some(_)))
));
}
#[rstest]
@ -291,10 +298,10 @@ fn index_creation_failures_are_unavailable(
bulk(b"document"),
redis::Value::Array(vec![bulk(b"vector_distance"), bulk(b"0.1")]),
]),
Err(Error::InvalidEntry)
Err(Error::InvalidEntry(None))
)]
#[case::unparsable_distance(search_hit(b"not-json", "abc"), Err(Error::InvalidEntry))]
#[case::undecodable_response(search_hit(b"not-json", "0.1"), Err(Error::InvalidEntry))]
#[case::unparsable_distance(search_hit(b"not-json", "abc"), Err(Error::InvalidEntry(None)))]
#[case::undecodable_response(search_hit(b"not-json", "0.1"), Err(Error::InvalidEntry(None)))]
fn get_applies_threshold_and_decodes_entry(
#[case] reply: redis::Value,
#[case] expected: Result<Option<Value>, Error>,
@ -302,7 +309,13 @@ fn get_applies_threshold_and_decodes_entry(
) {
let recording = recording([ok(), Ok(reply)]);
assert_eq!(recording.cache.get_cache("key", &context), expected);
let kind = |result: Result<Option<Value>, Error>| {
result.map_err(|error| std::mem::discriminant(&error))
};
assert_eq!(
kind(recording.cache.get_cache("key", &context)),
kind(expected)
);
}
#[rstest]

15
litellm-rust/crates/cache-valkey-semantic/tests/support/mod.rs
View file

@ -74,7 +74,7 @@ impl FakeSearch {
"FT.CREATE" => {
let name = text(1);
if state.indexes.contains_key(&name) {
return Err(error("Index already exists"));
return Err(server_error("Index already exists"));
}
let position = |token: &str| args.iter().position(|arg| arg == token.as_bytes());
let prefix = args[position("PREFIX").unwrap() + 2].clone();
@ -94,7 +94,7 @@ impl FakeSearch {
let index = state
.indexes
.get(&text(1))
.ok_or_else(|| error("Unknown index name"))?;
.ok_or_else(|| server_error("Unknown index name"))?;
Ok(index_info(index))
}
"FT.DROPINDEX" => {
@ -118,7 +118,7 @@ impl FakeSearch {
let index = state
.indexes
.get(&text(1))
.ok_or_else(|| error("no such index"))?;
.ok_or_else(|| server_error("no such index"))?;
let query = text(2);
let tag = query_tag(&query);
let params = args.iter().position(|arg| arg == b"PARAMS").unwrap();
@ -155,7 +155,7 @@ impl FakeSearch {
]))
}
"PING" => Ok(redis::Value::SimpleString("PONG".into())),
_ => Err(error("unsupported command")),
_ => Err(server_error("unsupported command")),
}
}
}
@ -192,8 +192,11 @@ impl redis::ConnectionLike for FakeSearch {
}
}
fn error(message: &'static str) -> redis::RedisError {
redis::RedisError::from((redis::ErrorKind::Extension, message))
pub fn server_error(reply: &str) -> redis::RedisError {
redis::parse_redis_value(format!("-{reply}\r\n").as_bytes())
.expect("error reply")
.extract_error()
.expect_err("server error")
}
fn bulk(bytes: &[u8]) -> redis::Value {

1
litellm-rust/crates/cache/Cargo.toml vendored
View file

@ -6,6 +6,7 @@ license.workspace = true
repository.workspace = true
[dependencies]
litellm-error.workspace = true
serde.workspace = true
serde_json = { workspace = true, features = ["preserve_order"] }
strum.workspace = true

4
litellm-rust/crates/cache/src/capabilities.rs vendored
View file

@ -48,7 +48,7 @@ pub trait BatchCache: BaseCache {
.map(|key| match self.get_cache(key, context) {
Ok(Some(value)) => Ok(BatchEntry::Hit(value)),
Ok(None) => Ok(BatchEntry::Miss),
Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid),
Err(Error::InvalidEntry(_)) => Ok(BatchEntry::Invalid),
Err(error) => Err(error),
})
.collect()
@ -65,7 +65,7 @@ pub trait BatchCache: BaseCache {
entries.push(match self.async_get_cache(&key, &context).await {
Ok(Some(value)) => BatchEntry::Hit(value),
Ok(None) => BatchEntry::Miss,
Err(Error::InvalidEntry) => BatchEntry::Invalid,
Err(Error::InvalidEntry(_)) => BatchEntry::Invalid,
Err(error) => return Err(error),
});
}

4
litellm-rust/crates/cache/src/codec.rs vendored
View file

@ -41,10 +41,10 @@ where
type Value = V;
fn encode(&self, value: &Self::Value) -> Result<Vec<u8>, Error> {
serde_json::to_vec(value).map_err(|_| Error::InvalidEntry)
serde_json::to_vec(value).map_err(Error::invalid_entry)
}
fn decode(&self, bytes: &[u8]) -> Result<Self::Value, Error> {
serde_json::from_slice(bytes).map_err(|_| Error::InvalidEntry)
serde_json::from_slice(bytes).map_err(Error::invalid_entry)
}
}

14
litellm-rust/crates/cache/src/dual.rs vendored
View file

@ -2,7 +2,7 @@ use std::{sync::Arc, time::Duration};
use crate::{
BaseCache, BatchCache, BatchEntry, BulkDeleteCache, CacheContext, ClaimCache, CounterCache,
DeleteCache, Error, FlushCache, IncrementOperation, SetCache, TtlCache,
DeleteCache, Error, FlushCache, IncrementOperation, SetCache, ShortBatch, TtlCache,
};
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
@ -98,7 +98,7 @@ impl<L1, L2> DualCache<L1, L2> {
fn remote<T>(&self, result: Result<T, Error>) -> Result<Option<T>, Error> {
match result {
Ok(value) => Ok(Some(value)),
Err(Error::Unavailable)
Err(Error::Unavailable(_))
if self.remote_failure_policy == RemoteFailurePolicy::UseLocal =>
{
Ok(None)
@ -136,7 +136,10 @@ where
remote: Vec<BatchEntry<V>>,
) -> Result<Vec<BatchEntry<V>>, Error> {
if missing.len() != remote.len() {
return Err(Error::Unavailable);
return Err(Error::unavailable(ShortBatch {
requested: missing.len(),
returned: remote.len(),
}));
}
for (index, entry) in missing.into_iter().zip(remote) {
if let BatchEntry::Hit(value) = &entry {
@ -389,7 +392,10 @@ where
}
let values = self.l2.async_increment_pipeline(operations.clone()).await?;
if values.len() != operations.len() {
return Err(Error::Unavailable);
return Err(Error::unavailable(ShortBatch {
requested: operations.len(),
returned: values.len(),
}));
}
for (operation, value) in operations.iter().zip(&values) {
self.l1

62
litellm-rust/crates/cache/src/error.rs vendored
View file

@ -1,9 +1,11 @@
use litellm_error::Source;
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
pub enum Error {
#[error("cache is unavailable")]
Unavailable,
Unavailable(#[source] Option<Source>),
#[error("invalid cache entry")]
InvalidEntry,
InvalidEntry(#[source] Option<Source>),
#[error("flushing Redis requires an explicit namespace")]
UnscopedFlush,
#[error("operation is not supported by this cache")]
@ -11,3 +13,59 @@ pub enum Error {
#[error("semantic cache requires request messages")]
MissingPrompt,
}
#[derive(Debug, thiserror::Error)]
#[error("cache backend returned {returned} results for {requested} keys")]
pub struct ShortBatch {
pub requested: usize,
pub returned: usize,
}
#[derive(Debug, thiserror::Error)]
#[error("cache lock is poisoned")]
struct Poisoned;
impl<T> From<std::sync::PoisonError<T>> for Error {
fn from(_: std::sync::PoisonError<T>) -> Self {
Self::unavailable(Poisoned)
}
}
impl Error {
pub fn unavailable(cause: impl std::error::Error + Send + Sync + 'static) -> Self {
Self::Unavailable(Some(Source::new(cause)))
}
pub fn invalid_entry(cause: impl std::error::Error + Send + Sync + 'static) -> Self {
Self::InvalidEntry(Some(Source::new(cause)))
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::Error;
#[rstest]
#[case::unavailable(
Error::unavailable(std::io::Error::other("connection refused")),
"cache is unavailable"
)]
#[case::invalid_entry(
Error::invalid_entry(std::io::Error::other("connection refused")),
"invalid cache entry"
)]
fn the_backend_cause_is_the_source_and_rendered_once(
#[case] error: Error,
#[case] message: &str,
) {
let report = litellm_error::Report::new(&error).to_string();
assert_eq!(report, format!("{message}: connection refused"));
assert!(!error.to_string().contains("connection refused"));
let cause = std::error::Error::source(&error)
.and_then(|cause| cause.downcast_ref::<litellm_error::Source>())
.expect("cause");
assert!(cause.downcast_ref::<std::io::Error>().is_some());
}
}

2
litellm-rust/crates/cache/src/lib.rs vendored
View file

@ -23,4 +23,4 @@ pub use codec::{CacheCodec, JsonCodec};
pub use dual::{
DEFAULT_DELETE_BATCH_SIZE, DualCache, ReadPolicy, RemoteFailurePolicy, WritePolicy,
};
pub use error::Error;
pub use error::{Error, ShortBatch};

10
litellm-rust/crates/cache/tests/caching.rs vendored
View file

@ -67,7 +67,7 @@ impl BaseCache for TestCache {
}
fn set_cache(&self, _: &str, _: Self::Value, _: &ExactCacheContext) -> Result<(), Error> {
Err(Error::Unavailable)
Err(Error::Unavailable(None))
}
async fn async_set_cache(
@ -77,7 +77,7 @@ impl BaseCache for TestCache {
context: ExactCacheContext,
) -> Result<(), Error> {
if key == "unavailable" {
return Err(Error::Unavailable);
return Err(Error::Unavailable(None));
}
self.writes
.lock()
@ -118,7 +118,7 @@ impl BaseCache for RecordingCounter {
impl CounterCache for RecordingCounter {
fn increment_cache(&self, key: &str, amount: f64, _: ExactCacheContext) -> Result<f64, Error> {
if key == "unavailable" {
return Err(Error::Unavailable);
return Err(Error::Unavailable(None));
}
let mut total = self.total.lock().unwrap();
*total += amount;
@ -226,7 +226,7 @@ async fn default_batch_operations_use_async_writes_and_stop_on_failure(cache: Te
context.clone(),
)
.await,
Err(Error::Unavailable)
Err(Error::Unavailable(None))
);
assert_eq!(
*cache.writes.lock().unwrap(),
@ -323,7 +323,7 @@ async fn default_increment_pipeline_stops_at_the_first_failure() {
operation("skipped", 1.0, None),
])
.await,
Err(Error::Unavailable)
Err(Error::Unavailable(None))
);
let keys = cache
.increments

10
litellm-rust/crates/cache/tests/codec.rs vendored
View file

@ -31,12 +31,18 @@ fn json_codec_round_trips_typed_domain_values() {
#[case::wrongly_typed(br#"{"deployment":12}"#)]
fn json_codec_rejects_malformed_and_wrongly_typed_entries(#[case] bytes: &[u8]) {
let codec = JsonCodec::<RoutingState>::new();
assert_eq!(codec.decode(bytes).unwrap_err(), Error::InvalidEntry);
assert!(matches!(
codec.decode(bytes),
Err(Error::InvalidEntry(Some(_)))
));
}
#[rstest]
fn json_codec_propagates_encoding_errors() {
let codec = JsonCodec::<BTreeMap<(u8, u8), String>>::new();
let value = BTreeMap::from([((1, 2), "invalid JSON object key".into())]);
assert_eq!(codec.encode(&value).unwrap_err(), Error::InvalidEntry);
assert!(matches!(
codec.encode(&value),
Err(Error::InvalidEntry(Some(_)))
));
}

35
litellm-rust/crates/cache/tests/dual.rs vendored
View file

@ -71,7 +71,7 @@ where
impl CounterCache for TestCache<f64> {
fn increment_cache(&self, _: &str, amount: f64, _: ExactCacheContext) -> Result<f64, Error> {
if self.fail {
return Err(Error::Unavailable);
return Err(Error::Unavailable(None));
}
let mut value = self.value.lock().unwrap();
let incremented = value.unwrap_or_default() + amount;
@ -92,7 +92,7 @@ where
_: ExactCacheContext,
) -> Result<V, Error> {
if self.fail {
return Err(Error::Unavailable);
return Err(Error::Unavailable(None));
}
let mut value = self.value.lock().unwrap();
let winner = match value.as_ref() {
@ -113,7 +113,7 @@ fn failed_l2_increment_leaves_l1_unchanged() {
assert_eq!(
cache.increment_cache("counter", 2.0, ExactCacheContext::default()),
Err(Error::Unavailable)
Err(Error::Unavailable(None))
);
assert_eq!(
l1.get_cache("counter", &ExactCacheContext::default())
@ -276,11 +276,11 @@ impl BaseCache for Unavailable {
}
fn set_cache(&self, _: &str, _: String, _: &ExactCacheContext) -> Result<(), Error> {
Err(Error::Unavailable)
Err(Error::Unavailable(None))
}
fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result<Option<String>, Error> {
Err(Error::Unavailable)
Err(Error::Unavailable(None))
}
}
@ -288,13 +288,13 @@ impl BatchCache for Unavailable {}
impl DeleteCache for Unavailable {
fn delete_cache(&self, _: &str) -> Result<(), Error> {
Err(Error::Unavailable)
Err(Error::Unavailable(None))
}
}
impl FlushCache for Unavailable {
fn flush_cache(&self) -> Result<(), Error> {
Err(Error::Unavailable)
Err(Error::Unavailable(None))
}
}
@ -306,7 +306,7 @@ impl ClaimCache for Unavailable {
_: &[String],
_: ExactCacheContext,
) -> Result<String, Error> {
Err(Error::InvalidEntry)
Err(Error::InvalidEntry(None))
}
}
@ -316,9 +316,12 @@ fn remote_failure_policy_selects_propagation_or_the_local_tier() {
let strict = DualCache::new(Arc::new(TestCache::new(None, false)), Arc::new(Unavailable));
assert_eq!(
strict.set_cache("key", "value".into(), &context),
Err(Error::Unavailable)
Err(Error::Unavailable(None))
);
assert_eq!(
strict.get_cache("key", &context),
Err(Error::Unavailable(None))
);
assert_eq!(strict.get_cache("key", &context), Err(Error::Unavailable));
let l1 = Arc::new(TestCache::new(None, false));
let degraded = DualCache::new(l1.clone(), Arc::new(Unavailable))
@ -347,7 +350,7 @@ fn claim_fallback_does_not_hide_non_availability_errors() {
&[],
ExactCacheContext::default()
),
Err(Error::InvalidEntry)
Err(Error::InvalidEntry(None))
);
}
@ -419,7 +422,7 @@ impl Tier {
fn check(&self) -> Result<(), Error> {
if self.fail {
return Err(Error::Unavailable);
return Err(Error::Unavailable(None));
}
Ok(())
}
@ -598,7 +601,7 @@ async fn failed_async_l2_increment_leaves_l1_unchanged(log: Log) {
cache
.async_increment("counter", 2.0, seconds(30), true)
.await,
Err(Error::Unavailable)
Err(Error::Unavailable(None))
);
assert_eq!(
l1.counter("counter"),
@ -669,7 +672,7 @@ async fn failed_l2_increment_pipeline_leaves_l1_unchanged(log: Log) {
ttl: None,
}])
.await,
Err(Error::Unavailable)
Err(Error::Unavailable(None))
);
assert_eq!(events(&log), ["l2 pipeline a"]);
assert_eq!(l1.counter("a"), Some((1.0, ExactCacheContext::default())));
@ -713,7 +716,7 @@ async fn failed_local_set_add_never_reaches_l2(log: Log) {
cache
.async_set_cache_sadd("members", vec!["a".into()], None)
.await,
Err(Error::Unavailable)
Err(Error::Unavailable(None))
);
assert!(events(&log).is_empty());
}
@ -768,7 +771,7 @@ async fn bulk_delete_stops_before_l2_when_the_local_delete_fails(log: Log) {
);
assert_eq!(
cache.delete_cache_keys(vec!["a".into()]).await,
Err(Error::Unavailable)
Err(Error::Unavailable(None))
);
assert!(events(&log).is_empty());
}

1
litellm-rust/crates/config/Cargo.toml
View file

@ -14,6 +14,7 @@ strum.workspace = true
thiserror.workspace = true
[dev-dependencies]
litellm-error.workspace = true
rstest.workspace = true
serde_json.workspace = true
tempfile.workspace = true

11
litellm-rust/crates/config/src/error.rs
View file

@ -1,7 +1,14 @@
use std::{io, path::PathBuf};
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[error("could not read config")]
Read(#[from] std::io::Error),
#[error("could not read config {}", path.display())]
Read { path: PathBuf, source: io::Error },
#[error("invalid YAML config")]
Parse(#[from] serde_yaml_ng::Error),
#[error("invalid YAML config {}", path.display())]
ParseFile {
path: PathBuf,
source: serde_yaml_ng::Error,
},
}

35
litellm-rust/crates/config/src/includes.rs
View file

@ -15,11 +15,31 @@ struct Includes {
}
fn read(path: &Path) -> Result<Mapping, Error> {
Ok(serde_yaml_ng::from_str(&std::fs::read_to_string(path)?)?)
let text = std::fs::read_to_string(path).map_err(|source| Error::Read {
path: path.to_owned(),
source,
})?;
serde_yaml_ng::from_str(&text).map_err(|source| Error::ParseFile {
path: path.to_owned(),
source,
})
}
fn canonical(path: &Path) -> Result<PathBuf, Error> {
path.canonicalize().map_err(|source| Error::Read {
path: path.to_owned(),
source,
})
}
fn entries(config: &Mapping, path: &Path) -> Result<Vec<(String, PathBuf)>, Error> {
let includes: Includes = serde_yaml_ng::from_value(Value::Mapping(config.clone()))?;
let includes: Includes =
serde_yaml_ng::from_value(Value::Mapping(config.clone())).map_err(|source| {
Error::ParseFile {
path: path.to_owned(),
source,
}
})?;
Ok(includes
.include
.into_iter()
@ -28,7 +48,7 @@ fn entries(config: &Mapping, path: &Path) -> Result<Vec<(String, PathBuf)>, Erro
}
pub(super) fn load(path: &Path) -> Result<Value, Error> {
let root = path.canonicalize()?;
let root = canonical(path)?;
let mut merged = read(&root)?;
let mut pending: VecDeque<_> = entries(&merged, &root)?.into();
let mut loaded = BTreeSet::from([root.clone()]);
@ -36,12 +56,11 @@ pub(super) fn load(path: &Path) -> Result<Value, Error> {
while let Some((entry, declaring)) = pending.pop_front() {
let declared = declaring.parent().unwrap_or(Path::new(".")).join(&entry);
let fallback = root.parent().unwrap_or(Path::new(".")).join(&entry);
let location = if declared.exists() {
declared
let location = canonical(if declared.exists() {
&declared
} else {
fallback
}
.canonicalize()?;
&fallback
})?;
if !loaded.insert(location.clone()) {
continue;
}

10
litellm-rust/crates/config/tests/config.rs
View file

@ -84,9 +84,15 @@ fn rejects_malformed_and_incomplete_config(#[case] yaml: &str) {
#[rstest]
fn distinguishes_read_errors_from_parse_errors(directory: TempDir) {
let error = Config::load(directory.path().join("missing.yaml")).unwrap_err();
assert!(
litellm_error::Report::new(&error)
.to_string()
.contains("missing.yaml")
);
assert!(matches!(
Config::load(directory.path().join("missing.yaml")),
Err(Error::Read(error)) if error.kind() == std::io::ErrorKind::NotFound
error,
Error::Read { source, .. } if source.kind() == std::io::ErrorKind::NotFound
));
}

1
litellm-rust/crates/core-utils/Cargo.toml
View file

@ -7,6 +7,7 @@ repository.workspace = true
[dependencies]
fancy-regex.workspace = true
litellm-error.workspace = true
litellm-tracing.workspace = true
litellm-llms-types.workspace = true
serde.workspace = true

19
litellm-rust/crates/core-utils/src/call_arguments.rs
View file

@ -1,5 +1,7 @@
use std::ops::Deref;
use litellm_error::Source;
use serde::{Deserialize, Serialize, de::DeserializeOwned};
use serde_json::{Map, Value};
@ -17,9 +19,11 @@ impl CallArguments {
}
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
#[error("invalid argument: {path}")]
#[error("invalid argument {path}")]
pub struct ArgumentError {
pub path: String,
#[source]
source: Source,
}
pub fn parse_options<T: DeserializeOwned>(arguments: &CallArguments) -> Result<T, ArgumentError> {
@ -28,6 +32,7 @@ pub fn parse_options<T: DeserializeOwned>(arguments: &CallArguments) -> Result<T
);
serde_path_to_error::deserialize(deserializer).map_err(|error| ArgumentError {
path: error.path().to_string(),
source: Source::new(error.into_inner()),
})
}
@ -42,8 +47,8 @@ pub fn compose_body<B: Serialize>(
body: &B,
consumed: &[&str],
) -> Result<Value, crate::params::Error> {
let Value::Object(fields) =
serde_json::to_value(body).map_err(|_| crate::params::Error::Body)?
let Value::Object(fields) = serde_json::to_value(body)
.map_err(|error| crate::params::Error::Serialize(Source::new(error)))?
else {
return Err(crate::params::Error::Body);
};
@ -183,9 +188,9 @@ mod tests {
assert_eq!(arguments.get("enabled"), Some(&Value::Null));
assert_eq!(arguments.get("missing"), None);
let invalid = serde_json::from_value(json!({"enabled":0})).unwrap();
assert_eq!(
parse_options::<Options>(&invalid).err().unwrap().path,
"enabled"
);
let error = parse_options::<Options>(&invalid).err().unwrap();
assert_eq!(error.path, "enabled");
let report = litellm_error::Report::new(&error).to_string();
assert!(report.starts_with("invalid argument enabled: invalid type"));
}
}

16
litellm-rust/crates/core-utils/src/exception_mapping_utils/mod.rs
View file

@ -342,15 +342,15 @@ mod tests {
#[case::request_timed_out("Request timed out")]
#[case::timed_out_generating("Timed out generating response")]
#[case::read_operation("The read operation timed out")]
fn timeout_markers_win_over_every_family(#[case] marker: &str) {
fn timeout_markers_win_over_every_family(
#[case] marker: &str,
#[values("mistral", "vertex_ai", "cohere", "reducto")] provider: &str,
) {
let body = format!("rate limit invalid api token {marker}");
for provider in ["mistral", "vertex_ai", "cohere", "reducto"] {
assert_eq!(
mapped(provider, &http(429, &body)).error,
PublicError::Timeout { status: 408 },
"{provider}"
);
}
assert_eq!(
mapped(provider, &http(429, &body)).error,
PublicError::Timeout { status: 408 }
);
}
#[test]

2
litellm-rust/crates/core-utils/src/params.rs
View file

@ -4,6 +4,8 @@ pub enum Error {
ExtraBody,
#[error("invalid request: body must be a JSON object")]
Body,
#[error("invalid request: body could not be serialized")]
Serialize(#[source] litellm_error::Source),
}
use std::ops::{Deref, DerefMut};

2
litellm-rust/crates/core-utils/src/url_utils.rs
View file

@ -4,7 +4,7 @@ use url::Url;
#[derive(Debug, thiserror::Error)]
pub enum ApiUrlError {
#[error("invalid URL: {0}")]
#[error("invalid URL")]
Parse(#[from] url::ParseError),
#[error("URL cannot be used as a base")]
CannotBeBase,

3
litellm-rust/crates/cost/Cargo.toml
View file

@ -5,6 +5,9 @@ edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
thiserror.workspace = true
[dev-dependencies]
criterion.workspace = true
proptest.workspace = true

27
litellm-rust/crates/cost/src/error.rs Normal file
View file

@ -0,0 +1,27 @@
#[derive(Clone, Copy, Debug, Eq, PartialEq, thiserror::Error)]
pub enum Error {
#[error("model has no input rate")]
MissingInputRate,
#[error("model has no output rate")]
MissingOutputRate,
#[error("rate must be finite and nonnegative")]
InvalidRate,
#[error("region multiplier must be finite and positive")]
InvalidRegionMultiplier,
#[error("off-peak pricing needs a billing minute within the day")]
InvalidBillingTime,
#[error("off-peak window must start before it ends within the day")]
InvalidOffPeakWindow,
#[error("cached tokens exceed prompt tokens")]
CacheExceedsPrompt,
#[error("cache write details do not add up to the cache write tokens")]
InvalidCacheWriteDetails,
#[error("token count or cost overflows")]
TokenCountOverflow,
#[error("service tier is repeated or cannot carry its own rates")]
DuplicateTier,
#[error("prompt token threshold is repeated")]
DuplicateThreshold,
#[error("service tier is repeated or cannot carry its own rates within a threshold")]
DuplicateThresholdTier,
}

88
litellm-rust/crates/cost/src/lib.rs
View file

@ -1,3 +1,7 @@
mod error;
pub use error::Error;
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum Rate {
Missing,
@ -156,22 +160,6 @@ pub struct EffectiveRates {
pub cache_write_1h: f64,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum PricingError {
MissingInputRate,
MissingOutputRate,
InvalidRate,
InvalidRegionMultiplier,
InvalidBillingTime,
InvalidOffPeakWindow,
CacheExceedsPrompt,
InvalidCacheWriteDetails,
TokenCountOverflow,
DuplicateTier,
DuplicateThreshold,
DuplicateThresholdTier,
}
fn selected_tier(tier: ServiceTier) -> ServiceTier {
if tier == ServiceTier::Fast {
ServiceTier::Priority
@ -210,9 +198,9 @@ fn valid_rates(rates: Rates) -> bool {
})
}
fn validate_tiers(tiers: &[TierRates], duplicate: PricingError) -> Result<(), PricingError> {
fn validate_tiers(tiers: &[TierRates], duplicate: Error) -> Result<(), Error> {
if tiers.iter().any(|entry| !valid_rates(entry.rates)) {
return Err(PricingError::InvalidRate);
return Err(Error::InvalidRate);
}
if tiers.iter().enumerate().any(|(index, entry)| {
matches!(
@ -227,20 +215,20 @@ fn validate_tiers(tiers: &[TierRates], duplicate: PricingError) -> Result<(), Pr
Ok(())
}
pub fn compile(pricing: &Pricing<'_>) -> Result<PricingPlan, PricingError> {
pub fn compile(pricing: &Pricing<'_>) -> Result<PricingPlan, Error> {
if !valid_rates(pricing.standard) {
return Err(PricingError::InvalidRate);
return Err(Error::InvalidRate);
}
validate_tiers(pricing.tiers, PricingError::DuplicateTier)?;
validate_tiers(pricing.tiers, Error::DuplicateTier)?;
if let Some(window) = pricing.off_peak {
if window.start_utc_minute >= 1440
|| window.end_utc_minute > 1440
|| window.start_utc_minute >= window.end_utc_minute
{
return Err(PricingError::InvalidOffPeakWindow);
return Err(Error::InvalidOffPeakWindow);
}
if !valid_rates(window.rates) {
return Err(PricingError::InvalidRate);
return Err(Error::InvalidRate);
}
}
let mut thresholds: Vec<_> = pricing
@ -248,9 +236,9 @@ pub fn compile(pricing: &Pricing<'_>) -> Result<PricingPlan, PricingError> {
.iter()
.map(|entry| {
if !valid_rates(entry.standard) {
return Err(PricingError::InvalidRate);
return Err(Error::InvalidRate);
}
validate_tiers(entry.tiers, PricingError::DuplicateThresholdTier)?;
validate_tiers(entry.tiers, Error::DuplicateThresholdTier)?;
Ok(CompiledThreshold {
above_prompt_tokens: entry.above_prompt_tokens,
standard: entry.standard,
@ -263,7 +251,7 @@ pub fn compile(pricing: &Pricing<'_>) -> Result<PricingPlan, PricingError> {
.windows(2)
.any(|pair| pair[0].above_prompt_tokens == pair[1].above_prompt_tokens)
{
return Err(PricingError::DuplicateThreshold);
return Err(Error::DuplicateThreshold);
}
Ok(PricingPlan {
standard: pricing.standard,
@ -274,11 +262,7 @@ pub fn compile(pricing: &Pricing<'_>) -> Result<PricingPlan, PricingError> {
}
impl PricingPlan {
fn resolve_rates(
&self,
request: &Request,
threshold_tokens: u64,
) -> Result<Rates, PricingError> {
fn resolve_rates(&self, request: &Request, threshold_tokens: u64) -> Result<Rates, Error> {
let tier = selected_tier(request.service_tier);
let base = self
.tiers
@ -305,13 +289,13 @@ impl PricingPlan {
|| window.end_utc_minute > 1440
|| window.start_utc_minute >= window.end_utc_minute
{
return Err(PricingError::InvalidOffPeakWindow);
return Err(Error::InvalidOffPeakWindow);
}
let minute = request
.billed_at_utc_minute
.ok_or(PricingError::InvalidBillingTime)?;
.ok_or(Error::InvalidBillingTime)?;
if minute >= 1440 {
return Err(PricingError::InvalidBillingTime);
return Err(Error::InvalidBillingTime);
}
if (window.start_utc_minute..window.end_utc_minute).contains(&minute) {
Ok(window.rates.overlay(selected))
@ -322,26 +306,26 @@ impl PricingPlan {
}
}
fn checked_rate(rate: Rate, missing: PricingError) -> Result<f64, PricingError> {
fn checked_rate(rate: Rate, missing: Error) -> Result<f64, Error> {
let value = rate.value().ok_or(missing)?;
if !value.is_finite() || value < 0.0 {
return Err(PricingError::InvalidRate);
return Err(Error::InvalidRate);
}
Ok(value)
}
pub fn calculate(&self, request: &Request) -> Result<Cost, PricingError> {
pub fn calculate(&self, request: &Request) -> Result<Cost, Error> {
let usage = request.usage;
let cached = usage
.cache_read_tokens
.checked_add(usage.cache_write_tokens)
.ok_or(PricingError::TokenCountOverflow)?;
.ok_or(Error::TokenCountOverflow)?;
let (regular, threshold_tokens) = match usage.prompt_convention {
PromptConvention::IncludesCache => (
usage
.prompt_tokens
.checked_sub(cached)
.ok_or(PricingError::CacheExceedsPrompt)?,
.ok_or(Error::CacheExceedsPrompt)?,
usage.prompt_tokens,
),
PromptConvention::ExcludesCache => (
@ -349,7 +333,7 @@ impl PricingPlan {
usage
.prompt_tokens
.checked_add(cached)
.ok_or(PricingError::TokenCountOverflow)?,
.ok_or(Error::TokenCountOverflow)?,
),
};
let writes = match (usage.cache_write_5m_tokens, usage.cache_write_1h_tokens) {
@ -357,26 +341,20 @@ impl PricingPlan {
(Some(five), Some(one)) if five.checked_add(one) == Some(usage.cache_write_tokens) => {
(five, one)
}
_ => return Err(PricingError::InvalidCacheWriteDetails),
_ => return Err(Error::InvalidCacheWriteDetails),
};
let rates = self.resolve_rates(request, threshold_tokens)?;
let input = Self::checked_rate(rates.input, PricingError::MissingInputRate)?;
let output = Self::checked_rate(rates.output, PricingError::MissingOutputRate)?;
let read = Self::checked_rate(
rates.cache_read.or(rates.input),
PricingError::MissingInputRate,
)?;
let write = Self::checked_rate(
rates.cache_write.or(rates.input),
PricingError::MissingInputRate,
)?;
let input = Self::checked_rate(rates.input, Error::MissingInputRate)?;
let output = Self::checked_rate(rates.output, Error::MissingOutputRate)?;
let read = Self::checked_rate(rates.cache_read.or(rates.input), Error::MissingInputRate)?;
let write = Self::checked_rate(rates.cache_write.or(rates.input), Error::MissingInputRate)?;
let write_1h = Self::checked_rate(
rates.cache_write_1h.or(rates.cache_write).or(rates.input),
PricingError::MissingInputRate,
Error::MissingInputRate,
)?;
let multiplier = request.region_multiplier.unwrap_or(1.0);
if !multiplier.is_finite() || multiplier <= 0.0 {
return Err(PricingError::InvalidRegionMultiplier);
return Err(Error::InvalidRegionMultiplier);
}
let cost = Cost {
uncached_input: regular as f64 * input,
@ -394,12 +372,12 @@ impl PricingPlan {
},
};
if !cost.total().is_finite() {
return Err(PricingError::TokenCountOverflow);
return Err(Error::TokenCountOverflow);
}
Ok(cost)
}
}
pub fn calculate(pricing: &Pricing<'_>, request: &Request) -> Result<Cost, PricingError> {
pub fn calculate(pricing: &Pricing<'_>, request: &Request) -> Result<Cost, Error> {
compile(pricing)?.calculate(request)
}

20
litellm-rust/crates/cost/tests/calculation.rs
View file

@ -1,5 +1,5 @@
use litellm_cost::{
OffPeakRates, Pricing, PricingError, PromptConvention, Rate, Rates, Request, ServiceTier,
Error, OffPeakRates, Pricing, PromptConvention, Rate, Rates, Request, ServiceTier,
ThresholdPolicy, ThresholdRates, TierRates, Usage, calculate, compile,
};
use rstest::rstest;
@ -132,7 +132,7 @@ fn split_writes_and_invalid_accounting() {
};
assert_eq!(
calculate(&pricing(standard), &overlapping),
Err(PricingError::CacheExceedsPrompt)
Err(Error::CacheExceedsPrompt)
);
let incomplete = Request {
usage: Usage {
@ -143,7 +143,7 @@ fn split_writes_and_invalid_accounting() {
};
assert_eq!(
calculate(&pricing(standard), &incomplete),
Err(PricingError::InvalidCacheWriteDetails)
Err(Error::InvalidCacheWriteDetails)
);
}
@ -297,10 +297,10 @@ fn compile_rejects_ambiguous_rates() {
};
assert_eq!(
compile(&specification).err(),
Some(PricingError::DuplicateThreshold)
Some(Error::DuplicateThreshold)
);
let invalid = pricing(rates(Rate::Value(f64::NAN), Rate::Value(1.0)));
assert_eq!(compile(&invalid).err(), Some(PricingError::InvalidRate));
assert_eq!(compile(&invalid).err(), Some(Error::InvalidRate));
}
#[test]
@ -326,7 +326,7 @@ fn off_peak_is_one_non_wrapping_utc_window() {
};
assert_eq!(
calculate(&specification, &base),
Err(PricingError::InvalidBillingTime)
Err(Error::InvalidBillingTime)
);
assert_eq!(calculate(&specification, &start).unwrap().input(), 100.0);
assert_eq!(calculate(&specification, &end).unwrap().input(), 200.0);
@ -347,11 +347,11 @@ fn missing_rates_and_free_rates_remain_distinct() {
};
assert_eq!(
calculate(&pricing(Rates::EMPTY), &empty),
Err(PricingError::MissingInputRate)
Err(Error::MissingInputRate)
);
assert_eq!(
calculate(&pricing(rates(Rate::Value(0.0), Rate::Missing)), &empty),
Err(PricingError::MissingOutputRate)
Err(Error::MissingOutputRate)
);
assert_eq!(
calculate(&pricing(rates(Rate::Value(0.0), Rate::Value(0.0))), &empty)
@ -481,7 +481,7 @@ fn regional_multiplier_applies_after_input_components_are_summed() {
};
assert_eq!(
calculate(&pricing(standard), &invalid),
Err(PricingError::InvalidRegionMultiplier)
Err(Error::InvalidRegionMultiplier)
);
}
@ -531,5 +531,5 @@ fn compilation_sorts_thresholds_and_rejects_duplicate_tiers() {
thresholds: &[],
..specification
};
assert_eq!(compile(&invalid).err(), Some(PricingError::DuplicateTier));
assert_eq!(compile(&invalid).err(), Some(Error::DuplicateTier));
}

Some files were not shown because too many files have changed in this diff Show more