scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-26 06:03:19 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 35

Merge litellm_customer_proxy into litellm_customer_proxy_redis_round_trips

Browse source 
Keeps one-round-trip admission; the model capacity hook's check_and_add,
requests and read stay on Store, with their scripts sent by hash too.
This commit is contained in:
narmstrong-berry 2026-10-19 17:40:52 +00:00
commit 0811b75845
74 changed files with 14090 additions and 432 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
litellm-rust/Cargo.lock generated
View file

@ -3750,14 +3750,17 @@ version = "0.1.0"
dependencies = [
"axum",
"base64 0.22.1",
"bytes",
"envy",
"futures-util",
"http 1.4.2",
"http-body-util",
"litellm-auth-aws",
"litellm-auth-types",
"litellm-config",
"litellm-db",
"litellm-db-testing",
"litellm-framer",
"litellm-gateway-auth",
"litellm-gateway-hooks",
"litellm-gateway-inference",
@ -3770,11 +3773,13 @@ dependencies = [
"litellm-router",
"litellm-secrets",
"litellm-tracing",
"md-5",
"rand 0.8.7",
"redis",
"rstest",
"serde",
"serde_json",
"sha2 0.10.9",
"sqlx",
"strum",
"tempfile",
@ -3785,6 +3790,7 @@ dependencies = [
"tower-sessions-moka-store",
"tracing",
"uuid",
"wiremock",
]
[[package]]
@ -4055,6 +4061,7 @@ dependencies = [
"litellm-cache-memory",
"litellm-cache-response",
"litellm-core-utils",
"litellm-framer",
"litellm-host",
"litellm-host-native",
"litellm-http",
@ -4068,7 +4075,9 @@ dependencies = [
"rstest",
"serde_json",
"tokio",
"tokio-util",
"tracing",
"uuid",
"wiremock",
]

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

@ -8,7 +8,8 @@ use sha2::{Digest, Sha256};
use aws_credential_types::Credentials;
use aws_credential_types::provider::ProvideCredentials;
use aws_sigv4::http_request::{
SignableBody, SignableRequest, SigningParams, SigningSettings, sign,
PercentEncodingMode, SignableBody, SignableRequest, SigningParams, SigningSettings,
UriPathNormalizationMode, sign,
};
use aws_sigv4::sign::v4;
use aws_smithy_runtime_api::client::identity::Identity;
@ -464,6 +465,58 @@ pub fn sign_post(
service: &str,
credentials: &Credentials,
signing_time: SystemTime,
) -> Result<BTreeMap<String, String>, Error> {
sign_request(
"POST",
url,
body,
headers,
region,
service,
credentials,
signing_time,
SigningSettings::default(),
)
}
/// Signs an S3 request as botocore's `S3SigV4Auth` does: the path is signed as sent, encoded
/// once and not normalized, and every header given is covered.
pub fn sign_s3(
method: &str,
url: &str,
body: &[u8],
headers: &BTreeMap<String, String>,
region: &str,
credentials: &Credentials,
signing_time: SystemTime,
) -> Result<BTreeMap<String, String>, Error> {
let mut settings = SigningSettings::default();
settings.percent_encoding_mode = PercentEncodingMode::Single;
settings.uri_path_normalization_mode = UriPathNormalizationMode::Disabled;
sign_request(
method,
url,
body,
headers,
region,
"s3",
credentials,
signing_time,
settings,
)
}
#[allow(clippy::too_many_arguments)]
fn sign_request(
method: &str,
url: &str,
body: &[u8],
headers: &BTreeMap<String, String>,
region: &str,
service: &str,
credentials: &Credentials,
signing_time: SystemTime,
settings: SigningSettings,
) -> Result<BTreeMap<String, String>, Error> {
let identity: Identity = credentials.clone().into();
let params = v4::SigningParams::builder()
@ -471,14 +524,14 @@ pub fn sign_post(
.region(region)
.name(service)
.time(signing_time)
.settings(SigningSettings::default())
.settings(settings)
.build()
.map(SigningParams::from)
.map_err(|error| Error::AwsSigningParameters(error.to_string()))?;
let header_refs = headers
.iter()
.map(|(name, value)| (name.as_str(), value.as_str()));
let request = SignableRequest::new("POST", url, header_refs, SignableBody::Bytes(body))
let request = SignableRequest::new(method, url, header_refs, SignableBody::Bytes(body))
.map_err(|error| Error::AwsSignableRequest(error.to_string()))?;
let (instructions, _) = sign(request, &params)
.map_err(|error| Error::AwsSigning(error.to_string()))?
@ -935,6 +988,49 @@ mod tests {
);
}
/// The PUT Object example of the S3 SigV4 guide, "Signature Calculations for the
/// Authorization Header: Transferring Payload in a Single Chunk" (checked 2026-10).
#[test]
fn s3_requests_are_signed_like_the_s3_guide_example() {
let headers = BTreeMap::from([
(
"Date".to_string(),
"Fri, 24 May 2013 00:00:00 GMT".to_string(),
),
(
"x-amz-content-sha256".to_string(),
"44ce7dd67c959e0d3524ffac1771dfbba87d2b6b4b4e99e42034a8b803f8b072".to_string(),
),
(
"x-amz-storage-class".to_string(),
"REDUCED_REDUNDANCY".to_string(),
),
]);
let signed = sign_s3(
"PUT",
"https://examplebucket.s3.amazonaws.com/test%24file.text",
b"Welcome to Amazon S3.",
&headers,
"us-east-1",
&Credentials::new(
"AKIAIOSFODNN7EXAMPLE",
"wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
None,
None,
"test",
),
UNIX_EPOCH + Duration::from_secs(1_369_353_600),
)
.unwrap();
assert_eq!(
signed["Authorization"],
"AWS4-HMAC-SHA256 Credential=AKIAIOSFODNN7EXAMPLE/20130524/us-east-1/s3/aws4_request, \
SignedHeaders=date;host;x-amz-content-sha256;x-amz-date;x-amz-storage-class, \
Signature=98ad721746da40c64f1a55b78f14c238d841ea1380cd77a1b5971af0ece108bd"
);
}
#[test]
fn same_role_comparison_matches_partition_account_and_role() {
assert!(same_role_arns(

6
litellm-rust/crates/core-utils/src/core_helpers.rs
View file

@ -11,6 +11,9 @@ const FINISH_REASONS: &[(&str, &str)] = &[
("end_turn", "stop"),
("stop_sequence", "stop"),
("max_tokens", "length"),
("tool_use", "tool_calls"),
("tool_calls", "tool_calls"),
("function_call", "function_call"),
("refusal", "content_filter"),
("compaction", "length"),
("guardrail_intervened", "content_filter"),
@ -44,7 +47,9 @@ pub fn usage_from_parts(
cached_tokens: cache_read_tokens,
cache_creation_tokens,
text_tokens: input_tokens,
..PromptTokensDetails::default()
},
..ChatCompletionsUsage::default()
}
}
@ -74,6 +79,7 @@ mod tests {
#[case::end_turn("end_turn", "stop")]
#[case::stop_sequence("stop_sequence", "stop")]
#[case::max_tokens("max_tokens", "length")]
#[case::tool_use("tool_use", "tool_calls")]
#[case::refusal("refusal", "content_filter")]
#[case::guardrail_intervened("guardrail_intervened", "content_filter")]
#[case::content_filtered("content_filtered", "content_filter")]

3
litellm-rust/crates/core-utils/src/get_llm_provider_logic.rs
View file

@ -12,6 +12,9 @@ pub enum LlmProviders {
Anthropic,
#[strum(serialize = "aws_textract")]
AwsTextract,
/// Azure OpenAI (`azure/<deployment>`), distinct from Azure AI Foundry's `azure_ai`.
#[strum(serialize = "azure")]
Azure,
#[strum(serialize = "azure_ai")]
AzureAi,
#[strum(serialize = "bedrock")]

2
litellm-rust/crates/gateway-hooks/src/facts.rs
View file

@ -221,6 +221,8 @@ pub struct Outcome {
/// In dollars, from usage and the deployment's prices; `None` when either is missing.
pub cost: Option<f64>,
pub time_to_first_byte: Option<Duration>,
/// When the attempt that served, or failed last, was sent to the provider.
pub upstream_started_at: Option<SystemTime>,
pub ended_at: SystemTime,
/// The stream stopped before its end: the client went away or the upstream failed.
pub interrupted: bool,

58
litellm-rust/crates/gateway-hooks/src/hook.rs
View file

@ -146,7 +146,7 @@ pub struct RequestEdit {
allowed: Targets,
writer: &'static str,
header_ops: Vec<HeaderOp>,
body_fields: Vec<(String, Value)>,
body_fields: Vec<(Vec<String>, Value)>,
response_headers: Vec<(HeaderName, HeaderValue)>,
annotations: Annotations,
}
@ -183,8 +183,15 @@ impl RequestEdit {
/// Sets a top-level field of the upstream body, replacing what the client sent.
pub fn set_upstream_body_field(&mut self, field: impl Into<String>, value: Value) {
if self.may(Targets::UPSTREAM_BODY) {
self.body_fields.push((field.into(), value));
self.set_upstream_body_path(vec![field.into()], value);
}
/// Sets a field inside a top-level object of the upstream body, such as
/// `["metadata", "user_id"]`, keeping the object's other fields. A missing or non-object
/// step becomes an object.
pub fn set_upstream_body_path(&mut self, path: Vec<String>, value: Value) {
if !path.is_empty() && self.may(Targets::UPSTREAM_BODY) {
self.body_fields.push((path, value));
}
}
@ -202,7 +209,11 @@ impl RequestEdit {
/// Whether any change reaches the provider.
pub fn changes_upstream(&self) -> bool {
!self.header_ops.is_empty() || !self.body_fields.is_empty()
self.changes_upstream_headers() || !self.body_fields.is_empty()
}
pub fn changes_upstream_headers(&self) -> bool {
!self.header_ops.is_empty()
}
/// Applies the header changes in order to the headers about to be sent.
@ -221,8 +232,24 @@ impl RequestEdit {
/// Applies the body changes in order to the body about to be sent.
pub fn apply_upstream_body(&self, body: &mut Map<String, Value>) {
for (field, value) in &self.body_fields {
body.insert(field.clone(), value.clone());
for (path, value) in &self.body_fields {
let Some((last, parents)) = path.split_last() else {
continue;
};
let mut object = &mut *body;
for step in parents {
let slot = object
.entry(step.clone())
.or_insert_with(|| Value::Object(Map::new()));
if !slot.is_object() {
*slot = Value::Object(Map::new());
}
let Value::Object(inner) = slot else {
unreachable!("just made an object")
};
object = inner;
}
object.insert(last.clone(), value.clone());
}
}
@ -328,6 +355,25 @@ mod tests {
);
}
#[rstest]
fn nested_body_fields_keep_their_siblings() {
let mut edit = RequestEdit::default();
edit.scope("hook", Targets::UPSTREAM_BODY);
edit.set_upstream_body_path(vec!["metadata".into(), "user_id".into()], Value::from("u"));
edit.set_upstream_body_path(vec!["extra".into(), "a".into()], Value::from(1));
let serde_json::Value::Object(mut body) =
serde_json::json!({"metadata": {"tag": "t"}, "extra": "not an object"})
else {
unreachable!()
};
edit.apply_upstream_body(&mut body);
assert_eq!(
Value::Object(body),
serde_json::json!({"metadata": {"tag": "t", "user_id": "u"}, "extra": {"a": 1}})
);
}
#[rstest]
fn undeclared_writes_are_dropped() {
let mut edit = RequestEdit::default();

1
litellm-rust/crates/gateway-hooks/src/observer.rs
View file

@ -146,6 +146,7 @@ mod tests {
usage: None,
cost: None,
time_to_first_byte: None,
upstream_started_at: None,
ended_at: SystemTime::now(),
interrupted: false,
},

2
litellm-rust/crates/gateway-hooks/src/pipeline.rs
View file

@ -80,6 +80,7 @@ impl Pipeline {
usage: None,
cost: None,
time_to_first_byte: None,
upstream_started_at: None,
ended_at: SystemTime::now(),
interrupted: false,
};
@ -185,6 +186,7 @@ mod tests {
}),
cost: None,
time_to_first_byte: None,
upstream_started_at: None,
ended_at: SystemTime::now(),
interrupted: false,
}

239
litellm-rust/crates/gateway-hooks/src/pricing.rs
View file

@ -18,6 +18,12 @@ use crate::usage::{ServiceTier, Usage};
pub struct Prices {
plans: HashMap<Arc<str>, PricingPlan>,
max_output_tokens: HashMap<Arc<str>, u64>,
/// Deployments priced by the model their response names, with the provider route the
/// cost map keys that model under. The proxy prices an Azure deployment that sets neither
/// `base_model` nor prices of its own this way.
by_response_model: HashMap<Arc<str>, Arc<str>>,
/// Cost map prices by model name, for the deployments priced by their response's model.
models: HashMap<String, PricingPlan>,
}
impl Prices {
@ -27,18 +33,29 @@ impl Prices {
if let Some(limit) = fields.get("max_output_tokens").and_then(Value::as_u64) {
self.max_output_tokens.insert(id.clone(), limit);
}
let parsed = ParsedRates::parse(fields);
if parsed.standard == Rates::EMPTY && parsed.thresholds.is_empty() {
return;
if let Some(plan) = plan(&id, fields) {
self.plans.insert(id, plan);
}
match parsed.compile() {
Ok(plan) => {
self.plans.insert(id, plan);
}
Err(error) => tracing::warn!(deployment = %id, ?error, "ignoring invalid prices"),
}
/// Prices `deployment` by the model its response names, looked up as
/// `{provider}/{model}`, before its own prices.
pub fn price_by_response_model(&mut self, deployment: Arc<str>, provider: &str) {
self.by_response_model.insert(deployment, provider.into());
}
/// Adds the cost map prices of `model`, for deployments priced by their response's model.
pub fn insert_model(&mut self, model: &str, fields: &Map<String, Value>) {
if let Some(plan) = plan(model, fields) {
self.models.insert(model.to_owned(), plan);
}
}
/// Whether some deployment is priced by its response's model.
pub fn prices_by_response_model(&self) -> bool {
!self.by_response_model.is_empty()
}
pub fn len(&self) -> usize {
self.plans.len()
}
@ -54,41 +71,78 @@ impl Prices {
/// The cost in dollars of `usage` on a deployment, or `None` when its prices are unknown.
pub fn cost(&self, deployment: &str, usage: &Usage) -> Option<f64> {
let plan = self.plans.get(deployment)?;
let one_hour = usage.cache_write_1h_tokens.min(usage.cache_write_tokens);
let split = (one_hour > 0).then_some(());
let request = Request {
usage: litellm_cost::Usage {
prompt_tokens: usage.prompt_tokens(),
completion_tokens: usage.output_tokens,
cache_read_tokens: usage.cache_read_tokens,
cache_write_tokens: usage.cache_write_tokens,
cache_write_5m_tokens: split.map(|()| usage.cache_write_tokens - one_hour),
cache_write_1h_tokens: split.map(|()| one_hour),
prompt_convention: PromptConvention::IncludesCache,
},
service_tier: match usage.service_tier {
ServiceTier::Priority => CostTier::Priority,
ServiceTier::Flex => CostTier::Flex,
ServiceTier::Standard | ServiceTier::Batch => CostTier::Standard,
},
threshold_policy: ThresholdPolicy::Exclusive,
region_multiplier: None,
billed_at_utc_minute: None,
};
match plan.calculate(&request) {
Ok(cost) => Some(cost.total()),
Err(error) => {
tracing::debug!(deployment, ?error, "no cost for response");
None
}
plan_cost(deployment, self.plans.get(deployment)?, usage)
}
/// [`Self::cost`], first by `response_model` for a deployment priced by its response's
/// model, as the proxy tries `{provider}/{response model}` ahead of the deployment.
pub fn response_cost(
&self,
deployment: &str,
response_model: Option<&str>,
usage: &Usage,
) -> Option<f64> {
self.by_response_model
.get(deployment)
.zip(response_model.filter(|model| !model.is_empty()))
.and_then(|(provider, model)| {
let name = match model.starts_with(&format!("{provider}/")) {
true => model.to_owned(),
false => format!("{provider}/{model}"),
};
plan_cost(deployment, self.models.get(&name)?, usage)
})
.or_else(|| self.cost(deployment, usage))
}
}
/// The plan `fields` price, or `None` when they hold no usable prices.
fn plan(name: &str, fields: &Map<String, Value>) -> Option<PricingPlan> {
let parsed = ParsedRates::parse(fields);
if parsed.standard == Rates::EMPTY && parsed.thresholds.is_empty() {
return None;
}
parsed
.compile()
.inspect_err(|error| tracing::warn!(deployment = %name, ?error, "ignoring invalid prices"))
.ok()
}
fn plan_cost(deployment: &str, plan: &PricingPlan, usage: &Usage) -> Option<f64> {
let one_hour = usage.cache_write_1h_tokens.min(usage.cache_write_tokens);
let split = (one_hour > 0).then_some(());
let request = Request {
usage: litellm_cost::Usage {
prompt_tokens: usage.prompt_tokens(),
completion_tokens: usage.output_tokens,
cache_read_tokens: usage.cache_read_tokens,
cache_write_tokens: usage.cache_write_tokens,
cache_write_5m_tokens: split.map(|()| usage.cache_write_tokens - one_hour),
cache_write_1h_tokens: split.map(|()| one_hour),
prompt_convention: PromptConvention::IncludesCache,
},
service_tier: match usage.service_tier {
ServiceTier::Priority => CostTier::Priority,
ServiceTier::Flex => CostTier::Flex,
ServiceTier::Standard | ServiceTier::Batch => CostTier::Standard,
},
threshold_policy: ThresholdPolicy::Exclusive,
region_multiplier: None,
billed_at_utc_minute: None,
};
match plan.calculate(&request) {
Ok(cost) => Some(cost.total()),
Err(error) => {
tracing::debug!(deployment, ?error, "no cost for response");
None
}
}
}
/// A deployment's cost map fields: the cost map entry for its model, overlaid by the prices its
/// config sets. The entry is found by `base_model`, then the model name, then the name without
/// its provider route.
/// its provider route. An Azure deployment's names are keyed under `azure/` first, and its bare
/// deployment name is never looked up, since it names a resource of the user's, not a model.
pub fn deployment_fields(
catalog: Option<&Catalog>,
model: &str,
@ -97,14 +151,10 @@ pub fn deployment_fields(
) -> Map<String, Value> {
let mut fields = catalog
.and_then(|catalog| {
let candidates = base_model.into_iter().chain([
model,
model.split_once('/').map_or(model, |(_, rest)| rest),
model.rsplit_once('/').map_or(model, |(_, rest)| rest),
]);
candidates
candidates(model, base_model)
.into_iter()
.filter(|name| !name.is_empty())
.find_map(|name| catalog.lookup(name))
.find_map(|name| catalog.lookup(&name))
})
.map(|found| found.entry.fields().clone())
.unwrap_or_default();
@ -118,6 +168,30 @@ pub fn deployment_fields(
fields
}
fn candidates(model: &str, base_model: Option<&str>) -> Vec<String> {
if model.starts_with(AZURE_ROUTE) {
return base_model
.into_iter()
.flat_map(|base| match base.starts_with(AZURE_ROUTE) {
true => vec![base.to_owned()],
false => vec![format!("{AZURE_ROUTE}{base}"), base.to_owned()],
})
.chain([model.to_owned()])
.collect();
}
base_model
.into_iter()
.chain([
model,
model.split_once('/').map_or(model, |(_, rest)| rest),
model.rsplit_once('/').map_or(model, |(_, rest)| rest),
])
.map(str::to_owned)
.collect()
}
const AZURE_ROUTE: &str = "azure/";
#[derive(Clone, Copy)]
enum Slot {
Input,
@ -382,4 +456,81 @@ mod tests {
assert_eq!(fields["output_cost_per_token"], json!(2e-6));
assert!(!fields.contains_key("api_key"));
}
fn azure_catalog() -> Catalog {
Catalog::parse(
br#"{
"gpt-4o": {"input_cost_per_token": 1e-6, "output_cost_per_token": 1e-6},
"azure/gpt-4o": {"input_cost_per_token": 2e-6, "output_cost_per_token": 2e-6},
"gpt-4o-mini": {"input_cost_per_token": 3e-6, "output_cost_per_token": 3e-6},
"azure/my-dep": {"input_cost_per_token": 4e-6, "output_cost_per_token": 4e-6},
"chat-dep": {"input_cost_per_token": 5e-6, "output_cost_per_token": 5e-6}
}"#,
Default::default(),
)
.unwrap()
}
#[rstest]
#[case::base_model_under_azure("azure/chat-dep", Some("gpt-4o"), Some(2e-6))]
#[case::base_model_already_routed("azure/chat-dep", Some("azure/gpt-4o"), Some(2e-6))]
#[case::bare_base_model("azure/chat-dep", Some("gpt-4o-mini"), Some(3e-6))]
#[case::deployment_under_azure("azure/my-dep", None, Some(4e-6))]
#[case::bare_deployment_is_not_a_model("azure/chat-dep", None, None)]
fn azure_deployments_are_priced_under_the_azure_route(
#[case] model: &str,
#[case] base_model: Option<&str>,
#[case] input: Option<f64>,
) {
let fields = deployment_fields(Some(&azure_catalog()), model, base_model, &[]);
assert_eq!(
fields.get("input_cost_per_token").and_then(Value::as_f64),
input
);
}
fn response_priced(own: Value) -> Prices {
let mut prices = Prices::default();
prices.insert("d".into(), own.as_object().unwrap());
prices.insert_model(
"azure/gpt-4o",
json!({"input_cost_per_token": 2e-6, "output_cost_per_token": 0.0})
.as_object()
.unwrap(),
);
prices.price_by_response_model("d".into(), "azure");
prices
}
#[rstest]
#[case::response_model_under_the_route(Some("gpt-4o"), Some(2e-6))]
#[case::response_model_already_routed(Some("azure/gpt-4o"), Some(2e-6))]
#[case::unknown_response_model_uses_the_deployment(Some("gpt-9"), Some(7e-6))]
#[case::no_response_model_uses_the_deployment(None, Some(7e-6))]
fn response_model_prices_come_first(
#[case] response_model: Option<&str>,
#[case] rate: Option<f64>,
) {
let prices =
response_priced(json!({"input_cost_per_token": 7e-6, "output_cost_per_token": 0.0}));
assert_eq!(
prices.response_cost("d", response_model, &usage(1, 0)),
rate
);
}
#[rstest]
fn other_deployments_ignore_the_response_model() {
let mut prices = response_priced(json!({}));
prices.insert(
"other".into(),
json!({"input_cost_per_token": 7e-6, "output_cost_per_token": 0.0})
.as_object()
.unwrap(),
);
assert_eq!(
prices.response_cost("other", Some("gpt-4o"), &usage(1, 0)),
Some(7e-6)
);
}
}

51
litellm-rust/crates/gateway-hooks/src/usage.rs
View file

@ -142,6 +142,8 @@ impl AnthropicUsage {
}
}
/// OpenAI's usage, with the fields LiteLLM adds when it translates another provider's:
/// cache writes by lifetime and server tool counts.
#[derive(Deserialize)]
struct OpenAiUsage {
#[serde(default)]
@ -149,23 +151,38 @@ struct OpenAiUsage {
#[serde(default)]
completion_tokens: u64,
prompt_tokens_details: Option<OpenAiPromptDetails>,
server_tool_use: Option<ServerToolUse>,
}
#[derive(Deserialize, Default)]
struct OpenAiPromptDetails {
cached_tokens: Option<u64>,
cache_creation_tokens: Option<u64>,
cache_creation_token_details: Option<CacheCreation>,
}
impl OpenAiUsage {
fn usage(self) -> Usage {
let details = self.prompt_tokens_details.unwrap_or_default();
Usage::from_prompt_counts(
let mut usage = Usage::from_prompt_counts(
self.prompt_tokens,
self.completion_tokens,
details.cached_tokens.unwrap_or(0),
details.cache_creation_tokens.unwrap_or(0),
)
);
if let Some(tokens) = details
.cache_creation_token_details
.and_then(|cache| cache.ephemeral_1h_input_tokens)
{
usage.cache_write_1h_tokens = tokens.min(usage.cache_write_tokens);
}
if let Some(requests) = self
.server_tool_use
.and_then(|tools| tools.web_search_requests)
{
usage.web_search_requests = requests;
}
usage
}
}
@ -197,6 +214,7 @@ pub struct StreamUsage {
line: Vec<u8>,
state: Line,
usage: Option<Usage>,
model: Option<String>,
}
impl StreamUsage {
@ -206,6 +224,7 @@ impl StreamUsage {
line: Vec::new(),
state: Line::Reading,
usage: None,
model: None,
}
}
@ -264,6 +283,11 @@ impl StreamUsage {
self.usage
}
/// The model named by the chunk that carried an OpenAI stream's usage.
pub fn model(&self) -> Option<&str> {
self.model.as_deref()
}
fn may_carry_usage(&self) -> bool {
let line = &self.line;
if line.len() > MAX_LINE_BYTES {
@ -304,8 +328,13 @@ impl StreamUsage {
if !contains(data, b"\"usage\":{") && !contains(data, b"\"usage\": {") {
return;
}
if let Ok(OpenAiChunk { usage: Some(usage) }) = serde_json::from_slice(data) {
if let Ok(OpenAiChunk {
usage: Some(usage),
model,
}) = serde_json::from_slice(data)
{
self.usage = Some(usage.usage());
self.model = model.filter(|model| !model.is_empty());
}
}
}
@ -326,6 +355,7 @@ struct AnthropicMessage {
#[derive(Deserialize)]
struct OpenAiChunk {
usage: Option<OpenAiUsage>,
model: Option<String>,
}
fn contains(haystack: &[u8], needle: &[u8]) -> bool {
@ -398,7 +428,7 @@ mod tests {
fn openai_usage_comes_from_the_last_chunk() {
let stream = concat!(
"data: {\"id\":\"c\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hi\"}}],\"usage\":null}\n\n",
"data: {\"id\":\"c\",\"choices\":[],\"usage\":{\"prompt_tokens\":30,\"completion_tokens\":5,\"total_tokens\":35,\"prompt_tokens_details\":{\"cached_tokens\":10}}}\n\n",
"data: {\"id\":\"c\",\"model\":\"gpt-4o-2024-11-20\",\"choices\":[],\"usage\":{\"prompt_tokens\":30,\"completion_tokens\":5,\"total_tokens\":35,\"prompt_tokens_details\":{\"cached_tokens\":10}}}\n\n",
"data: [DONE]\n\n",
);
let mut meter = StreamUsage::new(StreamFormat::OpenAi);
@ -414,6 +444,7 @@ mod tests {
..Usage::default()
})
);
assert_eq!(meter.model(), Some("gpt-4o-2024-11-20"));
}
#[rstest]
@ -431,6 +462,18 @@ mod tests {
assert_eq!(anthropic_usage().total_tokens(), 78);
}
#[rstest]
fn translated_openai_usage_keeps_cache_lifetimes_and_server_tools() {
let usage = Usage::from_openai(&json!({
"prompt_tokens": 36,
"completion_tokens": 42,
"prompt_tokens_details": {"cached_tokens": 20, "cache_creation_tokens": 6,
"cache_creation_token_details": {"ephemeral_5m_input_tokens": 2, "ephemeral_1h_input_tokens": 4}},
"server_tool_use": {"web_search_requests": 2},
}));
assert_eq!(usage, Some(anthropic_usage()));
}
#[rstest]
fn openai_body_usage() {
let usage = Usage::from_openai(&json!({

76
litellm-rust/crates/gateway-inference/src/chat_completions.rs
View file

@ -12,6 +12,7 @@ use litellm_host_http::Sse;
use litellm_inference_chat::{
ChatCompletionsRoute,
openai::{OpenAiChatCall, OpenAiChatCompletions},
stream::{ChatCompletionsStream, ChatStreamCall, ReportedUsage},
types::ChatCompletionsCall,
};
use serde_json::{Map, Value};
@ -93,6 +94,18 @@ async fn handle(
None => route,
};
// Chat forwards no client headers of its own; hooks may set some.
let extra_headers = tracking
.as_ref()
.and_then(|tracking| tracking.edit_headers(Vec::new()))
.filter(|headers| !headers.is_empty())
.map(|headers| {
headers
.into_iter()
.map(|(name, value)| (name, Value::from(value)))
.collect::<Map<String, Value>>()
});
let extra_headers = &extra_headers;
let cache_headers = crate::caching::CacheHeaders::default();
let mut body = Some(body);
let routed = gateway
@ -120,9 +133,10 @@ async fn handle(
&deployment.model,
deployment.custom_llm_provider.as_deref(),
) {
return openai(gateway, deployment, body).await;
return openai(gateway, deployment, body, extra_headers.clone()).await;
}
let (messages, optional_params) = split(body);
let stream = ChatCompletionsRoute::is_stream(&optional_params);
let call = ChatCompletionsCall {
model: deployment.model.clone(),
messages,
@ -130,9 +144,12 @@ async fn handle(
api_key: deployment.api_key.clone(),
api_base: deployment.api_base.clone(),
custom_llm_provider: deployment.custom_llm_provider.clone(),
extra_headers: None,
extra_headers: extra_headers.clone(),
timeout: deployment.timeout,
};
if stream {
return translated_stream(route, call, name).await;
}
Ok(litellm_host_http::serve_unary(
route.machine(call, cache_options.policy),
(),
@ -162,10 +179,29 @@ fn split(mut body: Map<String, Value>) -> (Value, Map<String, Value>) {
(messages, body)
}
/// The deployment's `litellm_params` a provider reads beyond its connection: its
/// `api_version`, `base_model` and every key the config adds.
fn provider_params(deployment: &Deployment) -> Map<String, Value> {
let params = &deployment.litellm_params;
let mut provider_params: Map<String, Value> = params
.extra
.iter()
.map(|(name, value)| (name.clone(), value.clone()))
.collect();
if let Some(api_version) = &params.api_version {
provider_params.insert("api_version".into(), api_version.clone().into());
}
if let Some(base_model) = &deployment.base_model {
provider_params.insert("base_model".into(), base_model.clone().into());
}
provider_params
}
async fn openai(
gateway: &Gateway,
deployment: &Deployment,
body: Map<String, Value>,
extra_headers: Option<Map<String, Value>>,
) -> Result<Response, Error> {
let machine = gateway.chat_completions.clone().openai_machine(
OpenAiChatCall {
@ -174,8 +210,16 @@ async fn openai(
api_key: deployment.api_key.clone(),
api_base: deployment.api_base.clone(),
custom_llm_provider: deployment.custom_llm_provider.clone(),
extra_headers: None,
extra_headers,
timeout: deployment.timeout,
litellm_params: provider_params(deployment),
drop_params: gateway.drop_params()
|| deployment
.litellm_params
.drop_params
.as_ref()
.is_some_and(|flag| flag.is_set()),
model_info: gateway.model_info(),
},
None,
);
@ -185,3 +229,29 @@ async fn openai(
);
Ok(litellm_host_http::serve(machine, (), (), encoder, None).await?)
}
/// A stream from a provider chat translates (Anthropic, Bedrock), sent as OpenAI chunks. The
/// usage it ends with rides on the response for the hooks, even when the client did not ask
/// to see it.
async fn translated_stream(
route: ChatCompletionsRoute,
call: ChatCompletionsCall,
model: &str,
) -> Result<Response, Error> {
let usage = ReportedUsage::default();
let machine = route.stream_machine(
ChatStreamCall {
call,
response_model: model.to_owned(),
usage: usage.clone(),
},
None,
);
let encoder = Sse::<ChatCompletionsStream, _, _>::new(
|body: Bytes| ([(CONTENT_TYPE, "application/json")], body),
|error| Bytes::from(Error::from(error).openai_sse_frame()),
);
let mut response = litellm_host_http::serve(machine, (), (), encoder, None).await?;
response.extensions_mut().insert(usage);
Ok(response)
}

78
litellm-rust/crates/gateway-inference/src/error.rs
View file

@ -121,12 +121,18 @@ impl Error {
}
_ => self.to_string(),
};
json!({"error": {
let mut error = json!({
"message": message,
"type": error_type(status),
"param": null,
"code": status.as_u16(),
}})
});
if let Self::Route(RouteError::Transport(TransportError::Http { body, .. })) = self
&& let Some(fields) = azure_content_filter_fields(body)
{
error["provider_specific_fields"] = fields;
}
json!({ "error": error })
}
/// The Anthropic error envelope Python's `AnthropicExceptionMapping` builds: an upstream
@ -233,6 +239,44 @@ fn openai_error_message(body: &str) -> Option<String> {
.map(str::to_owned)
}
/// Python's `AzureOpenAIExceptionMapping` keeps Azure's error and its content filter results
/// on a content policy violation. Only Azure nests an `inner_error`, which tells its errors
/// apart from another provider's.
fn azure_content_filter_fields(body: &str) -> Option<Value> {
const POLICY_PHRASES: [&str; 7] = [
"content_policy_violation",
"responsibleaipolicyviolation",
"the response was filtered due to the prompt triggering azure openai's content management",
"your task failed as a result of our safety system",
"the model produced invalid content",
"content_filter_policy",
"your request was rejected as a result of our safety system",
];
let parsed: Value = serde_json::from_str(body).ok()?;
let azure_error = parsed
.get("error")
.filter(|error| error.is_object())
.unwrap_or(&parsed);
let inner_error = ["inner_error", "innererror"].into_iter().find_map(|key| {
azure_error
.get(key)
.or_else(|| parsed.get(key))
.filter(|inner| !inner.is_null())
})?;
let lowered = body.to_lowercase();
let violation = azure_error.get("code").and_then(Value::as_str)
== Some("content_policy_violation")
|| inner_error.get("code").and_then(Value::as_str) == Some("ResponsibleAIPolicyViolation")
|| POLICY_PHRASES.iter().any(|phrase| lowered.contains(phrase));
violation.then(|| {
json!({
"innererror": inner_error,
"inner_error": inner_error,
"azure_error": azure_error,
})
})
}
fn provider_message(object: &Map<String, Value>) -> Option<&str> {
if let Some(detail) = object.get("detail").and_then(Value::as_object) {
return detail.get("message").and_then(Value::as_str);
@ -326,6 +370,36 @@ mod tests {
assert_eq!(frame, format!("data: {expected}\n\n"));
}
#[rstest]
#[case::content_filter(
r#"{"error":{"message":"The response was filtered due to the prompt triggering Azure OpenAI's content management policy.","type":null,"param":"prompt","code":"content_filter","status":400,"innererror":{"code":"ResponsibleAIPolicyViolation","content_filter_result":{"hate":{"filtered":true,"severity":"high"}}}}}"#,
true,
)]
#[case::inner_error_spelling(
r#"{"error":{"message":"blocked","code":"content_policy_violation","inner_error":{"code":"ResponsibleAIPolicyViolation"}}}"#,
true,
)]
#[case::other_azure_error(
r#"{"error":{"message":"bad value","code":"invalid_value","innererror":{"code":"Other"}}}"#,
false
)]
#[case::openai_policy_error(
r#"{"error":{"message":"rejected","code":"content_policy_violation","type":"invalid_request_error"}}"#,
false,
)]
fn azure_content_filter_errors_keep_the_azure_error(#[case] body: &str, #[case] kept: bool) {
let parsed: Value = serde_json::from_str(body).unwrap();
let error = upstream(400, body).openai_body();
let fields = &error["error"]["provider_specific_fields"];
assert_eq!(!fields.is_null(), kept);
if kept {
assert_eq!(fields["azure_error"], parsed["error"]);
assert_eq!(fields["innererror"], fields["inner_error"]);
assert_eq!(fields["innererror"]["code"], "ResponsibleAIPolicyViolation");
}
assert_eq!(error["error"]["message"], parsed["error"]["message"]);
}
#[rstest]
#[case::upstream_status(upstream(429, ""), StatusCode::TOO_MANY_REQUESTS)]
#[case::rejected_request(Error::Route(RouteError::InvalidRequest("top_k".into())), StatusCode::BAD_REQUEST)]

98
litellm-rust/crates/gateway-inference/src/hooks.rs
View file

@ -19,6 +19,7 @@ use litellm_gateway_hooks::{
Admission, DeploymentFacts, Facts, KeyFacts, Outcome, Pipeline, RequestFacts, Route,
StreamFormat, StreamUsage, Usage,
};
use litellm_inference_chat::stream::ReportedUsage;
use litellm_router::Attempt;
use serde_json::{Map, Value};
@ -52,7 +53,8 @@ pub(crate) struct Tracking {
pipeline: Arc<Pipeline>,
route: Route,
admission: Admission,
deployment: Mutex<Option<DeploymentFacts>>,
/// The deployment the latest attempt went to, and when it was sent.
deployment: Mutex<Option<(DeploymentFacts, SystemTime)>>,
started: Instant,
}
@ -154,6 +156,20 @@ impl Tracking {
self.admission.edit().apply_upstream_body(body);
}
/// The headers the gateway would forward, after the hooks' changes; `None` when no hook
/// changes upstream headers.
pub(crate) fn edit_headers(
&self,
forwarded: Vec<(String, String)>,
) -> Option<Vec<(String, String)>> {
let edit = self.admission.edit();
edit.changes_upstream_headers().then(|| {
let mut forwarded = forwarded;
edit.apply_upstream_headers(&mut forwarded);
forwarded
})
}
/// Notes the deployment an attempt went to; the last one is the one that served.
pub(crate) fn attempt(&self, name: &str, attempt: &Attempt<'_>) {
if !self.pipeline.reads(Facts::DEPLOYMENT) {
@ -167,7 +183,8 @@ impl Tracking {
provider: deployment.custom_llm_provider.clone(),
api_base: deployment.api_base.clone(),
};
*self.deployment.lock().unwrap_or_else(|e| e.into_inner()) = Some(facts);
*self.deployment.lock().unwrap_or_else(|e| e.into_inner()) =
Some((facts, SystemTime::now()));
}
/// Adds the hooks' response headers and follows the body to its end; an error is recorded
@ -180,7 +197,10 @@ impl Tracking {
deployment,
started,
} = self;
let deployment = deployment.into_inner().unwrap_or_else(|e| e.into_inner());
let (deployment, upstream_started_at) = deployment
.into_inner()
.unwrap_or_else(|e| e.into_inner())
.unzip();
let mut response = match result {
Ok(response) => response,
Err(error) => {
@ -190,6 +210,7 @@ impl Tracking {
usage: None,
cost: None,
time_to_first_byte: None,
upstream_started_at,
ended_at: SystemTime::now(),
interrupted: false,
};
@ -203,6 +224,7 @@ impl Tracking {
headers.insert(name.clone(), value.clone());
}
}
let reported = response.extensions().get::<ReportedUsage>().cloned();
let streaming = response
.headers()
.get(CONTENT_TYPE)
@ -223,8 +245,10 @@ impl Tracking {
route,
admission: Some(admission),
deployment,
upstream_started_at,
status: response.status().as_u16(),
meter,
reported,
started,
first_byte: None,
ended: false,
@ -262,8 +286,11 @@ struct Finish {
route: Route,
admission: Option<Admission>,
deployment: Option<DeploymentFacts>,
upstream_started_at: Option<SystemTime>,
status: u16,
meter: Meter,
/// The usage a translated chat stream ended with, sent to the client or not.
reported: Option<ReportedUsage>,
body: Option<Vec<Bytes>>,
started: Instant,
first_byte: Option<Duration>,
@ -284,12 +311,32 @@ impl Finish {
}
}
fn usage(&self) -> Option<Usage> {
match &self.meter {
Meter::Off => None,
Meter::Stream(meter) => meter.usage(),
Meter::Body(chunks) => body_usage(self.route, &chunks.concat()),
}
/// The usage the response reported, with its cost when the deployment is priced by the
/// model its response names, which only a chat completion carries.
fn usage_and_cost(&self) -> (Option<Usage>, Option<f64>) {
let by_model = self.pipeline.prices().prices_by_response_model()
&& self.route == Route::ChatCompletions;
let (usage, model) = match &self.meter {
Meter::Off => (None, None),
Meter::Stream(meter) => match self.reported.as_ref().and_then(reported_usage) {
Some(usage) => (Some(usage), None),
None => (meter.usage(), meter.model().map(str::to_owned)),
},
Meter::Body(chunks) => {
let body = chunks.concat();
let model = by_model.then(|| body_model(&body)).flatten();
(body_usage(self.route, &body), model)
}
};
let cost = by_model
.then_some(())
.zip(self.deployment.as_ref().zip(usage.as_ref()))
.and_then(|((), (deployment, usage))| {
self.pipeline
.prices()
.response_cost(&deployment.id, model.as_deref(), usage)
});
(usage, cost)
}
}
@ -298,12 +345,14 @@ impl Drop for Finish {
let Some(admission) = self.admission.take() else {
return;
};
let (usage, cost) = self.usage_and_cost();
let outcome = Outcome {
deployment: self.deployment.take(),
status: self.status,
usage: self.usage(),
cost: None,
usage,
cost,
time_to_first_byte: self.first_byte,
upstream_started_at: self.upstream_started_at,
ended_at: SystemTime::now(),
interrupted: !self.ended,
};
@ -312,6 +361,10 @@ impl Drop for Finish {
}
}
fn reported_usage(reported: &ReportedUsage) -> Option<Usage> {
Usage::from_openai(&serde_json::to_value(reported.get()?).ok()?)
}
/// The `usage` object of a complete JSON response. Only that object is parsed: it is found by
/// its key, which inside a string would be escaped.
fn body_usage(route: Route, body: &[u8]) -> Option<Usage> {
@ -327,6 +380,17 @@ fn body_usage(route: Route, body: &[u8]) -> Option<Usage> {
}
}
/// The `model` a complete chat completion names, found by its first key like [`body_usage`].
fn body_model(body: &[u8]) -> Option<String> {
const KEY: &[u8] = b"\"model\":";
let start = memchr::memmem::find(body, KEY)? + KEY.len();
serde_json::Deserializer::from_slice(&body[start..])
.into_iter::<String>()
.next()?
.ok()
.filter(|model| !model.is_empty())
}
struct Tapped {
inner: BodyDataStream,
finish: Finish,
@ -378,4 +442,16 @@ mod tests {
body_usage(Route::ChatCompletions, &serde_json::to_vec(&chat).unwrap()).unwrap();
assert_eq!((usage.prompt_tokens(), usage.output_tokens), (5, 6));
}
#[rstest]
#[case::top_level(
br#"{"id":"c","model":"gpt-4o-2024-11-20","choices":[]}"#,
Some("gpt-4o-2024-11-20")
)]
#[case::spaced(br#"{"model": "gpt-4o", "choices":[]}"#, Some("gpt-4o"))]
#[case::inside_content(br#"{"choices":[{"message":{"content":"\"model\":\"x\""}}]}"#, None)]
#[case::empty(br#"{"model":"","choices":[]}"#, None)]
fn body_model_reads_the_response_model(#[case] body: &[u8], #[case] expected: Option<&str>) {
assert_eq!(body_model(body).as_deref(), expected);
}
}

32
litellm-rust/crates/gateway-inference/src/lib.rs
View file

@ -25,7 +25,10 @@ use litellm_inference_messages::MessagesRoute;
use litellm_inference_ocr::OcrRoute;
use litellm_inference_responses::ResponsesRoute;
use litellm_inference_transcription::AudioTranscriptionRoute;
use litellm_llms::base_llm::ocr::{handler::OcrClient, settings::OcrSettings};
use litellm_llms::base_llm::{
model_info::ModelInfo,
ocr::{handler::OcrClient, settings::OcrSettings},
};
use litellm_secrets::source::SecretSource;
pub use error::Error;
@ -36,6 +39,8 @@ pub use request::{JsonBody, JsonObject, RequestId};
pub struct Gateway {
cache: Option<Arc<dyn litellm_cache_response::ResponseCacheService>>,
hooks: OnceLock<Hooks>,
model_info: OnceLock<Arc<dyn ModelInfo>>,
drop_params: bool,
allow_client_side_credentials: bool,
pub audio_transcription: AudioTranscriptionRoute,
pub chat_completions: ChatCompletionsRoute,
@ -67,6 +72,29 @@ impl Gateway {
self.hooks.get()
}
/// Gives providers whose param rules read model capabilities the model cost map. Only the
/// first call takes effect.
pub fn install_model_info(&self, model_info: Arc<dyn ModelInfo>) {
let _ = self.model_info.set(model_info);
}
pub(crate) fn model_info(&self) -> Option<Arc<dyn ModelInfo>> {
self.model_info.get().cloned()
}
/// `litellm_settings.drop_params`: drops params a provider does not support instead of
/// refusing the request.
pub fn with_drop_params(self, drop_params: bool) -> Self {
Self {
drop_params,
..self
}
}
pub(crate) fn drop_params(&self) -> bool {
self.drop_params
}
/// `general_settings.allow_client_side_credentials`: lets a request body name the
/// connection and credential params the proxy otherwise refuses.
pub fn with_client_side_credentials(self, allowed: bool) -> Self {
@ -98,6 +126,8 @@ impl Gateway {
Ok(Self {
cache: None,
hooks: OnceLock::new(),
model_info: OnceLock::new(),
drop_params: false,
allow_client_side_credentials: false,
audio_transcription: AudioTranscriptionRoute::new(
provider.clone(),

73
litellm-rust/crates/gateway-inference/src/messages.rs
View file

@ -82,6 +82,7 @@ async fn handle(
None => route,
};
let upstream_headers = &upstream_headers(headers, tracking.as_ref());
let cache_headers = crate::caching::CacheHeaders::default();
let mut body = Some(body);
let routed = gateway
@ -103,7 +104,7 @@ async fn handle(
async move {
crate::observe::record(name, attempt.deployment);
request::authorize(identity, name, attempt.deployment).await?;
let call = project(attempt.deployment, body?, headers)?;
let call = project(attempt.deployment, body?, upstream_headers)?;
let machine = route.machine(call, cache_options.policy);
let stream = Sse::<Messages, _, _>::new(Json, |error| {
Bytes::from(Error::from(error).sse_frame())
@ -123,10 +124,61 @@ async fn handle(
}
}
/// Client headers the provider gets: the API headers scoped to the providers that take them,
/// and any header a hook set.
#[derive(Default)]
struct UpstreamHeaders {
scoped: Option<ProviderSpecificHeaders>,
extra: Option<Map<String, Value>>,
}
fn upstream_headers(
headers: &HeaderMap,
tracking: Option<&crate::hooks::Tracking>,
) -> UpstreamHeaders {
let client = |name: &str| Some(headers.get(name)?.to_str().ok()?.to_owned());
let Some(tracking) = tracking else {
return UpstreamHeaders {
scoped: scoped_headers(client),
extra: None,
};
};
let mut names: Vec<&str> = header_scopes()
.iter()
.flat_map(|(_, names)| names.iter().copied())
.collect();
names.sort_unstable();
names.dedup();
let forwarded = names
.iter()
.filter_map(|name| Some(((*name).to_owned(), client(name)?)))
.collect();
let Some(sent) = tracking.edit_headers(forwarded) else {
return UpstreamHeaders {
scoped: scoped_headers(client),
extra: None,
};
};
let value = |name: &str| {
sent.iter()
.find(|(sent, _)| sent.eq_ignore_ascii_case(name))
.map(|(_, value)| value.clone())
};
let extra: Map<String, Value> = sent
.iter()
.filter(|(name, _)| !names.iter().any(|scoped| scoped.eq_ignore_ascii_case(name)))
.map(|(name, value)| (name.clone(), Value::from(value.as_str())))
.collect();
UpstreamHeaders {
scoped: scoped_headers(value),
extra: (!extra.is_empty()).then_some(extra),
}
}
fn project(
deployment: &Deployment,
body: Map<String, Value>,
headers: &HeaderMap,
headers: &UpstreamHeaders,
) -> Result<MessagesCall, Error> {
let body = body
.into_iter()
@ -141,28 +193,31 @@ fn project(
api_base: deployment.api_base.clone(),
custom_llm_provider: deployment.custom_llm_provider.clone(),
litellm_params: deployment.litellm_params.clone(),
extra_headers: None,
provider_specific_header: anthropic_api_headers(headers),
extra_headers: headers.extra.clone(),
provider_specific_header: headers.scoped.clone(),
timeout: deployment.timeout,
shaping: deployment.shaping.clone(),
})
}
fn anthropic_api_headers(headers: &HeaderMap) -> Option<ProviderSpecificHeaders> {
let scopes = [
fn header_scopes() -> [(&'static str, &'static [&'static str]); 2] {
[
(
ANTHROPIC_API_HEADER_PROVIDERS,
ANTHROPIC_API_HEADERS.as_slice(),
),
("openrouter", OPENROUTER_API_HEADERS.as_slice()),
];
let forwarded: Vec<ProviderSpecificHeader> = scopes
]
}
fn scoped_headers(get: impl Fn(&str) -> Option<String>) -> Option<ProviderSpecificHeaders> {
let forwarded: Vec<ProviderSpecificHeader> = header_scopes()
.into_iter()
.filter_map(|(provider, names)| {
let extra_headers: Map<String, Value> = names
.iter()
.filter_map(|name| {
let value = headers.get(*name)?.to_str().ok()?;
let value = get(name)?;
Some(((*name).to_owned(), Value::from(value)))
})
.collect();

1
litellm-rust/crates/gateway-inference/tests/routes.rs
View file

@ -68,7 +68,6 @@ async fn chat_aliases_call_core_and_use_the_body_model_before_the_path(
}
#[rstest]
#[case::streaming(json!({"messages": [{"role": "user", "content": "hi"}], "stream": true}), 501)]
#[case::missing_messages(json!({}), 400)]
#[case::malformed_messages(json!({"messages": "hi"}), 400)]
#[case::invalid_streaming_request(json!({"messages": [], "stream": true}), 400)]

10
litellm-rust/crates/gateway/Cargo.toml
View file

@ -12,7 +12,7 @@ http-body-util = "0.1"
litellm-inference.workspace = true
litellm-gateway-inference.workspace = true
litellm-gateway-mcp.workspace = true
tokio-util = "0.7"
tokio-util = { version = "0.7", features = ["codec"] }
thiserror.workspace = true
futures-util.workspace = true
http.workspace = true
@ -21,7 +21,11 @@ litellm-gateway-auth.workspace = true
litellm-gateway-hooks.workspace = true
litellm-model-catalog.workspace = true
litellm-gateway-ui.workspace = true
litellm-auth-aws.workspace = true
litellm-auth-types.workspace = true
base64.workspace = true
md-5 = "0.11.0"
sha2.workspace = true
litellm-config.workspace = true
litellm-db.workspace = true
litellm-http = { workspace = true, features = ["mcp"] }
@ -39,7 +43,9 @@ uuid.workspace = true
tower-sessions-moka-store = "0.15.0"
[dev-dependencies]
base64.workspace = true
wiremock.workspace = true
litellm-framer.workspace = true
bytes.workspace = true
litellm-db-testing.workspace = true
sqlx.workspace = true
rstest.workspace = true

4
litellm-rust/crates/gateway/src/error.rs
View file

@ -20,6 +20,10 @@ pub enum Error {
DatabaseUrl(String),
#[error("reading a gateway setting from the secret source failed")]
Secret(#[from] litellm_secrets::Error),
#[error("litellm_settings.hooks: {0}")]
Hooks(String),
#[error("litellm_settings.s3_callback_params: {0}")]
S3Log(String),
#[error("coordination Redis: {0}")]
Redis(String),
#[error("MCP requires a configured master key")]

98
litellm-rust/crates/gateway/src/hooks.rs
View file

@ -1,13 +1,14 @@
//! Builds the hooks and observers inference requests run through, and the per-deployment prices
//! observers that read cost need.
use std::sync::Arc;
use std::{collections::HashSet, sync::Arc};
use litellm_config::Config;
use litellm_config::{Config, Model};
use litellm_gateway_hooks::{
DEFAULT_QUEUE, Facts, Hook, Observer, Observers, Pipeline, Prices, deployment_fields,
};
use litellm_gateway_inference::{Gateway, Hooks};
use litellm_llms::base_llm::model_info::ModelInfo;
use litellm_model_catalog::{Catalog, Provenance};
use litellm_secrets::source::SecretSource;
use serde_json::{Map, Value};
@ -29,18 +30,26 @@ pub async fn install(
hooks: Vec<Arc<dyn Hook>>,
observers: Vec<Arc<dyn Observer>>,
) -> Result<(), Error> {
if hooks.is_empty() && observers.is_empty() {
let serves_azure = config.model_list.iter().any(is_azure);
if hooks.is_empty() && observers.is_empty() && !serves_azure {
return Ok(());
}
let needs_cost = observers
.iter()
.any(|observer| observer.reads().contains(Facts::COST))
|| hooks.iter().any(|hook| hook.reads().contains(Facts::COST));
let catalog = match needs_cost || serves_azure {
true => load_catalog(inference).await.map(Arc::new),
false => None,
};
if let Some(catalog) = catalog.as_ref().filter(|_| serves_azure) {
inference.install_model_info(Arc::new(CostMap::new(catalog.clone())));
}
if hooks.is_empty() && observers.is_empty() {
return Ok(());
}
let prices = match needs_cost {
true => {
let catalog = load_catalog(inference).await;
prices(config, inference, catalog.as_ref())
}
true => prices(config, inference, catalog.as_deref()),
false => Prices::default(),
};
let names: Vec<_> = hooks
@ -135,7 +144,9 @@ async fn fetch(inference: &Gateway, url: &str) -> Result<Vec<u8>, Box<dyn std::e
}
/// Every deployment's prices: the cost map entry for its model, overlaid by the prices its
/// `model_info` and `litellm_params` set, keyed by the router's deployment ids.
/// `model_info` and `litellm_params` set, keyed by the router's deployment ids. An Azure
/// deployment that names neither a `base_model` nor prices of its own is priced by the model
/// its response names, as the proxy prices it.
fn prices(config: &Config, inference: &Gateway, catalog: Option<&Catalog>) -> Prices {
let mut prices = Prices::default();
for (model, id) in config
@ -151,18 +162,75 @@ fn prices(config: &Config, inference: &Gateway, catalog: Option<&Catalog>) -> Pr
.iter()
.map(|(name, value)| (name, value.clone())),
);
let base_model = model_info.get("base_model").and_then(Value::as_str);
let fields = deployment_fields(
catalog,
&model.litellm_params.model,
base_model,
&[&model_info, &params],
);
let base_model = params
.get("base_model")
.or_else(|| model_info.get("base_model"))
.and_then(Value::as_str)
.filter(|name| !name.is_empty());
let name = match is_azure(model) && !model.litellm_params.model.starts_with(AZURE_ROUTE) {
true => format!("{AZURE_ROUTE}{}", model.litellm_params.model),
false => model.litellm_params.model.clone(),
};
let overrides = [&model_info, &params];
if is_azure(model)
&& base_model.is_none()
&& deployment_fields(None, &name, None, &overrides).is_empty()
{
prices.price_by_response_model(id.into(), AZURE_PROVIDER);
}
let fields = deployment_fields(catalog, &name, base_model, &overrides);
prices.insert(id.into(), &fields);
}
if let Some(catalog) = catalog.filter(|_| prices.prices_by_response_model()) {
for name in catalog
.model_names()
.filter(|name| name.starts_with(AZURE_ROUTE))
{
if let Some(found) = catalog.lookup(name) {
prices.insert_model(name, found.entry.fields());
}
}
}
prices
}
const AZURE_PROVIDER: &str = "azure";
const AZURE_ROUTE: &str = "azure/";
fn is_azure(model: &Model) -> bool {
match model.litellm_params.custom_llm_provider.as_deref() {
Some(provider) => provider == AZURE_PROVIDER,
None => model.litellm_params.model.starts_with(AZURE_ROUTE),
}
}
/// The model cost map as the capabilities Azure's param rules read.
struct CostMap {
catalog: Arc<Catalog>,
keys: HashSet<String>,
}
impl CostMap {
fn new(catalog: Arc<Catalog>) -> Self {
let keys = catalog
.model_names()
.filter_map(|name| catalog.lookup(name))
.flat_map(|found| found.entry.fields().keys().cloned().collect::<Vec<_>>())
.collect();
Self { catalog, keys }
}
}
impl ModelInfo for CostMap {
fn entry(&self, model: &str) -> Option<&Map<String, Value>> {
self.catalog.lookup(model).map(|found| found.entry.fields())
}
fn declares(&self, key: &str) -> bool {
self.keys.contains(key)
}
}
fn json_object<'a, V: serde::Serialize + 'a>(
entries: impl Iterator<Item = (&'a String, V)>,
) -> Map<String, Value> {

13
litellm-rust/crates/gateway/src/lib.rs
View file

@ -9,7 +9,11 @@ pub mod hooks;
mod key_status;
mod keys;
pub mod metrics;
pub mod model_capacity;
pub mod prometheus;
pub mod rate_limits;
pub mod request_hooks;
pub mod s3_log;
mod secrets;
pub mod spend;
pub use error::Error;
@ -75,7 +79,14 @@ pub fn build_inference(config: &Config) -> Result<Arc<Gateway>, Error> {
let resources = CoreResources::new(pool);
Ok(Arc::new(
Gateway::new(resources, http, secrets, ModelRouter::from_config(config)?)?
.with_client_side_credentials(allows_client_side_credentials(config)),
.with_client_side_credentials(allows_client_side_credentials(config))
.with_drop_params(
config
.litellm_settings
.drop_params
.as_ref()
.is_some_and(|flag| flag.is_set()),
),
))
}

16
litellm-rust/crates/gateway/src/main.rs
View file

@ -128,12 +128,25 @@ async fn main() -> Result<(), Box<dyn Error>> {
),
));
}
if let Some(capacity) =
litellm_gateway::model_capacity::ModelCapacity::for_gateway(&config, coordination.clone())
{
hooks.push(Arc::new(capacity));
}
hooks.extend(litellm_gateway::request_hooks::from_config(&config)?);
let mut observers: Vec<Arc<dyn litellm_gateway_hooks::Observer>> = Vec::new();
let spend = database.zip(litellm_gateway::spend::SpendSettings::from_config(&config));
let spend_tracking = spend.is_some();
if let Some((pool, settings)) = spend {
observers.push(litellm_gateway::spend::SpendWriter::start(pool, settings));
}
let prometheus = litellm_gateway::prometheus::Prometheus::from_config(&config)?.map(Arc::new);
if let Some(prometheus) = &prometheus {
observers.push(prometheus.clone());
}
if let Some(s3_log) = litellm_gateway::s3_log::S3Log::from_config(&config, &inference).await? {
observers.push(s3_log);
}
litellm_gateway::hooks::install(&inference, &config, hooks, observers).await?;
let listener = tokio::net::TcpListener::bind((settings.host.as_str(), settings.port)).await?;
let address = listener.local_addr()?;
@ -148,7 +161,8 @@ async fn main() -> Result<(), Box<dyn Error>> {
tracing::info!(%address, models = config.model_list.len(), virtual_keys, shared_state, spend_tracking, log_level = %level, "gateway listening");
let draining = litellm_gateway::health::Draining::default();
let metrics = litellm_gateway::metrics::Metrics::new(&config, draining.flag());
let metrics =
litellm_gateway::metrics::Metrics::new(&config, draining.flag()).serve(prometheus);
let metrics_auth = litellm_gateway::metrics_auth(&config, inference.secrets.clone());
let app = litellm_gateway::with_ops(
litellm_gateway::router_with_key_cache(inference.clone(), &config, ui, mcp_router, keys),

94
litellm-rust/crates/gateway/src/metrics.rs
View file

@ -9,7 +9,7 @@ use std::{
collections::HashMap,
fmt::Write as _,
sync::{
Arc, RwLock,
Arc, OnceLock, RwLock,
atomic::{AtomicBool, AtomicI64, AtomicU64, Ordering},
},
time::{Duration, Instant},
@ -26,6 +26,8 @@ use http_body_util::BodyExt;
use litellm_config::Config;
use litellm_gateway_inference::observe;
use crate::prometheus::Prometheus;
/// Upper bounds, in seconds, shared by the latency and time to first byte histograms. They reach
/// 30 minutes because streams that long are normal here.
const BUCKETS: [f64; 19] = [
@ -123,6 +125,8 @@ struct Inner {
failures: RwLock<HashMap<FailureKey, AtomicU64>>,
in_flight: AtomicI64,
draining: Arc<AtomicBool>,
/// The `prometheus` callback's metrics, served after these.
prometheus: OnceLock<Arc<Prometheus>>,
}
impl Metrics {
@ -164,9 +168,19 @@ impl Metrics {
failures: RwLock::default(),
in_flight: AtomicI64::new(0),
draining,
prometheus: OnceLock::new(),
}))
}
/// Serves `prometheus`'s metrics too. The ones it shares a name with are then served from it
/// alone, with the proxy's labels.
pub fn serve(self, prometheus: Option<Arc<Prometheus>>) -> Self {
if let Some(prometheus) = prometheus {
let _ = self.0.prometheus.set(prometheus);
}
self
}
fn series(&self, route: &str, key: SeriesKey) -> Arc<Series> {
if let Some(series) = self
.0
@ -249,21 +263,24 @@ impl Metrics {
})
.collect();
out.push_str("# HELP litellm_proxy_total_requests_metric Total requests the proxy answered, by route, requested model and status code\n");
out.push_str("# TYPE litellm_proxy_total_requests_metric counter\n");
for ((_, _, series), labels) in series.iter().zip(&labels) {
let _ = writeln!(
out,
"litellm_proxy_total_requests_metric_total{{{labels}}} {}",
series.latency.count.load(Ordering::Relaxed)
);
}
out.push_str("# HELP litellm_request_total_latency_metric Seconds from request start to the end of the response body\n");
out.push_str("# TYPE litellm_request_total_latency_metric histogram\n");
for ((_, _, series), labels) in series.iter().zip(&labels) {
series
.latency
.render(&mut out, "litellm_request_total_latency_metric", labels);
let replaced = inner.prometheus.get().is_some();
if !replaced {
out.push_str("# HELP litellm_proxy_total_requests_metric Total requests the proxy answered, by route, requested model and status code\n");
out.push_str("# TYPE litellm_proxy_total_requests_metric counter\n");
for ((_, _, series), labels) in series.iter().zip(&labels) {
let _ = writeln!(
out,
"litellm_proxy_total_requests_metric_total{{{labels}}} {}",
series.latency.count.load(Ordering::Relaxed)
);
}
out.push_str("# HELP litellm_request_total_latency_metric Seconds from request start to the end of the response body\n");
out.push_str("# TYPE litellm_request_total_latency_metric histogram\n");
for ((_, _, series), labels) in series.iter().zip(&labels) {
series
.latency
.render(&mut out, "litellm_request_total_latency_metric", labels);
}
}
out.push_str("# HELP litellm_request_time_to_first_byte_metric Seconds from request start to the first response body byte\n");
out.push_str("# TYPE litellm_request_time_to_first_byte_metric histogram\n");
@ -275,25 +292,27 @@ impl Metrics {
);
}
out.push_str("# HELP litellm_deployment_failure_responses Upstream calls that failed with 429 or 5xx, by model group and deployment api_base\n");
out.push_str("# TYPE litellm_deployment_failure_responses counter\n");
let mut failures: Vec<(FailureKey, u64)> = inner
.failures
.read()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.iter()
.map(|(key, count)| (*key, count.load(Ordering::Relaxed)))
.collect();
failures.sort_by_key(|(key, _)| *key);
for (key, count) in failures {
let deployment = &inner.deployments[key.deployment as usize];
let _ = writeln!(
out,
"litellm_deployment_failure_responses_total{{requested_model=\"{}\",api_base=\"{}\",exception_status=\"{}\"}} {count}",
escape(&inner.model_labels[deployment.model as usize]),
escape(&deployment.api_base),
key.status
);
if !replaced {
out.push_str("# HELP litellm_deployment_failure_responses Upstream calls that failed with 429 or 5xx, by model group and deployment api_base\n");
out.push_str("# TYPE litellm_deployment_failure_responses counter\n");
let mut failures: Vec<(FailureKey, u64)> = inner
.failures
.read()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.iter()
.map(|(key, count)| (*key, count.load(Ordering::Relaxed)))
.collect();
failures.sort_by_key(|(key, _)| *key);
for (key, count) in failures {
let deployment = &inner.deployments[key.deployment as usize];
let _ = writeln!(
out,
"litellm_deployment_failure_responses_total{{requested_model=\"{}\",api_base=\"{}\",exception_status=\"{}\"}} {count}",
escape(&inner.model_labels[deployment.model as usize]),
escape(&deployment.api_base),
key.status
);
}
}
out.push_str("# HELP litellm_in_flight_requests Requests whose response has not finished, streams included\n");
@ -310,6 +329,9 @@ impl Metrics {
"litellm_gateway_draining {}",
u8::from(inner.draining.load(Ordering::Relaxed))
);
if let Some(prometheus) = inner.prometheus.get() {
prometheus.render(&mut out);
}
out
}
}
@ -327,7 +349,7 @@ fn api_base_label(api_base: &str) -> String {
}
}
fn escape(value: &str) -> std::borrow::Cow<'_, str> {
pub(crate) fn escape(value: &str) -> std::borrow::Cow<'_, str> {
if value.contains(['\\', '"', '\n']) {
value
.replace('\\', "\\\\")

654
litellm-rust/crates/gateway/src/model_capacity.rs Normal file
View file

@ -0,0 +1,654 @@
//! The model capacity hook, the proxy's `dynamic_rate_limiter_v3`: caps each model group at the
//! sum of its deployments' `rpm` and `tpm` across every replica, and with `priority_reservation`
//! splits that capacity between priorities once the group is busy.
//!
//! Each request adds one to `{model_saturation_check:<group>}:requests` and checks the group's
//! tokens, all or nothing, with the proxy's keys and scripts. With priorities, the request's pool
//! `{priority_model:<group>:<priority>}` (or `:default_pool`) is held to its share of capacity
//! once the group's saturation reaches `saturation_threshold`, and only counted below it. Tokens
//! are added to both counters when the response ends.
use std::{
collections::HashMap,
sync::{Arc, Mutex},
time::{Duration, Instant},
};
use futures_util::future::BoxFuture;
use http::{HeaderName, HeaderValue};
use litellm_config::Config;
use litellm_gateway_hooks::{
Facts, Hook, KeyFacts, Record, Refusal, RequestEdit, RequestFacts, Settle, Targets,
};
use serde_json::Value;
use crate::{
coordination::Coordination,
rate_limits::{Kind, Pair, RateLimitSettings, Reserved, Store, Unavailable},
};
/// The names a config's `callbacks` list turns this hook on with.
const NAMES: [&str; 2] = ["dynamic_rate_limiter_v3", "model_capacity"];
/// A priority's share of a group's capacity, as `priority_reservation` writes it.
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum Share {
Percent(f64),
Rpm(f64),
Tpm(f64),
}
impl Share {
/// The share as a fraction of the group, as the proxy's `convert_priority_to_percent`.
fn fraction(self, capacity: &Capacity) -> f64 {
match self {
Self::Percent(value) => value,
Self::Rpm(value) => match capacity.rpm {
Some(rpm) if rpm > 0 => value / rpm as f64,
_ => value,
},
Self::Tpm(value) => match capacity.tpm {
Some(tpm) if tpm > 0 => value / tpm as f64,
_ => value,
},
}
}
}
/// Which tokens count: `general_settings.token_rate_limit_type`.
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub enum TokenCount {
#[default]
Total,
Input,
Output,
}
#[derive(Clone, Debug, PartialEq)]
pub struct CapacitySettings {
pub window: u64,
pub fail_closed: bool,
/// `litellm_settings.priority_reservation`, in config order.
pub reservation: Option<Vec<(String, Share)>>,
/// `priority_reservation_settings.default_priority`: the share of keys without a listed
/// priority, together.
pub default_priority: f64,
/// `priority_reservation_settings.saturation_threshold`.
pub saturation_threshold: f64,
/// `priority_reservation_settings.saturation_check_cache_ttl`: how long a group's
/// saturation is reused before it is read again.
pub saturation_ttl: Duration,
pub tokens: TokenCount,
}
impl Default for CapacitySettings {
fn default() -> Self {
let rate_limits = RateLimitSettings::default();
Self {
window: rate_limits.window,
fail_closed: rate_limits.fail_closed,
reservation: None,
default_priority: 0.25,
saturation_threshold: 0.5,
saturation_ttl: Duration::from_secs(60),
tokens: TokenCount::Total,
}
}
}
impl CapacitySettings {
/// `None` when the config's `callbacks` do not list the hook.
pub fn from_config(config: &Config) -> Option<Self> {
let litellm = &config.litellm_settings;
let listed = litellm.callbacks.iter().flatten().any(|callback| {
callback
.as_str()
.is_some_and(|name| NAMES.contains(&name.trim()))
});
if !listed {
return None;
}
let rate_limits = RateLimitSettings::from_config(config);
let number = |value: &litellm_config::Value| match value {
litellm_config::Value::Number(number) => number.as_f64(),
litellm_config::Value::String(text) => text.trim().parse().ok(),
_ => None,
};
let reservation = match litellm.additional_fields.get("priority_reservation") {
Some(litellm_config::Value::Mapping(mapping)) => Some(
mapping
.iter()
.filter_map(|(priority, share)| {
let priority = match priority {
litellm_config::Value::String(text) => text.clone(),
litellm_config::Value::Number(number) => number.to_string(),
_ => return None,
};
let share = match share {
litellm_config::Value::Mapping(share) => {
let value = share.get("value").and_then(number).unwrap_or(1.0);
match share.get("type").and_then(litellm_config::Value::as_str) {
Some("rpm") => Share::Rpm(value),
Some("tpm") => Share::Tpm(value),
_ => Share::Percent(value),
}
}
other => Share::Percent(number(other)?),
};
Some((priority, share))
})
.collect(),
),
_ => None,
};
let settings = litellm
.additional_fields
.get("priority_reservation_settings");
let setting = |name: &str| {
settings
.and_then(|settings| settings.get(name))
.and_then(number)
};
let defaults = Self::default();
Some(Self {
window: rate_limits.window,
fail_closed: rate_limits.fail_closed,
reservation,
default_priority: setting("default_priority").unwrap_or(defaults.default_priority),
saturation_threshold: setting("saturation_threshold")
.unwrap_or(defaults.saturation_threshold),
saturation_ttl: setting("saturation_check_cache_ttl")
.filter(|seconds| *seconds >= 0.0)
.map_or(defaults.saturation_ttl, Duration::from_secs_f64),
tokens: match config
.general_settings
.additional_fields
.get("token_rate_limit_type")
.and_then(litellm_config::Value::as_str)
{
Some("input") => TokenCount::Input,
Some("output") => TokenCount::Output,
_ => TokenCount::Total,
},
})
}
}
/// A model group's requests and tokens per window: its deployments' `rpm` and `tpm`, summed.
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub struct Capacity {
pub rpm: Option<i64>,
pub tpm: Option<i64>,
}
/// Each configured group's capacity, from `rpm`/`tpm` on the deployment, its `litellm_params`
/// or its `model_info`, in that order, as the proxy's router sums them.
pub fn capacities(config: &Config) -> HashMap<String, Capacity> {
let mut groups: HashMap<String, Capacity> = HashMap::new();
for model in &config.model_list {
let configured = |field: &str| -> Option<f64> {
let from_value = |value: &litellm_config::Value| match value {
litellm_config::Value::Number(number) => number.as_f64(),
litellm_config::Value::String(text) => text.trim().parse().ok(),
_ => None,
};
let spelled = |value: &litellm_config::Spelled<f64>| match value {
litellm_config::Spelled::Value(value) => Some(*value),
litellm_config::Spelled::Text(text) => text.trim().parse().ok(),
};
let params = match field {
"tpm" => model.litellm_params.tpm.as_ref(),
_ => model.litellm_params.rpm.as_ref(),
};
model
.additional_fields
.get(field)
.and_then(from_value)
.or_else(|| params.and_then(spelled))
.or_else(|| model.model_info.get(field).and_then(from_value))
};
let group = groups.entry(model.model_name.clone()).or_default();
let add = |total: &mut Option<i64>, value: Option<f64>| {
if let Some(value) = value {
*total = Some(total.unwrap_or(0) + value as i64);
}
};
add(&mut group.tpm, configured("tpm"));
add(&mut group.rpm, configured("rpm"));
}
groups
}
/// The pool a request draws on and its share of the group.
#[derive(Clone, Debug, PartialEq)]
struct Pool {
name: String,
share: f64,
}
/// The `model_capacity` hook.
pub struct ModelCapacity {
settings: CapacitySettings,
groups: HashMap<String, Capacity>,
store: Arc<Store>,
saturation: Mutex<HashMap<String, (Option<i64>, Instant)>>,
}
impl ModelCapacity {
pub fn new(
settings: CapacitySettings,
groups: HashMap<String, Capacity>,
redis: Option<Arc<Coordination>>,
) -> Self {
Self {
store: Arc::new(Store::new(redis, settings.window, settings.fail_closed)),
settings,
groups,
saturation: Mutex::default(),
}
}
/// The hook for a gateway whose config lists it, else `None`.
pub fn for_gateway(config: &Config, redis: Option<Arc<Coordination>>) -> Option<Self> {
Some(Self::new(
CapacitySettings::from_config(config)?,
capacities(config),
redis,
))
}
/// The pool the proxy's `_get_priority_allocation` picks, with its normalized share.
fn pool(&self, group: &str, capacity: &Capacity, priority: Option<&str>) -> Pool {
let explicit = self.settings.reservation.as_ref().and_then(|reservation| {
let priority = priority?;
reservation
.iter()
.any(|(name, _)| name == priority)
.then(|| {
let fractions: Vec<(&str, f64)> = reservation
.iter()
.map(|(name, share)| (name.as_str(), share.fraction(capacity)))
.collect();
let total: f64 = fractions.iter().map(|(_, fraction)| fraction).sum();
let fraction = fractions
.iter()
.find(|(name, _)| *name == priority)
.map_or(0.0, |(_, fraction)| *fraction);
(
priority,
if total > 1.0 {
fraction / total
} else {
fraction
},
)
})
});
match explicit {
Some((priority, share)) => Pool {
name: format!("{group}:{priority}"),
share,
},
None => Pool {
name: format!("{group}:default_pool"),
share: self.settings.default_priority,
},
}
}
/// A counter's value, reused for `saturation_check_cache_ttl`.
async fn cached(&self, prefix: &str, kind: Kind) -> Option<i64> {
let key = format!("{prefix}:{}", kind.name());
if let Some((value, read_at)) = self
.saturation
.lock()
.unwrap_or_else(|e| e.into_inner())
.get(&key)
&& read_at.elapsed() < self.settings.saturation_ttl
{
return *value;
}
let value = self.store.read(prefix, kind).await;
let mut cache = self.saturation.lock().unwrap_or_else(|e| e.into_inner());
if cache.len() >= 10_000 {
cache.clear();
}
cache.insert(key, (value, Instant::now()));
value
}
/// How full the group is: its busiest of requests and tokens against capacity.
async fn saturation(&self, prefix: &str, capacity: &Capacity) -> f64 {
let mut saturation = 0.0_f64;
for (kind, limit) in [(Kind::Requests, capacity.rpm), (Kind::Tokens, capacity.tpm)] {
if let Some(limit) = limit.filter(|limit| *limit > 0)
&& let Some(count) = self.cached(prefix, kind).await
{
saturation = saturation.max(count as f64 / limit as f64);
}
}
saturation
}
fn unavailable(&self, request: &RequestFacts) -> Refusal {
crate::rate_limits::unavailable(request.route)
}
}
/// `team_metadata.priority`, else the key's `metadata.priority`.
fn priority(key: Option<&KeyFacts>) -> Option<String> {
let key = key?;
let read = |value: &Value| match value {
Value::String(text) => Some(text.clone()),
Value::Number(number) => Some(number.to_string()),
_ => None,
};
key.team_metadata
.get("priority")
.and_then(read)
.or_else(|| key.metadata.get("priority").and_then(read))
}
fn limit_text(limit: Option<i64>) -> String {
limit.map_or_else(|| "not configured".to_owned(), |limit| limit.to_string())
}
/// `{:.<places>%}` as Python formats a fraction.
fn percent(fraction: f64, places: usize) -> String {
format!("{:.*}%", places, fraction * 100.0)
}
impl Hook for ModelCapacity {
fn name(&self) -> &'static str {
"model_capacity"
}
fn reads(&self) -> Facts {
Facts::KEY | Facts::USAGE
}
fn writes(&self) -> Targets {
Targets::ADMISSION | Targets::RESPONSE_HEADERS
}
fn run<'a>(
&'a self,
request: &'a Arc<RequestFacts>,
edit: &'a mut RequestEdit,
) -> BoxFuture<'a, Result<Option<Box<dyn Settle>>, Refusal>> {
Box::pin(async move {
let priority = priority(request.key.as_deref());
let priority_header = priority.as_deref().unwrap_or("default");
if !request.stream {
if priority_header.chars().all(|char| (char as u32) < 256)
&& let Ok(value) = HeaderValue::from_str(priority_header)
{
edit.add_response_header(HeaderName::from_static("x-litellm-priority"), value);
}
edit.add_response_header(
HeaderName::from_static("x-litellm-rate-limiter-version"),
HeaderValue::from_static("v3"),
);
}
let group = &*request.requested_model;
let Some(capacity) = self.groups.get(group).copied() else {
return Ok(None);
};
if capacity.rpm.is_none() && capacity.tpm.is_none() {
return Ok(None);
}
let model_prefix = format!("{{model_saturation_check:{group}}}");
let saturation = self.saturation(&model_prefix, &capacity).await;
let pool = self.pool(group, &capacity, priority.as_deref());
let pool_prefix = format!("{{priority_model:{}}}", pool.name);
// The group counts every request against its whole capacity; a zero limit is off.
let mut pairs = Vec::new();
let push = |pairs: &mut Vec<(&'static str, Pair)>,
key: &'static str,
prefix: &str,
limits: [Option<i64>; 2]| {
for (kind, limit, increment) in
[(Kind::Requests, limits[0], 1), (Kind::Tokens, limits[1], 0)]
{
if let Some(limit) = limit {
pairs.push((
key,
Pair {
prefix: prefix.to_owned(),
kind,
limit,
increment,
},
));
}
}
};
push(
&mut pairs,
"model_saturation_check",
&model_prefix,
[
capacity.rpm.filter(|rpm| *rpm != 0),
capacity.tpm.filter(|tpm| *tpm != 0),
],
);
let share = |limit: Option<i64>| limit.map(|limit| (limit as f64 * pool.share) as i64);
let pool_limits = [share(capacity.rpm), share(capacity.tpm)];
let reserving =
self.settings.reservation.is_some() && pool_limits.iter().any(Option::is_some);
let enforced = reserving && saturation >= self.settings.saturation_threshold;
if enforced {
push(&mut pairs, "priority_model", &pool_prefix, pool_limits);
}
let checks: Vec<Pair> = pairs.iter().map(|(_, pair)| pair.clone()).collect();
let (reserved, shared) = self
.store
.check_and_add(&checks)
.await
.map_err(|Unavailable| self.unavailable(request))?;
let counts = match reserved {
Reserved::Counts(counts) => counts,
Reserved::Over { index, current } => {
let (key, pair) = &pairs[index];
let remaining = pair.limit - current;
let kind = pair.kind.name();
let priority_text = priority.as_deref().unwrap_or("None");
let (message, saturation_header) = match *key {
"model_saturation_check" => (
format!(
"Model capacity reached for {group}. Priority: {priority_text}, Rate limit type: {kind}, Model TPM: {}, Model RPM: {}, Remaining: {remaining}",
limit_text(capacity.tpm),
limit_text(capacity.rpm),
),
None,
),
_ => (
format!(
"Priority-based rate limit exceeded. Model: {group}, Priority: {priority_text}, Rate limit type: {kind}, Model TPM: {}, Model RPM: {}, Remaining: {remaining}, Model saturation: {}",
limit_text(capacity.tpm),
limit_text(capacity.rpm),
percent(saturation, 1),
),
Some(percent(saturation, 2)),
),
};
let mut headers = vec![
("retry-after", self.settings.window.to_string()),
("rate_limit_type", kind.to_owned()),
("x-litellm-priority", priority_header.to_owned()),
];
if let Some(saturation) = saturation_header {
headers.push(("x-litellm-saturation", saturation));
}
return Err(crate::rate_limits::throttled(
request.route,
message,
headers,
));
}
};
let mut statuses: Vec<(&str, Kind, i64, i64)> = pairs
.iter()
.zip(&counts)
.map(|((key, pair), count)| (*key, pair.kind, pair.limit, pair.limit - count))
.collect();
// Below the threshold the pool's requests are only counted, so its share is known
// once it is enforced.
if reserving
&& !enforced
&& let Some(rpm) = pool_limits[0]
&& let Ok(counts) = self
.store
.requests(std::slice::from_ref(&pool_prefix))
.await
{
statuses.push(("priority_model", Kind::Requests, rpm, rpm - counts[0]));
}
if !request.stream {
for (key, kind, limit, remaining) in statuses {
let kind = kind.name();
for (name, value) in [
(format!("x-ratelimit-{key}-remaining-{kind}"), remaining),
(format!("x-ratelimit-{key}-limit-{kind}"), limit),
] {
if let Ok(name) = HeaderName::from_bytes(name.as_bytes()) {
edit.add_response_header(name, HeaderValue::from(value));
}
}
}
}
Ok(Some(Box::new(Charge {
store: self.store.clone(),
prefixes: [model_prefix, pool_prefix],
tokens: self.settings.tokens,
shared,
}) as Box<dyn Settle>))
})
}
}
/// Adds a successful response's tokens to the group and its pool.
struct Charge {
store: Arc<Store>,
prefixes: [String; 2],
tokens: TokenCount,
shared: bool,
}
impl Settle for Charge {
fn settle(self: Box<Self>, record: &Record) {
if !(200..300).contains(&record.outcome.status) {
return;
}
let Some(usage) = record.outcome.usage else {
return;
};
let tokens = match self.tokens {
TokenCount::Total => usage.total_tokens(),
TokenCount::Input => usage.prompt_tokens(),
TokenCount::Output => usage.output_tokens,
};
let Ok(tokens) = i64::try_from(tokens) else {
return;
};
if tokens == 0 {
return;
}
let Charge {
store,
prefixes,
shared,
..
} = *self;
tokio::spawn(async move {
let deltas = prefixes
.into_iter()
.map(|prefix| (prefix, tokens))
.collect();
store.charge(deltas, shared).await;
});
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
fn capacity_hook(reservation: Option<Vec<(String, Share)>>) -> ModelCapacity {
ModelCapacity::new(
CapacitySettings {
reservation,
..CapacitySettings::default()
},
HashMap::new(),
None,
)
}
#[rstest]
fn pools_follow_the_proxy_allocation() {
let capacity = Capacity {
rpm: Some(100),
tpm: Some(1_000),
};
let hook = capacity_hook(Some(vec![
("prod".into(), Share::Percent(0.6)),
("dev".into(), Share::Rpm(80.0)),
]));
// 0.6 + 0.8 is over-allocated, so both shrink to their proportion.
let prod = hook.pool("m", &capacity, Some("prod"));
assert_eq!(prod.name, "m:prod");
assert!((prod.share - 0.6 / 1.4).abs() < 1e-12);
assert_eq!(
hook.pool("m", &capacity, Some("other")),
Pool {
name: "m:default_pool".into(),
share: 0.25
}
);
assert_eq!(
capacity_hook(None).pool("m", &capacity, Some("prod")).name,
"m:default_pool"
);
}
#[rstest]
#[case(0.625, 1, "62.5%")]
#[case(0.5, 2, "50.00%")]
fn formats_percentages_like_python(
#[case] fraction: f64,
#[case] places: usize,
#[case] expected: &str,
) {
assert_eq!(percent(fraction, places), expected);
}
#[rstest]
fn capacities_sum_deployments() {
let config = Config::from_yaml(
"model_list:
- model_name: a
litellm_params: {model: openai/x, rpm: 10, tpm: 1000}
- model_name: a
rpm: 5
litellm_params: {model: openai/y, rpm: 99}
model_info: {tpm: 500}
- model_name: b
litellm_params: {model: openai/z}
",
)
.unwrap();
let groups = capacities(&config);
assert_eq!(
groups["a"],
Capacity {
rpm: Some(15),
tpm: Some(1_500)
}
);
assert_eq!(groups["b"], Capacity::default());
}
}

1253
litellm-rust/crates/gateway/src/prometheus.rs Normal file
View file

File diff suppressed because it is too large Load diff

394
litellm-rust/crates/gateway/src/rate_limits.rs
View file

@ -181,6 +181,102 @@ end
return reply
"#;
/// Counts a request on every `(window, requests)` pair, starting a window that is missing or
/// over, and returns each window's start and count. The caller compares; the count always rises.
const REQUESTS: &str = r#"
local results = {}
local now = tonumber(ARGV[1])
local window_size = tonumber(ARGV[2])
for i = 1, #KEYS, 2 do
local window_key = KEYS[i]
local counter_key = KEYS[i + 1]
local window_start = redis.call('GET', window_key)
if not window_start or (now - tonumber(window_start)) >= window_size then
local prefix = string.sub(window_key, 1, -(#':window') - 1)
redis.call('DEL', prefix .. ':requests', prefix .. ':tokens')
redis.call('SET', window_key, tostring(now))
redis.call('SET', counter_key, 1)
redis.call('EXPIRE', window_key, window_size)
redis.call('EXPIRE', counter_key, window_size)
table.insert(results, tostring(now))
table.insert(results, 1)
else
local counter = redis.call('INCR', counter_key)
if redis.call('TTL', counter_key) == -1 then
redis.call('EXPIRE', counter_key, window_size)
end
table.insert(results, window_start)
table.insert(results, counter)
end
end
return results
"#;
/// Adds the increment to every `(window, tokens)` pair, or to none when one would go over its
/// limit: `{0, count, start, ...}` or `{1, index, current, limit}`.
const RESERVE: &str = r#"
local now = tonumber(redis.call('TIME')[1])
local count = #KEYS / 2
local state = {}
for i = 1, count do
local window_key = KEYS[(i - 1) * 2 + 1]
local counter_key = KEYS[(i - 1) * 2 + 2]
local base = (i - 1) * 4 + 1
local limit = tonumber(ARGV[base])
local increment = tonumber(ARGV[base + 1])
local window_size = tonumber(ARGV[base + 3])
local window_start = redis.call('GET', window_key)
local expired = (not window_start) or ((now - tonumber(window_start)) >= window_size)
local current = 0
if not expired then
current = tonumber(redis.call('GET', counter_key) or 0)
end
local blocked
if increment > 0 then
blocked = current + increment > limit
else
blocked = current >= limit
end
if blocked then
return { 1, i, current, limit }
end
state[i] = { expired, window_start }
end
local results = { 0 }
local reset = {}
for i = 1, count do
local window_key = KEYS[(i - 1) * 2 + 1]
local counter_key = KEYS[(i - 1) * 2 + 2]
local base = (i - 1) * 4 + 1
local increment = tonumber(ARGV[base + 1])
local ttl = tonumber(ARGV[base + 2])
local window_size = tonumber(ARGV[base + 3])
if state[i][1] then
if not reset[window_key] then
local prefix = string.sub(window_key, 1, -(#':window') - 1)
redis.call('DEL', prefix .. ':requests', prefix .. ':tokens')
reset[window_key] = true
end
redis.call('SET', window_key, tostring(now))
redis.call('SET', counter_key, increment)
redis.call('EXPIRE', window_key, window_size)
if ttl > 0 then
redis.call('EXPIRE', counter_key, ttl)
end
table.insert(results, increment)
table.insert(results, now)
else
local counter = redis.call('INCRBY', counter_key, increment)
if redis.call('TTL', counter_key) == -1 and ttl > 0 then
redis.call('EXPIRE', counter_key, ttl)
end
table.insert(results, counter)
table.insert(results, tonumber(state[i][2]))
end
end
return results
"#;
/// Adds each token delta, giving a counter without an expiry the window's.
const CHARGE: &str = r#"
local results = {}
@ -197,6 +293,8 @@ return results
/// Sent by hash, loaded once per Redis.
static ADMIT_SCRIPT: LazyLock<redis::Script> = LazyLock::new(|| redis::Script::new(ADMIT));
static REQUESTS_SCRIPT: LazyLock<redis::Script> = LazyLock::new(|| redis::Script::new(REQUESTS));
static RESERVE_SCRIPT: LazyLock<redis::Script> = LazyLock::new(|| redis::Script::new(RESERVE));
static CHARGE_SCRIPT: LazyLock<redis::Script> = LazyLock::new(|| redis::Script::new(CHARGE));
#[derive(Clone, Debug, PartialEq)]
@ -261,14 +359,14 @@ impl Descriptor {
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum Kind {
pub(crate) enum Kind {
Requests,
Tokens,
Parallel,
}
impl Kind {
fn name(self) -> &'static str {
pub(crate) fn name(self) -> &'static str {
match self {
Self::Requests => "requests",
Self::Tokens => "tokens",
@ -286,15 +384,25 @@ struct Status {
remaining: i64,
}
/// The token reservation went over on the descriptor at `index`, which held `current`.
/// A counter to check and add to: `{prefix}:requests` or `{prefix}:tokens` in its window.
#[derive(Clone, Debug, PartialEq)]
pub(crate) struct Pair {
pub(crate) prefix: String,
pub(crate) kind: Kind,
pub(crate) limit: i64,
pub(crate) increment: i64,
}
/// Every pair's count after adding, or the pair at `index` that would go over, which held
/// `current`.
#[derive(Debug, PartialEq)]
enum Reserved {
pub(crate) enum Reserved {
Counts(Vec<i64>),
Over { index: usize, current: i64 },
}
/// Redis could not be used and the settings say to refuse rather than fall back.
struct Unavailable;
pub(crate) struct Unavailable;
#[derive(Debug)]
struct Window {
@ -343,27 +451,24 @@ impl Local {
.collect()
}
fn reserve(
&mut self,
scopes: &[(String, i64)],
increment: i64,
now: u64,
size: u64,
) -> Reserved {
for (index, (prefix, limit)) in scopes.iter().enumerate() {
let current = match self.windows.get(prefix) {
Some(window) if now.saturating_sub(window.start) < size => window.tokens,
fn check_and_add(&mut self, pairs: &[Pair], now: u64, size: u64) -> Reserved {
for (index, pair) in pairs.iter().enumerate() {
let current = match self.windows.get(&pair.prefix) {
Some(window) if now.saturating_sub(window.start) < size => match pair.kind {
Kind::Requests => window.requests,
_ => window.tokens,
},
_ => 0,
};
if blocked(current, increment, *limit) {
if blocked(current, pair.increment, pair.limit) {
return Reserved::Over { index, current };
}
}
Reserved::Counts(
scopes
pairs
.iter()
.map(|(prefix, _)| {
let (window, expired) = self.window(prefix, now, size);
.map(|pair| {
let (window, expired) = self.window(&pair.prefix, now, size);
if expired {
*window = Window {
start: now,
@ -371,13 +476,26 @@ impl Local {
tokens: 0,
};
}
window.tokens += increment;
window.tokens
let count = match pair.kind {
Kind::Requests => &mut window.requests,
_ => &mut window.tokens,
};
*count += pair.increment;
*count
})
.collect(),
)
}
/// A counter's value in its current window.
fn read(&self, prefix: &str, kind: Kind, now: u64, size: u64) -> Option<i64> {
let window = self.windows.get(prefix)?;
(now.saturating_sub(window.start) < size).then_some(match kind {
Kind::Requests => window.requests,
_ => window.tokens,
})
}
fn charge(&mut self, deltas: &[(String, i64)]) {
for (prefix, delta) in deltas {
if let Some(window) = self.windows.get_mut(prefix) {
@ -416,7 +534,7 @@ fn blocked(current: i64, increment: i64, limit: i64) -> bool {
}
/// Where the counters live.
struct Store {
pub(crate) struct Store {
redis: Option<Arc<Coordination>>,
local: Mutex<Local>,
window: u64,
@ -424,6 +542,15 @@ struct Store {
}
impl Store {
pub(crate) fn new(redis: Option<Arc<Coordination>>, window: u64, fail_closed: bool) -> Self {
Self {
redis,
local: Mutex::default(),
window,
fail_closed,
}
}
fn connection(&self) -> Option<(ConnectionManager, &Coordination)> {
let redis = self.redis.as_deref()?;
Some((redis.connection()?, redis))
@ -449,6 +576,34 @@ impl Store {
}
}
/// Counts the request on each prefix's window and returns the counts.
pub(crate) async fn requests(&self, prefixes: &[String]) -> Result<Vec<i64>, Unavailable> {
let now = epoch_seconds();
if self.redis.is_some() {
let error = match self.connection() {
Some((mut connection, redis)) => {
let mut invocation = REQUESTS_SCRIPT.prepare_invoke();
for prefix in prefixes {
invocation
.key(redis.namespaced(&format!("{prefix}:window")))
.key(redis.namespaced(&format!("{prefix}:requests")));
}
invocation.arg(now).arg(self.window);
match integers(invocation.invoke_async(&mut connection).await) {
Ok(reply) if reply.len() == prefixes.len() * 2 => {
return Ok(reply.iter().skip(1).step_by(2).copied().collect());
}
Ok(_) => Some(malformed()),
Err(error) => Some(error),
}
}
None => None,
};
self.fall_back(error)?;
}
Ok(self.local().requests(prefixes, now, self.window))
}
/// Counts the request, takes its in-flight slot and reserves its tokens, stopping at the first
/// limit it is over: one script with Redis, else this replica's counters.
async fn admit(&self, plan: &Plan) -> Result<(Admission, bool), Unavailable> {
@ -501,7 +656,7 @@ impl Store {
};
let tokens = match plan.tokens.is_empty() {
true => Vec::new(),
false => match local.reserve(&plan.tokens, plan.increment, now, self.window) {
false => match local.check_and_add(&plan.token_pairs(), now, self.window) {
Reserved::Counts(counts) => counts,
Reserved::Over { index, current } => {
if let Some(slot) = &plan.slot {
@ -536,8 +691,90 @@ impl Store {
self.local().release(prefix, slot);
}
/// Adds every pair's increment, or none when one would go over its limit. `true` when the
/// counters are in Redis.
pub(crate) async fn check_and_add(
&self,
pairs: &[Pair],
) -> Result<(Reserved, bool), Unavailable> {
if self.redis.is_some() {
let error = match self.connection() {
Some((mut connection, redis)) => {
let mut invocation = RESERVE_SCRIPT.prepare_invoke();
for pair in pairs {
invocation
.key(redis.namespaced(&format!("{}:window", pair.prefix)))
.key(redis.namespaced(&format!(
"{}:{}",
pair.prefix,
pair.kind.name()
)));
}
for pair in pairs {
invocation
.arg(pair.limit)
.arg(pair.increment)
.arg(self.window)
.arg(self.window);
}
match integers(invocation.invoke_async(&mut connection).await) {
Ok(reply) if reply.first() == Some(&0) => {
return Ok((
Reserved::Counts(
reply.iter().skip(1).step_by(2).copied().collect(),
),
true,
));
}
Ok(reply) => match reply.as_slice() {
[1, index, current, _] if (1..=pairs.len() as i64).contains(index) => {
return Ok((
Reserved::Over {
index: (*index - 1) as usize,
current: *current,
},
true,
));
}
_ => Some(malformed()),
},
Err(error) => Some(error),
}
}
None => None,
};
self.fall_back(error)?;
}
Ok((
self.local()
.check_and_add(pairs, epoch_seconds(), self.window),
false,
))
}
/// A counter's current value, `None` when it is not set.
pub(crate) async fn read(&self, prefix: &str, kind: Kind) -> Option<i64> {
if let Some((mut connection, redis)) = self.connection() {
let key = redis.namespaced(&format!("{prefix}:{}", kind.name()));
match redis::cmd("GET")
.arg(key)
.query_async::<Option<String>>(&mut connection)
.await
{
Ok(value) => {
return value
.and_then(|value| value.parse::<f64>().ok())
.map(|value| value as i64);
}
Err(error) => tracing::warn!(%error, "reading a rate limit counter failed"),
}
}
self.local()
.read(prefix, kind, epoch_seconds(), self.window)
}
/// Adds each delta to its prefix's tokens.
async fn charge(&self, deltas: Vec<(String, i64)>, shared: bool) {
pub(crate) async fn charge(&self, deltas: Vec<(String, i64)>, shared: bool) {
if shared && let Some((mut connection, redis)) = self.connection() {
let mut invocation = CHARGE_SCRIPT.prepare_invoke();
for (prefix, _) in &deltas {
@ -568,6 +805,21 @@ struct Plan {
increment: i64,
}
impl Plan {
/// The tokens limits as the pairs to reserve the increment on.
fn token_pairs(&self) -> Vec<Pair> {
self.tokens
.iter()
.map(|(prefix, limit)| Pair {
prefix: prefix.clone(),
kind: Kind::Tokens,
limit: *limit,
increment: self.increment,
})
.collect()
}
}
struct SlotPlan {
prefix: String,
limit: i64,
@ -798,12 +1050,7 @@ impl RateLimiter {
database: Option<PgPool>,
) -> Self {
Self {
store: Arc::new(Store {
redis,
local: Mutex::default(),
window: settings.window,
fail_closed: settings.fail_closed,
}),
store: Arc::new(Store::new(redis, settings.window, settings.fail_closed)),
settings,
end_users: database.map(EndUsers::new),
}
@ -888,22 +1135,15 @@ impl RateLimiter {
status.limit,
status.remaining.max(0),
);
let mut refusal = match request.route {
Route::Messages => Refusal::new(429, "rate_limit_error", message).anthropic(),
Route::ChatCompletions => {
let mut refusal = Refusal::new(429, "throttling_error", message).without_param();
for (name, value) in [
("retry-after", self.settings.window.to_string()),
("rate_limit_type", status.kind.name().to_owned()),
("reset_at", reset_at),
] {
if let Ok(value) = HeaderValue::from_str(&value) {
refusal = refusal.with_header(HeaderName::from_static(name), value);
}
}
refusal
}
};
let mut refusal = throttled(
request.route,
message,
vec![
("retry-after", self.settings.window.to_string()),
("rate_limit_type", status.kind.name().to_owned()),
("reset_at", reset_at),
],
);
for (name, value) in key_limit_headers(request.key.as_deref()) {
refusal = refusal.with_header(name, value);
}
@ -911,11 +1151,37 @@ impl RateLimiter {
}
fn unavailable(&self, request: &RequestFacts) -> Refusal {
match request.route {
Route::Messages => Refusal::new(503, "api_error", FAIL_CLOSED).anthropic(),
Route::ChatCompletions => {
Refusal::new(503, "internal_server_error", FAIL_CLOSED).without_param()
unavailable(request.route)
}
}
/// The proxy's 429 from one of its own limiters: on Chat Completions its error body with the
/// limiter's headers, on `/v1/messages` Anthropic's envelope, where the proxy drops them.
pub(crate) fn throttled(
route: Route,
message: String,
headers: Vec<(&'static str, String)>,
) -> Refusal {
match route {
Route::Messages => Refusal::new(429, "rate_limit_error", message).anthropic(),
Route::ChatCompletions => {
let mut refusal = Refusal::new(429, "throttling_error", message).without_param();
for (name, value) in headers {
if let Ok(value) = HeaderValue::from_str(&value) {
refusal = refusal.with_header(HeaderName::from_static(name), value);
}
}
refusal
}
}
}
/// The proxy's 503 when counters cannot be verified and enforcement fails closed.
pub(crate) fn unavailable(route: Route) -> Refusal {
match route {
Route::Messages => Refusal::new(503, "api_error", FAIL_CLOSED).anthropic(),
Route::ChatCompletions => {
Refusal::new(503, "internal_server_error", FAIL_CLOSED).without_param()
}
}
}
@ -1238,35 +1504,35 @@ mod tests {
assert_eq!(estimate(&request, Some(&key), "other", 100_000).tokens, 10);
}
fn reserve(local: &mut Local, increment: i64, now: u64) -> Reserved {
let pair = Pair {
prefix: "{api_key:h}".into(),
kind: Kind::Tokens,
limit: 10,
increment,
};
local.check_and_add(&[pair], now, 60)
}
#[rstest]
fn local_windows_count_reserve_and_reset() {
let mut local = Local::default();
let prefixes = vec!["{api_key:h}".to_owned()];
assert_eq!(local.requests(&prefixes, 100, 60), [1]);
assert_eq!(local.requests(&prefixes, 159, 60), [2]);
let scopes = vec![("{api_key:h}".to_owned(), 10)];
assert_eq!(reserve(&mut local, 6, 159), Reserved::Counts(vec![6]));
assert_eq!(
local.reserve(&scopes, 6, 159, 60),
Reserved::Counts(vec![6])
);
assert_eq!(
local.reserve(&scopes, 6, 159, 60),
reserve(&mut local, 6, 159),
Reserved::Over {
index: 0,
current: 6
}
);
local.charge(&[("{api_key:h}".to_owned(), -4)]);
assert_eq!(
local.reserve(&scopes, 8, 159, 60),
Reserved::Counts(vec![10])
);
assert_eq!(reserve(&mut local, 8, 159), Reserved::Counts(vec![10]));
// The window started at 100 is over at 160: requests and tokens start again.
assert_eq!(local.requests(&prefixes, 160, 60), [1]);
assert_eq!(
local.reserve(&scopes, 10, 160, 60),
Reserved::Counts(vec![10])
);
assert_eq!(reserve(&mut local, 10, 160), Reserved::Counts(vec![10]));
}
#[rstest]

556
litellm-rust/crates/gateway/src/request_hooks.rs Normal file
View file

@ -0,0 +1,556 @@
//! Built-in hooks a config turns on and configures under `litellm_settings.hooks`, run in the
//! order listed:
//!
//! ```yaml
//! litellm_settings:
//! hooks:
//! - forward_headers:
//! headers: ["x-request-origin", "x-trace-*"]
//! - spend_metadata:
//! from_headers:
//! x-litellm-spend-logs-metadata: json
//! x-cost-center: cost_center
//! from_key_metadata: [cost_center]
//! - strip_headers:
//! headers: ["x-litellm-spend-logs-metadata", "x-cost-center"]
//! - provider_metadata:
//! body:
//! metadata.user_id: "{end_user|key.user_id}"
//! headers:
//! x-caller-team: "{key.team_alias}"
//! ```
use std::sync::Arc;
use futures_util::future::BoxFuture;
use http::{HeaderName, HeaderValue};
use litellm_config::Config;
use litellm_gateway_hooks::{
Facts, HeaderPattern, Hook, Refusal, RequestEdit, RequestFacts, Route, Settle, Targets,
};
use serde::Deserialize;
use serde_json::{Map, Value};
use crate::{Error, spend::SPEND_LOGS_METADATA};
/// Headers a hook never forwards: credentials, and those the HTTP client or a request signer
/// sets itself.
const PROTECTED: [&str; 13] = [
"authorization",
"proxy-authorization",
"x-api-key",
"api-key",
"x-litellm-api-key",
"x-goog-api-key",
"cookie",
"host",
"content-length",
"content-type",
"transfer-encoding",
"connection",
"date",
];
fn protected(name: &str) -> bool {
PROTECTED
.iter()
.any(|protected| protected.eq_ignore_ascii_case(name))
|| name
.get(..6)
.is_some_and(|prefix| prefix.eq_ignore_ascii_case("x-amz-"))
}
/// The hooks `litellm_settings.hooks` lists, in its order.
pub fn from_config(config: &Config) -> Result<Vec<Arc<dyn Hook>>, Error> {
let Some(entries) = config.litellm_settings.additional_fields.get("hooks") else {
return Ok(Vec::new());
};
let litellm_config::Value::Sequence(entries) = entries else {
return Err(Error::Hooks("must be a list".into()));
};
entries.iter().map(hook).collect()
}
fn hook(entry: &litellm_config::Value) -> Result<Arc<dyn Hook>, Error> {
let (name, settings) = match entry {
litellm_config::Value::Mapping(mapping) if mapping.len() == 1 => {
let (name, settings) = mapping.iter().next().expect("one entry");
(name.as_str().unwrap_or_default(), settings.clone())
}
litellm_config::Value::String(name) => (name.as_str(), litellm_config::Value::Null),
_ => {
return Err(Error::Hooks(
"each entry is a hook name with its settings".into(),
));
}
};
let built: Result<Arc<dyn Hook>, String> = match name {
"forward_headers" => parse(settings)
.and_then(patterns)
.map(|patterns| Arc::new(ForwardHeaders { patterns }) as Arc<dyn Hook>),
"strip_headers" => parse(settings)
.and_then(patterns)
.map(|patterns| Arc::new(StripHeaders { patterns }) as Arc<dyn Hook>),
"spend_metadata" => parse(settings)
.and_then(SpendMetadata::new)
.map(|hook| Arc::new(hook) as Arc<dyn Hook>),
"provider_metadata" => parse(settings)
.and_then(ProviderMetadata::new)
.map(|hook| Arc::new(hook) as Arc<dyn Hook>),
other => return Err(Error::Hooks(format!("unknown hook {other}"))),
};
built.map_err(|error| Error::Hooks(format!("{name}: {error}")))
}
fn parse<T: serde::de::DeserializeOwned>(settings: litellm_config::Value) -> Result<T, String> {
T::deserialize(settings).map_err(|error| error.to_string())
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct HeaderList {
headers: Vec<String>,
}
fn patterns(list: HeaderList) -> Result<Vec<HeaderPattern>, String> {
list.headers
.iter()
.map(|pattern| {
HeaderPattern::parse(pattern).map_err(|_| format!("{pattern} is not a header name"))
})
.collect()
}
/// `forward_headers`: copies the client headers it names onto the provider request.
struct ForwardHeaders {
patterns: Vec<HeaderPattern>,
}
impl Hook for ForwardHeaders {
fn name(&self) -> &'static str {
"forward_headers"
}
fn reads(&self) -> Facts {
Facts::CLIENT_HEADERS
}
fn writes(&self) -> Targets {
Targets::UPSTREAM_HEADERS
}
fn run<'a>(
&'a self,
request: &'a Arc<RequestFacts>,
edit: &'a mut RequestEdit,
) -> BoxFuture<'a, Result<Option<Box<dyn Settle>>, Refusal>> {
for (name, value) in request.client_headers.iter() {
if !protected(name.as_str())
&& self
.patterns
.iter()
.any(|pattern| pattern.matches(name.as_str()))
{
edit.set_upstream_header(name.clone(), value.clone());
}
}
Box::pin(async { Ok(None) })
}
}
/// `strip_headers`: removes the headers it names from the provider request, whether a hook set
/// them or the gateway forwards them by default.
struct StripHeaders {
patterns: Vec<HeaderPattern>,
}
impl Hook for StripHeaders {
fn name(&self) -> &'static str {
"strip_headers"
}
fn reads(&self) -> Facts {
Facts::NONE
}
fn writes(&self) -> Targets {
Targets::UPSTREAM_HEADERS
}
fn run<'a>(
&'a self,
_request: &'a Arc<RequestFacts>,
edit: &'a mut RequestEdit,
) -> BoxFuture<'a, Result<Option<Box<dyn Settle>>, Refusal>> {
for pattern in &self.patterns {
edit.remove_upstream_headers(pattern.clone());
}
Box::pin(async { Ok(None) })
}
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct SpendMetadataSettings {
/// Header name to `json`, merged as an object, or to the key it is stored under.
#[serde(default)]
from_headers: litellm_config::Value,
/// Fields of the key's metadata copied under the same names.
#[serde(default)]
from_key_metadata: Vec<String>,
}
enum HeaderRule {
Json,
Field(String),
}
/// `spend_metadata`: adds caller-supplied values to the request's `spend_logs_metadata`
/// without changing the request.
struct SpendMetadata {
headers: Vec<(HeaderName, HeaderRule)>,
key_fields: Vec<String>,
}
impl SpendMetadata {
fn new(settings: SpendMetadataSettings) -> Result<Self, String> {
let headers = match settings.from_headers {
litellm_config::Value::Null => Vec::new(),
litellm_config::Value::Mapping(mapping) => mapping
.iter()
.map(|(name, rule)| {
let (Some(name), Some(rule)) = (name.as_str(), rule.as_str()) else {
return Err("from_headers maps header names to json or a key".to_owned());
};
let name = HeaderName::from_bytes(name.as_bytes())
.map_err(|_| format!("{name} is not a header name"))?;
let rule = match rule {
"json" => HeaderRule::Json,
field => HeaderRule::Field(field.to_owned()),
};
Ok((name, rule))
})
.collect::<Result<_, _>>()?,
_ => return Err("from_headers maps header names to json or a key".into()),
};
Ok(Self {
headers,
key_fields: settings.from_key_metadata,
})
}
}
impl Hook for SpendMetadata {
fn name(&self) -> &'static str {
"spend_metadata"
}
fn reads(&self) -> Facts {
Facts::CLIENT_HEADERS | Facts::KEY
}
fn writes(&self) -> Targets {
Targets::ANNOTATIONS
}
fn run<'a>(
&'a self,
request: &'a Arc<RequestFacts>,
edit: &'a mut RequestEdit,
) -> BoxFuture<'a, Result<Option<Box<dyn Settle>>, Refusal>> {
let mut metadata = Map::new();
for (name, rule) in &self.headers {
let Some(value) = request
.client_headers
.get(name)
.and_then(|value| value.to_str().ok())
else {
continue;
};
match rule {
HeaderRule::Json => match serde_json::from_str::<Value>(value) {
Ok(Value::Object(fields)) => metadata.extend(fields),
_ => tracing::warn!(
header = name.as_str(),
"spend metadata header is not a JSON object; ignored"
),
},
HeaderRule::Field(field) => {
metadata.insert(field.clone(), Value::from(value));
}
}
}
// The key's own values win over what the caller sent.
if let Some(key) = request.key.as_deref() {
for field in &self.key_fields {
if let Some(value) = key.metadata.get(field) {
metadata.insert(field.clone(), value.clone());
}
}
}
if !metadata.is_empty() {
edit.annotate(SPEND_LOGS_METADATA, Value::Object(metadata));
}
Box::pin(async { Ok(None) })
}
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct ProviderMetadataSettings {
/// Dotted body paths, such as `metadata.user_id`, to templates.
#[serde(default)]
body: std::collections::BTreeMap<String, String>,
#[serde(default)]
headers: std::collections::BTreeMap<String, String>,
}
/// A fact a template can name.
#[derive(Clone, Debug, PartialEq)]
enum Fact {
EndUser,
Model,
RequestId,
KeyAlias,
KeyUserId,
KeyTeamId,
KeyTeamAlias,
KeyMetadata(String),
TeamMetadata(String),
Header(HeaderName),
}
impl Fact {
fn parse(name: &str) -> Result<Self, String> {
Ok(match name {
"end_user" => Self::EndUser,
"model" => Self::Model,
"request_id" => Self::RequestId,
"key.alias" => Self::KeyAlias,
"key.user_id" => Self::KeyUserId,
"key.team_id" => Self::KeyTeamId,
"key.team_alias" => Self::KeyTeamAlias,
other => {
if let Some(field) = other.strip_prefix("key.metadata.") {
Self::KeyMetadata(field.to_owned())
} else if let Some(field) = other.strip_prefix("team.metadata.") {
Self::TeamMetadata(field.to_owned())
} else if let Some(header) = other.strip_prefix("header.") {
Self::Header(
HeaderName::from_bytes(header.as_bytes())
.map_err(|_| format!("{header} is not a header name"))?,
)
} else {
return Err(format!("unknown fact {other}"));
}
}
})
}
fn value(&self, request: &RequestFacts) -> Option<String> {
let key = request.key.as_deref();
let scalar = |value: &Value| match value {
Value::String(text) => Some(text.clone()),
Value::Number(number) => Some(number.to_string()),
Value::Bool(flag) => Some(flag.to_string()),
_ => None,
};
match self {
Self::EndUser => request.end_user.as_deref().map(str::to_owned),
Self::Model => Some(request.requested_model.to_string()),
Self::RequestId => Some(request.request_id.to_string()),
Self::KeyAlias => key?.alias.clone(),
Self::KeyUserId => key?.user_id.clone(),
Self::KeyTeamId => key?.team_id.clone(),
Self::KeyTeamAlias => key?.team_alias.clone(),
Self::KeyMetadata(field) => key?.metadata.get(field).and_then(scalar),
Self::TeamMetadata(field) => key?.team_metadata.get(field).and_then(scalar),
Self::Header(name) => request
.client_headers
.get(name)
.and_then(|value| value.to_str().ok())
.map(str::to_owned),
}
.filter(|value| !value.is_empty())
}
}
#[derive(Debug, PartialEq)]
enum Part {
Text(String),
/// The first of these facts that is set.
Facts(Vec<Fact>),
}
/// Text with `{fact|fallback}` placeholders, parsed once. A placeholder none of whose facts are
/// set leaves the whole value unset.
#[derive(Debug, PartialEq)]
struct Template(Vec<Part>);
impl Template {
fn parse(source: &str) -> Result<Self, String> {
let mut parts = Vec::new();
let mut rest = source;
while let Some(open) = rest.find('{') {
if open > 0 {
parts.push(Part::Text(rest[..open].to_owned()));
}
let close = rest[open..]
.find('}')
.ok_or_else(|| format!("unclosed placeholder in {source}"))?;
let facts = rest[open + 1..open + close]
.split('|')
.map(|fact| Fact::parse(fact.trim()))
.collect::<Result<_, _>>()?;
parts.push(Part::Facts(facts));
rest = &rest[open + close + 1..];
}
if !rest.is_empty() {
parts.push(Part::Text(rest.to_owned()));
}
Ok(Self(parts))
}
fn reads(&self) -> Facts {
let mut reads = Facts::NONE;
for part in &self.0 {
let Part::Facts(facts) = part else { continue };
for fact in facts {
reads |= match fact {
Fact::EndUser => Facts::END_USER,
Fact::Header(_) => Facts::CLIENT_HEADERS,
Fact::Model | Fact::RequestId => Facts::NONE,
_ => Facts::KEY,
};
}
}
reads
}
fn render(&self, request: &RequestFacts) -> Option<String> {
let mut rendered = String::new();
for part in &self.0 {
match part {
Part::Text(text) => rendered.push_str(text),
Part::Facts(facts) => {
rendered.push_str(&facts.iter().find_map(|fact| fact.value(request))?)
}
}
}
Some(rendered)
}
}
/// `provider_metadata`: sets body fields and headers on the provider request from facts about
/// the caller. Body fields are set on `/v1/messages` requests only.
struct ProviderMetadata {
body: Vec<(Vec<String>, Template)>,
headers: Vec<(HeaderName, Template)>,
reads: Facts,
}
impl ProviderMetadata {
fn new(settings: ProviderMetadataSettings) -> Result<Self, String> {
let body = settings
.body
.iter()
.map(|(path, template)| {
let path: Vec<String> = path.split('.').map(str::to_owned).collect();
if path.iter().any(String::is_empty) {
return Err(format!("{} is not a body path", path.join(".")));
}
Ok((path, Template::parse(template)?))
})
.collect::<Result<Vec<_>, String>>()?;
let headers = settings
.headers
.iter()
.map(|(name, template)| {
let header = HeaderName::from_bytes(name.as_bytes())
.map_err(|_| format!("{name} is not a header name"))?;
if protected(header.as_str()) {
return Err(format!("{name} cannot be set by a hook"));
}
Ok((header, Template::parse(template)?))
})
.collect::<Result<Vec<_>, String>>()?;
let reads = body
.iter()
.map(|(_, template)| template)
.chain(headers.iter().map(|(_, template)| template))
.fold(Facts::NONE, |reads, template| reads | template.reads());
Ok(Self {
body,
headers,
reads,
})
}
}
impl Hook for ProviderMetadata {
fn name(&self) -> &'static str {
"provider_metadata"
}
fn reads(&self) -> Facts {
self.reads
}
fn writes(&self) -> Targets {
Targets::UPSTREAM_BODY | Targets::UPSTREAM_HEADERS
}
fn run<'a>(
&'a self,
request: &'a Arc<RequestFacts>,
edit: &'a mut RequestEdit,
) -> BoxFuture<'a, Result<Option<Box<dyn Settle>>, Refusal>> {
// Body paths name Messages API fields; Chat Completions takes the headers only.
let body = match request.route {
Route::Messages => self.body.as_slice(),
Route::ChatCompletions => &[],
};
for (path, template) in body {
if let Some(value) = template.render(request) {
edit.set_upstream_body_path(path.clone(), Value::from(value));
}
}
for (name, template) in &self.headers {
if let Some(value) = template
.render(request)
.and_then(|value| HeaderValue::from_str(&value).ok())
{
edit.set_upstream_header(name.clone(), value);
}
}
Box::pin(async { Ok(None) })
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
#[rstest]
fn templates_name_facts_with_fallbacks() {
assert_eq!(
Template::parse("team-{key.team_alias|key.team_id}").unwrap(),
Template(vec![
Part::Text("team-".into()),
Part::Facts(vec![Fact::KeyTeamAlias, Fact::KeyTeamId]),
])
);
assert!(Template::parse("{nope}").is_err());
assert!(Template::parse("{end_user").is_err());
}
#[rstest]
#[case::credential("Authorization", true)]
#[case::signed("X-Amz-Date", true)]
#[case::trace("x-trace-id", false)]
fn credentials_are_never_forwarded(#[case] name: &str, #[case] expected: bool) {
assert_eq!(protected(name), expected);
}
}

676
litellm-rust/crates/gateway/src/s3_log/config.rs Normal file
View file

@ -0,0 +1,676 @@
//! `litellm_settings.s3_callback_params`, read the way the proxy's `S3Logger` reads it: values
//! may be `os.environ/` references, and a number or flag the proxy would not parse falls back to
//! its default with a warning instead of stopping startup.
use std::time::Duration;
use litellm_auth_aws::AwsAuthConfig;
use litellm_config::{Config, Value};
use litellm_http::SslVerify;
use litellm_secrets::source::SecretSource;
use super::key::Granularity;
use crate::Error;
/// The names a config's callback lists turn this observer on with.
const NAMES: [&str; 2] = ["s3_v2", "s3_log"];
const DEFAULT_FLUSH_INTERVAL_SECONDS: u64 = 10;
const DEFAULT_MAX_CONCURRENT_UPLOADS: usize = 16;
const DEFAULT_MAX_QUEUE_SIZE: usize = 50_000;
const DEFAULT_MAX_RETRY_AGE_SECONDS: u64 = 3600;
const DEFAULT_UPLOAD_ATTEMPT_TIMEOUT_SECONDS: u64 = 30;
/// Where the proxy signs non-LLM AWS calls when no region is configured.
const DEFAULT_REGION: &str = "us-west-2";
/// A configured value once its `os.environ/` reference is resolved.
#[derive(Clone, Debug, PartialEq)]
enum Param {
Bool(bool),
Int(i64),
Float(f64),
Text(String),
}
impl Param {
/// Python's `bool(value)`, which the proxy applies to the prefix and style flags: any
/// non-empty string is true, `"false"` included.
fn truthy(&self) -> bool {
match self {
Self::Bool(flag) => *flag,
Self::Int(number) => *number != 0,
Self::Float(number) => *number != 0.0,
Self::Text(text) => !text.is_empty(),
}
}
/// Pydantic's lax `bool`.
fn boolean(&self) -> Option<bool> {
match self {
Self::Bool(flag) => Some(*flag),
Self::Int(0) => Some(false),
Self::Int(1) => Some(true),
Self::Float(number) if *number == 0.0 => Some(false),
Self::Float(number) if *number == 1.0 => Some(true),
Self::Text(text) => match text.trim().to_ascii_lowercase().as_str() {
"0" | "off" | "f" | "false" | "n" | "no" => Some(false),
"1" | "on" | "t" | "true" | "y" | "yes" => Some(true),
_ => None,
},
_ => None,
}
}
/// Pydantic's lax `int`.
fn integer(&self) -> Option<i64> {
match self {
Self::Bool(flag) => Some(i64::from(*flag)),
Self::Int(number) => Some(*number),
Self::Float(number) if number.fract() == 0.0 && number.is_finite() => {
Some(*number as i64)
}
Self::Text(text) => text.trim().parse().ok(),
Self::Float(_) => None,
}
}
fn is_blank(&self) -> bool {
matches!(self, Self::Text(text) if text.is_empty())
}
}
/// The resolved `s3_callback_params`.
struct Params(Vec<(String, Param)>);
impl Params {
fn get(&self, name: &str) -> Option<&Param> {
self.0
.iter()
.find(|(key, _)| key == name)
.map(|(_, value)| value)
}
/// `params.get(name) or None`: an empty or false value counts as unset.
fn text(&self, name: &str) -> Result<Option<String>, Error> {
match self.get(name) {
None => Ok(None),
Some(Param::Text(text)) => Ok(Some(text.clone()).filter(|text| !text.is_empty())),
Some(Param::Int(number)) => Ok(Some(number.to_string())),
Some(Param::Float(number)) => Ok(Some(number.to_string())),
Some(Param::Bool(_)) => Err(Error::S3Log(format!("{name} must be a string"))),
}
}
fn truthy(&self, name: &str) -> bool {
self.get(name).is_some_and(Param::truthy)
}
/// A flag the proxy parses with pydantic; an invalid value warns and takes `invalid`.
fn flag(&self, name: &str, default: bool, invalid: bool) -> bool {
match self.get(name) {
None => default,
Some(value) if value.is_blank() => default,
Some(value) => value.boolean().unwrap_or_else(|| {
tracing::warn!(
setting = name,
?value,
"s3 logging: not a boolean, using {invalid}"
);
invalid
}),
}
}
/// A positive integer; anything else warns and takes `default`.
fn positive(&self, name: &str, default: usize, reject_bool: bool) -> usize {
let Some(value) = self.get(name).filter(|value| !value.is_blank()) else {
return default;
};
let parsed = match value {
Param::Bool(_) if reject_bool => None,
value => value.integer(),
};
match parsed.and_then(|number| usize::try_from(number).ok()) {
Some(number) if number >= 1 => number,
_ => {
tracing::warn!(
setting = name,
?value,
"s3 logging: not a positive integer, using {default}"
);
default
}
}
}
}
/// How each request is logged and where it goes.
#[derive(Clone, Debug, PartialEq)]
pub(super) struct Settings {
pub bucket: String,
/// `s3_path`, the prefix every key starts with.
pub path: String,
/// The region requests are signed for and the default endpoint names.
pub region: String,
pub endpoint_url: Option<String>,
pub virtual_hosted_style: bool,
/// `s3_verify`; unset keeps the gateway's own TLS settings.
pub verify: Option<SslVerify>,
pub auth: AwsAuthConfig,
pub team_prefix: bool,
pub key_prefix: bool,
pub strip_base64_files: bool,
pub prompts_only: bool,
pub granularity: Granularity,
pub server_side_encryption: Option<String>,
pub sse_kms_key_id: Option<String>,
pub max_concurrent_uploads: usize,
pub max_queue_size: usize,
/// How long an upload that keeps failing is retried; `None` retries it until the queue
/// overflows.
pub max_retry_age: Option<Duration>,
pub drop_on_terminal_error: bool,
pub batch_file_upload: bool,
pub flush_interval: Duration,
pub attempt_timeout: Duration,
pub log_success: bool,
pub log_failure: bool,
pub turn_off_message_logging: bool,
pub redact_user_api_key_info: bool,
pub user_agent_tags: bool,
}
impl Settings {
/// `Ok(None)` when no callback list names `s3_v2` or `s3_log`. A config the proxy could
/// never upload with stops the gateway.
pub async fn from_config(
config: &Config,
secrets: &dyn SecretSource,
env: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Option<Self>, Error> {
let litellm = &config.litellm_settings;
let listed = |list: &Option<Vec<Value>>| {
list.iter().flatten().any(|callback| {
callback
.as_str()
.is_some_and(|name| NAMES.contains(&name.trim()))
})
};
let everywhere = listed(&litellm.callbacks);
let log_success = everywhere || listed(&litellm.success_callback);
let log_failure = everywhere || listed(&litellm.failure_callback);
if !log_success && !log_failure {
return Ok(None);
}
let params = resolve(litellm.additional_fields.get("s3_callback_params"), secrets).await?;
for (name, why) in [
(
"s3_flush_interval",
"the proxy reads DEFAULT_S3_FLUSH_INTERVAL_SECONDS instead",
),
("s3_batch_size", "the proxy flushes on its interval alone"),
(
"s3_adaptive_concurrency",
"this gateway uploads s3_max_concurrent_uploads at a time",
),
] {
if params.get(name).is_some() {
tracing::warn!(setting = name, "s3 logging: ignored, {why}");
}
}
let bucket = params
.text("s3_bucket_name")?
.ok_or_else(|| Error::S3Log("s3_bucket_name is required".into()))?;
let configured_region = params.text("s3_region_name")?;
let region = configured_region
.clone()
.or_else(|| env("AWS_REGION"))
.or_else(|| env("AWS_REGION_NAME"))
.unwrap_or_else(|| DEFAULT_REGION.into());
let (server_side_encryption, sse_kms_key_id) = sse(
params.text("s3_server_side_encryption")?,
params.text("s3_sse_kms_key_id")?,
);
let verify = match params.get("s3_verify") {
None => None,
Some(Param::Bool(true)) => Some(SslVerify::Enabled),
Some(Param::Bool(false)) => Some(SslVerify::Disabled),
Some(Param::Text(text)) => Some(SslVerify::parse(text)),
Some(other) => {
return Err(Error::S3Log(format!(
"s3_verify must be a boolean or a CA bundle path, not {other:?}"
)));
}
};
let env_number = |name: &str, default: u64| {
env(name)
.and_then(|value| value.trim().parse::<u64>().ok())
.unwrap_or(default)
};
let max_retry_age = match params.get("s3_max_retry_age_seconds") {
None => Some(DEFAULT_MAX_RETRY_AGE_SECONDS),
Some(value) if value.is_blank() => Some(DEFAULT_MAX_RETRY_AGE_SECONDS),
Some(value) => match value {
Param::Bool(_) => None,
value => value.integer(),
}
.filter(|seconds| *seconds >= 0)
.map(|seconds| seconds as u64)
.or_else(|| {
tracing::warn!(?value, "s3 logging: s3_max_retry_age_seconds is not a non-negative integer, using {DEFAULT_MAX_RETRY_AGE_SECONDS}");
Some(DEFAULT_MAX_RETRY_AGE_SECONDS)
}),
}
.filter(|seconds| *seconds > 0)
.map(Duration::from_secs);
let prompts_only = match params.get("s3_log_prompts_only").cloned() {
None => env("S3_LOG_PROMPTS_ONLY").map(Param::Text),
value => value,
}
.filter(|value| !value.is_blank())
.is_some_and(|value| {
value.boolean().unwrap_or_else(|| {
tracing::warn!(
?value,
"s3 logging: s3_log_prompts_only is not a boolean, logging prompts only"
);
true
})
});
let granularity = match params.get("s3_partition_granularity").cloned() {
None => env("S3_PARTITION_GRANULARITY").map(Param::Text),
value => value,
};
let granularity = match granularity {
Some(Param::Text(text)) if text == "hour" => Granularity::Hour,
None => Granularity::Day,
Some(Param::Text(text)) if text.is_empty() || text == "day" => Granularity::Day,
Some(value) => {
tracing::warn!(
?value,
"s3 logging: s3_partition_granularity is not one of day, hour, using day"
);
Granularity::Day
}
};
let litellm_flag =
|name: &str| matches!(litellm.additional_fields.get(name), Some(Value::Bool(true)));
Ok(Some(Self {
bucket,
path: params.text("s3_path")?.unwrap_or_default(),
region: region.clone(),
endpoint_url: params.text("s3_endpoint_url")?,
virtual_hosted_style: params.truthy("s3_use_virtual_hosted_style"),
verify,
auth: AwsAuthConfig {
access_key_id: params.text("s3_aws_access_key_id")?,
secret_access_key: params.text("s3_aws_secret_access_key")?,
session_token: params.text("s3_aws_session_token")?,
region_name: Some(region),
session_name: params.text("s3_aws_session_name")?,
profile_name: params.text("s3_aws_profile_name")?,
role_name: params.text("s3_aws_role_name")?,
web_identity_token: params.text("s3_aws_web_identity_token")?,
sts_endpoint: params.text("s3_aws_sts_endpoint")?,
external_id: None,
},
team_prefix: params.truthy("s3_use_team_prefix"),
key_prefix: params.truthy("s3_use_key_prefix"),
strip_base64_files: params.truthy("s3_strip_base64_files"),
prompts_only,
granularity,
server_side_encryption,
sse_kms_key_id,
max_concurrent_uploads: params.positive(
"s3_max_concurrent_uploads",
env_number(
"DEFAULT_S3_MAX_CONCURRENT_UPLOADS",
DEFAULT_MAX_CONCURRENT_UPLOADS as u64,
) as usize,
false,
),
max_queue_size: params.positive("s3_max_queue_size", DEFAULT_MAX_QUEUE_SIZE, true),
max_retry_age,
drop_on_terminal_error: params.flag("s3_drop_on_terminal_error", true, true),
batch_file_upload: params.flag("s3_batch_file_upload", false, false),
flush_interval: Duration::from_secs(
env_number(
"DEFAULT_S3_FLUSH_INTERVAL_SECONDS",
DEFAULT_FLUSH_INTERVAL_SECONDS,
)
.max(1),
),
attempt_timeout: Duration::from_secs(
env_number(
"S3_UPLOAD_ATTEMPT_TIMEOUT_SECONDS",
DEFAULT_UPLOAD_ATTEMPT_TIMEOUT_SECONDS,
)
.clamp(1, 3600),
),
log_success,
log_failure,
turn_off_message_logging: litellm_flag("turn_off_message_logging"),
redact_user_api_key_info: litellm_flag("redact_user_api_key_info"),
user_agent_tags: !litellm_flag("disable_add_user_agent_to_request_tags"),
}))
}
}
/// The encryption headers to send: a KMS key alone means `aws:kms`, and a key next to another
/// algorithm is dropped, as the proxy's `resolve_sse_params` does.
fn sse(algorithm: Option<String>, key_id: Option<String>) -> (Option<String>, Option<String>) {
let algorithm = algorithm.or_else(|| key_id.as_ref().map(|_| "aws:kms".into()));
match (algorithm, key_id) {
(None, _) => (None, None),
(Some(algorithm), Some(_)) if !algorithm.starts_with("aws:kms") => {
tracing::warn!(%algorithm, "s3 logging: ignoring s3_sse_kms_key_id, set s3_server_side_encryption to aws:kms to use it");
(Some(algorithm), None)
}
(algorithm, key_id) => (algorithm, key_id),
}
}
async fn resolve(value: Option<&Value>, secrets: &dyn SecretSource) -> Result<Params, Error> {
let entries = match value {
None | Some(Value::Null) => return Ok(Params(Vec::new())),
Some(Value::Mapping(entries)) => entries,
Some(_) => return Err(Error::S3Log("s3_callback_params must be a mapping".into())),
};
let mut params = Vec::with_capacity(entries.len());
for (key, value) in entries {
let Some(name) = key.as_str() else {
return Err(Error::S3Log(
"s3_callback_params keys must be strings".into(),
));
};
let param = match value {
Value::Null => continue,
Value::Bool(flag) => Param::Bool(*flag),
Value::Number(number) => match number.as_i64() {
Some(number) => Param::Int(number),
None => Param::Float(number.as_f64().unwrap_or(f64::NAN)),
},
Value::String(text) => match text.strip_prefix("os.environ/") {
Some(variable) => match secrets.get_secret_str(variable).await? {
Some(secret) => Param::Text(secret.expose().to_owned()),
None => continue,
},
None => Param::Text(text.clone()),
},
_ => return Err(Error::S3Log(format!("{name} must be a scalar"))),
};
params.push((name.to_owned(), param));
}
Ok(Params(params))
}
#[cfg(test)]
pub(super) mod tests {
use std::collections::HashMap;
use futures_util::future::BoxFuture;
use litellm_auth_types::SecretValue;
use rstest::rstest;
use super::*;
struct Secrets(HashMap<&'static str, &'static str>);
impl SecretSource for Secrets {
fn get_secret_str<'a>(
&'a self,
name: &'a str,
) -> BoxFuture<'a, Result<Option<SecretValue>, litellm_secrets::Error>> {
Box::pin(async move { Ok(self.0.get(name).map(|value| SecretValue::new(*value))) })
}
}
pub(crate) async fn settings_with(
yaml: &str,
environment: &[(&'static str, &'static str)],
) -> Result<Option<Settings>, String> {
let config = Config::from_yaml(&format!("model_list: []\n{yaml}")).unwrap();
let environment: HashMap<&str, &str> = environment.iter().copied().collect();
let env = |name: &str| environment.get(name).map(|value| (*value).to_owned());
Settings::from_config(&config, &Secrets(environment.clone()), &env)
.await
.map_err(|error| error.to_string())
}
async fn settings(params: &str) -> Settings {
settings_with(
&format!(
"litellm_settings:\n callbacks: [s3_v2]\n s3_callback_params:\n s3_bucket_name: logs\n{params}"
),
&[],
)
.await
.unwrap()
.unwrap()
}
#[rstest]
#[case::none("", None)]
#[case::other_callback(" success_callback: [prometheus]\n", None)]
#[case::callbacks(" callbacks: [s3_v2]\n", Some((true, true)))]
#[case::gateway_name(" callbacks: [s3_log]\n", Some((true, true)))]
#[case::success_only(" success_callback: [s3_v2]\n", Some((true, false)))]
#[case::failure_only(" failure_callback: s3_v2\n", Some((false, true)))]
#[tokio::test]
async fn on_for_the_lists_that_name_it(
#[case] lists: &str,
#[case] expected: Option<(bool, bool)>,
) {
let yaml =
format!("litellm_settings:\n{lists} s3_callback_params:\n s3_bucket_name: logs\n");
let settings = settings_with(&yaml, &[]).await.unwrap();
assert_eq!(
settings.map(|settings| (settings.log_success, settings.log_failure)),
expected
);
}
#[tokio::test]
async fn defaults_are_the_proxys() {
let settings = settings("").await;
assert_eq!(settings.path, "");
assert_eq!(settings.region, "us-west-2");
assert_eq!(settings.granularity, Granularity::Day);
assert_eq!(settings.max_concurrent_uploads, 16);
assert_eq!(settings.max_queue_size, 50_000);
assert_eq!(settings.max_retry_age, Some(Duration::from_secs(3600)));
assert!(settings.drop_on_terminal_error);
assert!(!settings.batch_file_upload);
assert!(!settings.prompts_only);
assert_eq!(settings.flush_interval, Duration::from_secs(10));
assert_eq!(settings.attempt_timeout, Duration::from_secs(30));
assert_eq!(settings.verify, None);
assert_eq!(settings.server_side_encryption, None);
}
#[tokio::test]
async fn references_are_resolved_from_the_secret_source() {
let settings = settings_with(
"litellm_settings:
callbacks: [s3_v2]
s3_callback_params:
s3_bucket_name: os.environ/BUCKET
s3_region_name: eu-west-1
s3_aws_access_key_id: os.environ/KEY_ID
s3_aws_secret_access_key: os.environ/SECRET
s3_aws_session_token: os.environ/UNSET
s3_use_team_prefix: os.environ/TEAM_PREFIX
",
&[
("BUCKET", "audit"),
("KEY_ID", "AKID"),
("SECRET", "shh"),
("TEAM_PREFIX", "false"),
],
)
.await
.unwrap()
.unwrap();
assert_eq!(settings.bucket, "audit");
assert_eq!(settings.region, "eu-west-1");
assert_eq!(settings.auth.access_key_id.as_deref(), Some("AKID"));
assert_eq!(settings.auth.secret_access_key.as_deref(), Some("shh"));
assert_eq!(settings.auth.session_token, None);
assert_eq!(settings.auth.region_name.as_deref(), Some("eu-west-1"));
// The proxy takes `bool()` of the string, so any non-empty text turns it on.
assert!(settings.team_prefix);
}
#[rstest]
#[case::configured(" s3_region_name: ap-south-1\n", &[("AWS_REGION", "us-east-2")], "ap-south-1")]
#[case::aws_region_wins("", &[("AWS_REGION", "us-east-2"), ("AWS_REGION_NAME", "eu-west-3")], "us-east-2")]
#[case::aws_region_name("", &[("AWS_REGION_NAME", "eu-west-3")], "eu-west-3")]
#[tokio::test]
async fn the_signing_region_falls_back_like_the_proxy(
#[case] params: &str,
#[case] environment: &[(&'static str, &'static str)],
#[case] expected: &str,
) {
let settings = settings_with(
&format!(
"litellm_settings:\n callbacks: [s3_v2]\n s3_callback_params:\n s3_bucket_name: logs\n{params}"
),
environment,
)
.await
.unwrap()
.unwrap();
assert_eq!(settings.region, expected);
}
#[rstest]
#[case::kms_key_alone(" s3_sse_kms_key_id: key-1\n", Some("aws:kms"), Some("key-1"))]
#[case::aes(" s3_server_side_encryption: AES256\n", Some("AES256"), None)]
#[case::key_with_aes(
" s3_server_side_encryption: AES256\n s3_sse_kms_key_id: key-1\n",
Some("AES256"),
None
)]
#[case::dsse(
" s3_server_side_encryption: aws:kms:dsse\n s3_sse_kms_key_id: key-1\n",
Some("aws:kms:dsse"),
Some("key-1")
)]
#[tokio::test]
async fn encryption_settings_combine_like_the_proxy(
#[case] params: &str,
#[case] algorithm: Option<&str>,
#[case] key: Option<&str>,
) {
let settings = settings(params).await;
assert_eq!(settings.server_side_encryption.as_deref(), algorithm);
assert_eq!(settings.sse_kms_key_id.as_deref(), key);
}
#[rstest]
#[case::text_number(" s3_max_concurrent_uploads: ' 4 '\n", 4, 50_000)]
#[case::zero_falls_back(" s3_max_concurrent_uploads: 0\n", 16, 50_000)]
#[case::garbage_falls_back(" s3_max_queue_size: lots\n", 16, 50_000)]
#[case::bool_queue_falls_back(" s3_max_queue_size: true\n", 16, 50_000)]
#[case::bool_width_is_one(" s3_max_concurrent_uploads: true\n", 1, 50_000)]
#[case::queue(" s3_max_queue_size: 10\n", 16, 10)]
#[tokio::test]
async fn numbers_parse_or_fall_back(
#[case] params: &str,
#[case] uploads: usize,
#[case] queue: usize,
) {
let settings = settings(params).await;
assert_eq!(settings.max_concurrent_uploads, uploads);
assert_eq!(settings.max_queue_size, queue);
}
#[rstest]
#[case::zero_disables(" s3_max_retry_age_seconds: 0\n", None)]
#[case::seconds(" s3_max_retry_age_seconds: '60'\n", Some(60))]
#[case::negative_falls_back(" s3_max_retry_age_seconds: -1\n", Some(3600))]
#[case::blank_is_default(" s3_max_retry_age_seconds: ''\n", Some(3600))]
#[tokio::test]
async fn retry_age(#[case] params: &str, #[case] seconds: Option<u64>) {
assert_eq!(
settings(params).await.max_retry_age,
seconds.map(Duration::from_secs)
);
}
#[rstest]
#[case::configured(" s3_log_prompts_only: 'yes'\n", &[], true)]
#[case::invalid_logs_prompts_only(" s3_log_prompts_only: sometimes\n", &[], true)]
#[case::environment("", &[("S3_LOG_PROMPTS_ONLY", "true")], true)]
#[case::configured_beats_environment(" s3_log_prompts_only: false\n", &[("S3_LOG_PROMPTS_ONLY", "true")], false)]
#[case::unset("", &[], false)]
#[tokio::test]
async fn prompts_only(
#[case] params: &str,
#[case] environment: &[(&'static str, &'static str)],
#[case] expected: bool,
) {
let settings = settings_with(
&format!(
"litellm_settings:\n callbacks: [s3_v2]\n s3_callback_params:\n s3_bucket_name: logs\n{params}"
),
environment,
)
.await
.unwrap()
.unwrap();
assert_eq!(settings.prompts_only, expected);
}
#[rstest]
#[case::hour(" s3_partition_granularity: hour\n", &[], Granularity::Hour)]
#[case::environment("", &[("S3_PARTITION_GRANULARITY", "hour")], Granularity::Hour)]
#[case::unknown_is_day(" s3_partition_granularity: minute\n", &[], Granularity::Day)]
#[tokio::test]
async fn partition_granularity(
#[case] params: &str,
#[case] environment: &[(&'static str, &'static str)],
#[case] expected: Granularity,
) {
let settings = settings_with(
&format!(
"litellm_settings:\n callbacks: [s3_v2]\n s3_callback_params:\n s3_bucket_name: logs\n{params}"
),
environment,
)
.await
.unwrap()
.unwrap();
assert_eq!(settings.granularity, expected);
}
#[rstest]
#[case::no_bucket(
"litellm_settings:\n callbacks: [s3_v2]\n",
"s3_bucket_name is required"
)]
#[case::unset_bucket_reference(
"litellm_settings:\n callbacks: [s3_v2]\n s3_callback_params:\n s3_bucket_name: os.environ/NOPE\n",
"s3_bucket_name is required"
)]
#[case::not_a_mapping(
"litellm_settings:\n callbacks: [s3_v2]\n s3_callback_params: [a]\n",
"s3_callback_params must be a mapping"
)]
#[case::nested_value(
"litellm_settings:\n callbacks: [s3_v2]\n s3_callback_params:\n s3_bucket_name: logs\n s3_path: [a, b]\n",
"s3_path must be a scalar"
)]
#[case::bad_verify(
"litellm_settings:\n callbacks: [s3_v2]\n s3_callback_params:\n s3_bucket_name: logs\n s3_verify: 3\n",
"s3_verify must be a boolean or a CA bundle path"
)]
#[tokio::test]
async fn a_config_that_cannot_upload_stops_the_gateway(
#[case] yaml: &str,
#[case] message: &str,
) {
let error = settings_with(yaml, &[]).await.unwrap_err();
assert!(error.contains(message), "{error}");
}
}

389
litellm-rust/crates/gateway/src/s3_log/key.rs Normal file
View file

@ -0,0 +1,389 @@
//! Object keys, download names and URLs, laid out as the proxy's `get_s3_object_key` and
//! `S3Logger._build_object_url` lay them out, length bounds included.
use std::time::{SystemTime, UNIX_EPOCH};
use sha2::{Digest, Sha256};
use super::config::Settings;
use crate::spend::civil;
/// S3's limit on a key, in UTF-8 bytes.
const MAX_KEY_BYTES: usize = 1024;
/// The proxy's limit on the last segment of a key.
const MAX_FILE_NAME_BYTES: usize = 255;
/// The proxy's limit on the `Content-Disposition` file name.
const MAX_DOWNLOAD_NAME_BYTES: usize = 1024;
/// How much of an over-long name is kept before its digest.
const BOUNDED_HEAD_BYTES: i64 = 64;
const MIN_BOUNDED_FILE_NAME_BYTES: usize = 64;
const PREFIX_DIGEST_CHARS: usize = 16;
const EXTENSION: &str = ".json";
/// `s3_partition_granularity`: a folder per day, or per hour within it.
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub(super) enum Granularity {
#[default]
Day,
Hour,
}
/// A wall-clock time, UTC, split the way the key names it.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(super) struct Utc {
year: i64,
month: u32,
day: u32,
hour: u64,
minute: u64,
second: u64,
micros: u32,
}
impl Utc {
pub fn new(time: SystemTime) -> Self {
let elapsed = time.duration_since(UNIX_EPOCH).unwrap_or_default();
let seconds = elapsed.as_secs();
let (year, month, day) = civil((seconds / 86_400) as i64);
let of_day = seconds % 86_400;
Self {
year,
month,
day,
hour: of_day / 3600,
minute: of_day / 60 % 60,
second: of_day % 60,
micros: elapsed.subsec_micros(),
}
}
fn date(&self) -> String {
format!("{:04}-{:02}-{:02}", self.year, self.month, self.day)
}
/// `%H-%M-%S`.
pub fn clock(&self) -> String {
format!("{:02}-{:02}-{:02}", self.hour, self.minute, self.second)
}
/// `%H-%M-%S-%f`.
fn clock_micros(&self) -> String {
format!("{}-{:06}", self.clock(), self.micros)
}
}
/// The proxy's `get_logging_id`: the request's time and the payload id.
pub(super) fn logging_id(start: &Utc, id: &str) -> String {
format!("time-{}_{id}", start.clock_micros())
}
/// `{s3_path}/[prefix]YYYY-MM-DD[/HH]/{file_name}.json`, shortened as the proxy shortens it when
/// it would pass S3's limits: the file name first, the configured prefix only if that is what
/// does not fit.
pub(super) fn object_key(
path: &str,
prefix: &str,
start: &Utc,
file_name: &str,
granularity: Granularity,
) -> String {
let sanitized = file_name.replace(['/', ':'], "_");
let configured = match path.is_empty() {
true => prefix.to_owned(),
false => format!("{}/{prefix}", path.trim_end_matches('/')),
};
let date = match granularity {
Granularity::Day => format!("{}/", start.date()),
Granularity::Hour => format!("{}/{:02}/", start.date(), start.hour),
};
let file_name_budget = MAX_FILE_NAME_BYTES - EXTENSION.len();
let key = format!("{configured}{date}{sanitized}{EXTENSION}");
if key.len() <= MAX_KEY_BYTES && sanitized.len() <= file_name_budget {
return key;
}
let budget = MAX_KEY_BYTES - date.len() - EXTENSION.len();
if configured.len() + MIN_BOUNDED_FILE_NAME_BYTES <= budget {
let bounded = bounded_name(
file_name,
&sanitized,
file_name_budget.min(budget - configured.len()),
);
return format!("{configured}{date}{bounded}{EXTENSION}");
}
let shortest = bounded_name(file_name, &sanitized, MIN_BOUNDED_FILE_NAME_BYTES);
let prefix = bounded_prefix(&configured, budget.saturating_sub(shortest.len()));
format!("{prefix}{date}{shortest}{EXTENSION}")
}
/// The proxy's `get_s3_object_download_filename`.
pub(super) fn download_name(start: &Utc, id: &str) -> String {
let stamp = format!("time-{}T{}", start.date(), start.clock_micros());
let name = format!("{stamp}_{id}");
let sanitized = format!("{stamp}_{}", id.replace(['/', '"'], "_"));
let budget = MAX_DOWNLOAD_NAME_BYTES - EXTENSION.len();
match sanitized.len() <= budget {
true => format!("{sanitized}{EXTENSION}"),
false => format!("{}{EXTENSION}", bounded_name(&name, &sanitized, budget)),
}
}
/// As much of `sanitized` as fits, then the sha256 of the unsanitized name.
fn bounded_name(name: &str, sanitized: &str, max_bytes: usize) -> String {
let digest = hex_sha256(name);
let head_budget = BOUNDED_HEAD_BYTES.min(max_bytes as i64 - digest.len() as i64 - 1);
let head = truncate(sanitized, head_budget);
match head.is_empty() {
true => digest,
false => format!("{head}_{digest}"),
}
}
/// As much of the configured prefix as fits, then a digest segment naming all of it.
fn bounded_prefix(prefix: &str, max_bytes: usize) -> String {
let digest = format!("{}/", &hex_sha256(prefix)[..PREFIX_DIGEST_CHARS]);
if max_bytes < digest.len() {
return String::new();
}
let head = truncate(prefix, max_bytes as i64 - digest.len() as i64 - 1).trim_end_matches('/');
match head.is_empty() {
true => digest,
false => format!("{head}/{digest}"),
}
}
/// `value` cut to at most `max_bytes` of UTF-8, never inside a character.
fn truncate(value: &str, max_bytes: i64) -> &str {
let Ok(max_bytes) = usize::try_from(max_bytes) else {
return "";
};
if value.len() <= max_bytes {
return value;
}
let end = (0..=max_bytes)
.rev()
.find(|end| value.is_char_boundary(*end))
.unwrap_or(0);
&value[..end]
}
fn hex_sha256(value: &str) -> String {
Sha256::digest(value.as_bytes())
.iter()
.map(|byte| format!("{byte:02x}"))
.collect()
}
/// Everything before the last `/`, the folder batch files are grouped by.
pub(super) fn parent(key: &str) -> &str {
key.rsplit_once('/').map_or("", |(parent, _)| parent)
}
/// The URL that is signed and sent, with the key percent-encoded once, as the proxy builds it.
pub(super) fn object_url(settings: &Settings, key: &str) -> String {
let key = encode(key);
let bucket = &settings.bucket;
match &settings.endpoint_url {
Some(endpoint) if settings.virtual_hosted_style => {
let (scheme, host) = match endpoint.strip_prefix("https://") {
Some(host) => ("https://", host),
None => ("http://", endpoint.trim_start_matches("http://")),
};
format!("{scheme}{bucket}.{host}/{key}")
}
Some(endpoint) => format!("{endpoint}/{bucket}/{key}"),
None => format!(
"https://{bucket}.s3.{region}.{suffix}/{key}",
region = settings.region,
suffix = dns_suffix(&settings.region)
),
}
}
/// Python's `quote(key, safe="/")`.
fn encode(key: &str) -> String {
key.bytes()
.map(|byte| match byte {
b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'.' | b'_' | b'~' | b'/' => {
char::from(byte).to_string()
}
byte => format!("%{byte:02X}"),
})
.collect()
}
/// The AWS partition's domain for a region.
fn dns_suffix(region: &str) -> &'static str {
[
("cn-", "amazonaws.com.cn"),
("us-isob-", "sc2s.sgov.gov"),
("us-isof-", "csp.hci.ic.gov"),
("us-iso-", "c2s.ic.gov"),
("eu-isoe-", "cloud.adc-e.uk"),
]
.into_iter()
.find(|(prefix, _)| region.starts_with(prefix))
.map_or("amazonaws.com", |(_, suffix)| suffix)
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use rstest::rstest;
use super::*;
/// 2026-03-04 05:06:07.000089 UTC.
fn start() -> Utc {
Utc::new(UNIX_EPOCH + Duration::from_micros(1_772_600_767_000_089))
}
#[rstest]
#[case::bare(
"",
"",
Granularity::Day,
"2026-03-04/time-05-06-07-000089_chatcmpl-1.json"
)]
#[case::path(
"logs/",
"",
Granularity::Day,
"logs/2026-03-04/time-05-06-07-000089_chatcmpl-1.json"
)]
#[case::prefixes(
"logs",
"research/team-key/",
Granularity::Day,
"logs/research/team-key/2026-03-04/time-05-06-07-000089_chatcmpl-1.json"
)]
#[case::hour(
"logs",
"",
Granularity::Hour,
"logs/2026-03-04/05/time-05-06-07-000089_chatcmpl-1.json"
)]
fn keys_are_laid_out_like_the_proxy(
#[case] path: &str,
#[case] prefix: &str,
#[case] granularity: Granularity,
#[case] expected: &str,
) {
let file_name = logging_id(&start(), "chatcmpl-1");
assert_eq!(
object_key(path, prefix, &start(), &file_name, granularity),
expected
);
}
#[test]
fn slashes_and_colons_in_the_id_stay_in_one_segment() {
let file_name = logging_id(&start(), "a/b:c");
assert_eq!(
object_key("", "", &start(), &file_name, Granularity::Day),
"2026-03-04/time-05-06-07-000089_a_b_c.json"
);
assert_eq!(
download_name(&start(), "a/b\"c"),
"time-2026-03-04T05-06-07-000089_a_b_c.json"
);
}
#[test]
fn a_long_id_is_shortened_to_a_digest_within_s3_limits() {
let id = "é".repeat(400);
let file_name = logging_id(&start(), &id);
let key = object_key("logs", "", &start(), &file_name, Granularity::Day);
let (folder, last) = key.rsplit_once('/').unwrap();
assert_eq!(folder, "logs/2026-03-04");
assert!(last.len() <= MAX_FILE_NAME_BYTES, "{last}");
assert!(last.ends_with(&format!("_{}.json", hex_sha256(&file_name))));
assert!(last.starts_with("time-05-06-07-000089_é"));
}
#[test]
fn a_long_prefix_keeps_its_head_and_a_digest_of_all_of_it() {
let path = "p".repeat(2000);
let file_name = logging_id(&start(), "chatcmpl-1");
let key = object_key(&path, "", &start(), &file_name, Granularity::Day);
assert!(key.len() <= MAX_KEY_BYTES, "{}", key.len());
let digest = &hex_sha256(&format!("{path}/"))[..PREFIX_DIGEST_CHARS];
assert!(key.starts_with("ppp"), "{key}");
assert!(key.contains(&format!("/{digest}/2026-03-04/")), "{key}");
}
fn settings(endpoint: Option<&str>, virtual_hosted: bool, region: &str) -> Settings {
Settings {
bucket: "logs".into(),
path: String::new(),
region: region.into(),
endpoint_url: endpoint.map(str::to_owned),
virtual_hosted_style: virtual_hosted,
verify: None,
auth: Default::default(),
team_prefix: false,
key_prefix: false,
strip_base64_files: false,
prompts_only: false,
granularity: Granularity::Day,
server_side_encryption: None,
sse_kms_key_id: None,
max_concurrent_uploads: 1,
max_queue_size: 1,
max_retry_age: None,
drop_on_terminal_error: true,
batch_file_upload: false,
flush_interval: Duration::from_secs(1),
attempt_timeout: Duration::from_secs(1),
log_success: true,
log_failure: true,
turn_off_message_logging: false,
redact_user_api_key_info: false,
user_agent_tags: true,
}
}
#[rstest]
#[case::aws(
None,
false,
"eu-west-1",
"https://logs.s3.eu-west-1.amazonaws.com/a/b%20c%3D%2B.json"
)]
#[case::china(
None,
false,
"cn-north-1",
"https://logs.s3.cn-north-1.amazonaws.com.cn/a/b%20c%3D%2B.json"
)]
#[case::path_style(
Some("http://minio:9000"),
false,
"us-east-1",
"http://minio:9000/logs/a/b%20c%3D%2B.json"
)]
#[case::virtual_hosted(
Some("https://s3.example"),
true,
"us-east-1",
"https://logs.s3.example/a/b%20c%3D%2B.json"
)]
fn urls_are_built_like_the_proxy(
#[case] endpoint: Option<&str>,
#[case] virtual_hosted: bool,
#[case] region: &str,
#[case] expected: &str,
) {
assert_eq!(
object_url(&settings(endpoint, virtual_hosted, region), "a/b c=+.json"),
expected
);
}
#[rstest]
#[case("a/b/c.json", "a/b")]
#[case("c.json", "")]
fn parents(#[case] key: &str, #[case] expected: &str) {
assert_eq!(parent(key), expected);
}
}

506
litellm-rust/crates/gateway/src/s3_log/mod.rs Normal file
View file

@ -0,0 +1,506 @@
//! The `s3_v2` callback: every request's `StandardLoggingPayload` written to S3 as a JSON
//! object, under the proxy's key layout, as the proxy's `S3Logger` writes it.
//!
//! It is an observer, so nothing here runs on the request path. Each finished request is turned
//! into its object and queued; every flush interval the queue is uploaded with bounded
//! concurrency, or written as one NDJSON object per folder with `s3_batch_file_upload`. What
//! fails for now goes back on the queue for the next flush, until it is older than
//! `s3_max_retry_age_seconds` or the queue overflows.
mod config;
mod key;
mod payload;
mod upload;
use std::{
collections::{BTreeMap, VecDeque},
sync::{Arc, Mutex, Weak},
time::{Instant, SystemTime},
};
use axum::body::Bytes;
use futures_util::{StreamExt, future::BoxFuture};
use litellm_config::Config;
use litellm_gateway_hooks::{Facts, Observer, Record};
use litellm_gateway_inference::Gateway;
use serde_json::Value;
use self::{
config::Settings,
key::{Utc, download_name, logging_id, object_key, parent},
payload::{Redaction, prompts_only, standard_logging_payload, strip_base64_files},
upload::{Upload, Uploader},
};
use crate::Error;
const JSON: &str = "application/json";
const NDJSON: &str = "application/x-ndjson";
/// One object waiting to be written.
#[derive(Clone, Debug)]
struct Element {
key: String,
download_name: String,
body: Bytes,
content_type: &'static str,
/// A batch file, already holding many requests; never grouped again.
batch: bool,
/// When this object first failed to upload.
retrying_since: Option<Instant>,
}
#[derive(Default)]
struct Queue {
elements: VecDeque<Element>,
/// How many at the front were put back by the last failed flush.
requeued: usize,
}
/// The `s3_v2` observer.
pub struct S3Log {
settings: Arc<Settings>,
uploader: Uploader,
queue: Mutex<Queue>,
/// Serializes flushes, so the periodic one and the one at shutdown never interleave.
flushing: tokio::sync::Mutex<()>,
}
impl S3Log {
/// Starts the observer and its flush timer, or `Ok(None)` when no callback list names
/// `s3_v2`. Must be called inside a tokio runtime.
pub async fn from_config(
config: &Config,
inference: &Gateway,
) -> Result<Option<Arc<Self>>, Error> {
let env = |name: &str| std::env::var(name).ok();
let Some(settings) =
Settings::from_config(config, inference.secrets.as_ref(), &env).await?
else {
return Ok(None);
};
let settings = Arc::new(settings);
let log = Arc::new(Self {
uploader: Uploader::new(
settings.clone(),
inference.resources.pool.clone(),
&inference.http,
),
settings,
queue: Mutex::default(),
flushing: tokio::sync::Mutex::new(()),
});
tokio::spawn(run(Arc::downgrade(&log)));
Ok(Some(log))
}
fn element(&self, record: &Record) -> Element {
let settings = &self.settings;
let redaction = Redaction {
messages: settings.turn_off_message_logging,
user_api_key_info: settings.redact_user_api_key_info,
user_agent_tags: settings.user_agent_tags,
};
let mut payload = standard_logging_payload(record, redaction);
if settings.strip_base64_files {
strip_base64_files(&mut payload);
}
let metadata = |field: &str| {
payload["metadata"][field]
.as_str()
.filter(|value| !value.is_empty())
.map(str::to_owned)
};
let prefix: String = [
settings
.team_prefix
.then(|| metadata("user_api_key_team_alias")),
settings.key_prefix.then(|| metadata("user_api_key_alias")),
]
.into_iter()
.flatten()
.flatten()
.map(|segment| format!("{segment}/"))
.collect();
let id = payload["id"].as_str().unwrap_or_default().to_owned();
let start = Utc::new(record.request.started_at);
if settings.prompts_only {
prompts_only(&mut payload);
}
Element {
key: object_key(
&settings.path,
&prefix,
&start,
&logging_id(&start, &id),
settings.granularity,
),
download_name: download_name(&start, &id),
body: Bytes::from(Value::Object(payload).to_string()),
content_type: JSON,
batch: false,
retrying_since: None,
}
}
fn enqueue(&self, element: Element) {
let mut queue = self.queue.lock().unwrap_or_else(|e| e.into_inner());
queue.elements.push_back(element);
let excess = queue
.elements
.len()
.saturating_sub(self.settings.max_queue_size);
if excess > 0 {
queue.elements.drain(..excess);
queue.requeued = queue.requeued.saturating_sub(excess);
tracing::warn!(
dropped = excess,
max_queue_size = self.settings.max_queue_size,
"s3 logging: queue full; dropped the oldest events"
);
}
}
/// Uploads everything queued; what fails for now is put back at the front of the queue.
pub async fn flush(&self) {
let _flushing = self.flushing.lock().await;
let (batch, requeued) = {
let mut queue = self.queue.lock().unwrap_or_else(|e| e.into_inner());
let requeued = std::mem::take(&mut queue.requeued);
(Vec::from(std::mem::take(&mut queue.elements)), requeued)
};
if batch.is_empty() {
return;
}
let batch_len = batch.len();
let uploads = match self.settings.batch_file_upload {
true => batch_files(batch),
false => batch,
};
// What failed last time goes last, so a backlog never starves new requests.
let stale = match uploads.len() == batch_len {
true => requeued.min(uploads.len()),
false => 0,
};
let order: Vec<usize> = (stale..uploads.len()).chain(0..stale).collect();
let mut outcomes = vec![Upload::Retry; uploads.len()];
let mut finished = futures_util::stream::iter(order)
.map(|index| {
let uploads = &uploads;
async move { (index, self.uploader.upload(&uploads[index]).await) }
})
.buffer_unordered(self.settings.max_concurrent_uploads);
while let Some((index, outcome)) = finished.next().await {
outcomes[index] = outcome;
}
drop(finished);
let delivered = outcomes
.iter()
.filter(|outcome| **outcome == Upload::Delivered)
.count();
let total = uploads.len();
let failed = total - delivered;
if failed == 0 {
return;
}
let now = Instant::now();
let max_age = self.settings.max_retry_age;
let failures = uploads
.into_iter()
.zip(outcomes)
.filter(|(_, outcome)| *outcome != Upload::Delivered);
// When nothing got through the bucket itself is failing: keep everything, unaged.
let kept: Vec<Element> = match delivered {
0 => failures.map(|(element, _)| element).collect(),
_ => failures
.filter(|(element, outcome)| {
*outcome == Upload::Retry
&& !matches!(
(max_age, element.retrying_since),
(Some(max_age), Some(since)) if now.duration_since(since) > max_age
)
})
.map(|(element, _)| Element {
retrying_since: element.retrying_since.or(max_age.map(|_| now)),
..element
})
.collect(),
};
tracing::warn!(
failed,
kept = kept.len(),
total,
"s3 logging: uploads failed this flush; the kept ones are retried on the next, the rest \
were rejected by S3 or retried longer than s3_max_retry_age_seconds"
);
if kept.is_empty() {
return;
}
let mut queue = self.queue.lock().unwrap_or_else(|e| e.into_inner());
let arrivals = std::mem::take(&mut queue.elements);
let kept_len = kept.len();
let mut elements: VecDeque<Element> = kept.into_iter().chain(arrivals).collect();
let overflow = elements.len().saturating_sub(self.settings.max_queue_size);
if overflow > 0 {
elements.drain(..overflow);
tracing::warn!(
dropped = overflow,
max_queue_size = self.settings.max_queue_size,
"s3 logging: queue exceeded max_queue_size after a failed flush; dropped the oldest events"
);
}
queue.elements = elements;
queue.requeued = kept_len.saturating_sub(overflow);
}
}
/// `s3_batch_file_upload`: one NDJSON object per folder, named for the flush.
fn batch_files(batch: Vec<Element>) -> Vec<Element> {
let now = Utc::new(SystemTime::now());
let (ready, singles): (Vec<Element>, Vec<Element>) =
batch.into_iter().partition(|element| element.batch);
let mut folders: BTreeMap<String, Vec<Element>> = BTreeMap::new();
for element in singles {
folders
.entry(parent(&element.key).to_owned())
.or_default()
.push(element);
}
let files = folders.into_iter().map(|(folder, elements)| {
let name = format!("batch_{}_{}", now.clock(), uuid::Uuid::new_v4().simple());
let lines: Vec<&[u8]> = elements.iter().map(|element| &element.body[..]).collect();
Element {
key: match folder.is_empty() {
true => format!("{name}.jsonl"),
false => format!("{folder}/{name}.jsonl"),
},
download_name: format!("{name}.jsonl"),
body: Bytes::from(lines.join(&b'\n')),
content_type: NDJSON,
batch: true,
retrying_since: elements
.iter()
.filter_map(|element| element.retrying_since)
.min(),
}
});
ready.into_iter().chain(files).collect()
}
impl Observer for S3Log {
fn name(&self) -> &'static str {
"s3_v2"
}
fn reads(&self) -> Facts {
Facts::KEY
| Facts::END_USER
| Facts::CLIENT_HEADERS
| Facts::DEPLOYMENT
| Facts::USAGE
| Facts::COST
| Facts::TIMINGS
| Facts::ANNOTATIONS
| Facts::REQUEST_BODY
| Facts::RESPONSE_BODY
}
fn observe<'a>(&'a self, record: &'a Record) -> BoxFuture<'a, ()> {
let logged = match record.succeeded() {
true => self.settings.log_success,
false => self.settings.log_failure,
};
if logged {
self.enqueue(self.element(record));
}
Box::pin(async {})
}
fn flush(&self) -> BoxFuture<'_, ()> {
Box::pin(S3Log::flush(self))
}
}
/// Flushes every flush interval until the observer is dropped.
async fn run(log: Weak<S3Log>) {
let interval = match log.upgrade() {
Some(log) => log.settings.flush_interval,
None => return,
};
loop {
tokio::time::sleep(interval).await;
let Some(log) = log.upgrade() else {
return;
};
log.flush().await;
}
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use litellm_gateway_hooks::Route;
use serde_json::json;
use super::*;
fn element(key: &str, body: &str) -> Element {
Element {
key: key.into(),
download_name: "d.json".into(),
body: Bytes::from(body.to_owned()),
content_type: JSON,
batch: false,
retrying_since: None,
}
}
#[test]
fn batch_files_hold_one_line_per_request_per_folder() {
let earlier = Instant::now();
let mut retried = element("logs/2026-03-04/b.json", "{\"b\":1}");
retried.retrying_since = Some(earlier);
let files = batch_files(vec![
element("logs/2026-03-04/a.json", "{\"a\":1}"),
element("logs/team/2026-03-04/c.json", "{\"c\":1}"),
retried,
]);
assert_eq!(files.len(), 2);
let day = &files[0];
let (folder, name) = day.key.rsplit_once('/').unwrap();
assert_eq!(folder, "logs/2026-03-04");
assert!(
name.starts_with("batch_") && name.ends_with(".jsonl"),
"{name}"
);
assert_eq!(day.download_name, name);
assert_eq!(&day.body[..], b"{\"a\":1}\n{\"b\":1}");
assert_eq!(day.content_type, NDJSON);
assert_eq!(day.retrying_since, Some(earlier));
assert!(files[1].key.starts_with("logs/team/2026-03-04/batch_"));
// A batch file kept for a retry is sent again as it is.
let again = batch_files(files.clone());
assert_eq!(
again.iter().map(|file| &file.key).collect::<Vec<_>>(),
files.iter().map(|file| &file.key).collect::<Vec<_>>()
);
}
async fn log(params: &str) -> S3Log {
let settings = config::tests::settings_with(
&format!(
"litellm_settings:\n callbacks: [s3_v2]\n s3_callback_params:\n s3_bucket_name: logs\n s3_path: audit\n{params}"
),
&[],
)
.await
.unwrap()
.unwrap();
let settings = Arc::new(settings);
S3Log {
uploader: Uploader::new(
settings.clone(),
Arc::new(litellm_http::HttpClientPool::new(Arc::new(
litellm_http::media::PublicDnsResolver,
))),
&litellm_http::Resolution::from(&litellm_http::HttpSettings::default()).config,
),
settings,
queue: Mutex::default(),
flushing: tokio::sync::Mutex::new(()),
}
}
fn record() -> Record {
payload::tests::record(
Route::ChatCompletions,
200,
false,
json!({"messages": [{"role": "user", "content": "hi"}]}),
Some(payload::tests::COMPLETION),
)
}
#[rstest::rstest]
#[case::plain("", "audit/1970-01-01/time-00-16-40-000000_chatcmpl-1.json")]
#[case::team(
" s3_use_team_prefix: true\n",
"audit/research/1970-01-01/time-00-16-40-000000_chatcmpl-1.json"
)]
#[case::team_and_key(
" s3_use_team_prefix: true\n s3_use_key_prefix: true\n s3_partition_granularity: hour\n",
"audit/research/team-key/1970-01-01/00/time-00-16-40-000000_chatcmpl-1.json"
)]
#[tokio::test]
async fn objects_are_keyed_by_team_key_and_time(#[case] params: &str, #[case] expected: &str) {
let element = log(params).await.element(&record());
assert_eq!(element.key, expected);
assert_eq!(
element.download_name,
"time-1970-01-01T00-16-40-000000_chatcmpl-1.json"
);
}
#[tokio::test]
async fn prompts_only_objects_have_no_response() {
let element = log(" s3_log_prompts_only: true\n")
.await
.element(&record());
let body: Value = serde_json::from_slice(&element.body).unwrap();
assert_eq!(body["response"], Value::Null);
assert_eq!(body["messages"], json!([{"role": "user", "content": "hi"}]));
assert!(element.key.ends_with("_chatcmpl-1.json"));
}
#[tokio::test]
async fn only_the_listed_outcomes_are_queued() {
let settings = config::tests::settings_with(
"litellm_settings:\n failure_callback: [s3_v2]\n s3_callback_params:\n s3_bucket_name: logs\n",
&[],
)
.await
.unwrap()
.unwrap();
let mut log = log("").await;
log.settings = Arc::new(settings);
log.observe(&record()).await;
assert!(log.queue.lock().unwrap().elements.is_empty());
let mut failed = record();
failed.outcome.status = 500;
log.observe(&failed).await;
assert_eq!(log.queue.lock().unwrap().elements.len(), 1);
}
#[tokio::test]
async fn a_full_queue_drops_the_oldest() {
let log = log(" s3_max_queue_size: 2\n").await;
for key in ["a", "b", "c"] {
log.enqueue(element(key, "{}"));
}
let queue = log.queue.lock().unwrap();
let keys: Vec<_> = queue
.elements
.iter()
.map(|element| element.key.as_str())
.collect();
assert_eq!(keys, ["b", "c"]);
}
#[tokio::test]
async fn an_unreachable_bucket_keeps_everything_for_the_next_flush() {
let log = log(" s3_endpoint_url: http://127.0.0.1:9\n s3_aws_access_key_id: k\n s3_aws_secret_access_key: s\n s3_max_retry_age_seconds: 1\n").await;
let mut old = element("a", "{}");
old.retrying_since = Some(Instant::now() - Duration::from_secs(60));
log.enqueue(old);
log.enqueue(element("b", "{}"));
log.flush().await;
let queue = log.queue.lock().unwrap();
let kept: Vec<_> = queue
.elements
.iter()
.map(|element| element.key.as_str())
.collect();
assert_eq!(kept, ["a", "b"]);
assert_eq!(queue.requeued, 2);
}
}

1204
litellm-rust/crates/gateway/src/s3_log/payload.rs Normal file
View file

File diff suppressed because it is too large Load diff

278
litellm-rust/crates/gateway/src/s3_log/upload.rs Normal file
View file

@ -0,0 +1,278 @@
//! One signed PUT per object, retried as the proxy's `_put_with_retries` retries it: three
//! attempts, a transport error or a 403, 429, 500 or 503 retried after a jittered backoff, and
//! a 400 or 403 that names an object-specific error code dropped instead of kept for later.
use std::{
collections::BTreeMap,
sync::Arc,
time::{Duration, Instant, SystemTime},
};
use base64::Engine;
use litellm_auth_aws::{AwsAuthService, Credentials, sign_s3};
use litellm_http::{ClientVariant, HttpClientConfig, HttpClientPool, Verify};
use md5::Md5;
use sha2::{Digest, Sha256};
use super::{Element, config::Settings, key::object_url};
const ATTEMPTS: u32 = 3;
const RETRYABLE: [u16; 4] = [403, 429, 500, 503];
/// Codes S3 answers that mean this object can never be stored, so retrying it is pointless.
const TERMINAL_CODES: [&str; 7] = [
"EntityTooLarge",
"InvalidArgument",
"MalformedXML",
"InvalidDigest",
"KeyTooLongError",
"BadDigest",
"InvalidRequest",
];
/// Credentials without an expiry are resolved again after this long.
const CREDENTIALS_TTL: Duration = Duration::from_secs(600);
/// Credentials are refreshed this long before they expire.
const CREDENTIALS_MARGIN: Duration = Duration::from_secs(300);
/// How one object's upload ended.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(super) enum Upload {
Delivered,
/// Failed for now; it goes back on the queue.
Retry,
/// S3 rejected the object itself.
Dropped,
}
/// What one attempt got back.
enum Attempt {
Response { status: u16, body: String },
Transport(String),
}
impl Attempt {
fn is_terminal(&self) -> bool {
matches!(self, Self::Response { status, body } if is_terminal(*status, body))
}
}
/// A 400 or 403 whose `<Code>` S3 reserves for a rejection of this object.
fn is_terminal(status: u16, body: &str) -> bool {
matches!(status, 400 | 403)
&& error_code(body).is_some_and(|code| TERMINAL_CODES.contains(&code))
}
fn error_code(body: &str) -> Option<&str> {
let start = body.find("<Code>")? + "<Code>".len();
let length = body[start..].find("</Code>")?;
Some(&body[start..start + length]).filter(|code| !code.is_empty() && !code.contains('<'))
}
/// `(1 << attempt) * (0.5 + random())` seconds.
fn backoff(attempt: u32) -> Duration {
Duration::from_secs_f64(f64::from(1u32 << attempt) * (0.5 + rand::random::<f64>()))
}
pub(super) struct Uploader {
settings: Arc<Settings>,
pool: Arc<HttpClientPool>,
http: HttpClientConfig,
auth: AwsAuthService,
credentials: tokio::sync::Mutex<Option<(Credentials, Instant)>>,
}
impl Uploader {
pub fn new(
settings: Arc<Settings>,
pool: Arc<HttpClientPool>,
http: &HttpClientConfig,
) -> Self {
let verify = match &settings.verify {
None => http.verify.clone(),
Some(litellm_http::SslVerify::Enabled) => Verify::BuiltInRoots,
Some(litellm_http::SslVerify::Disabled) => Verify::Disabled,
Some(litellm_http::SslVerify::CaBundle(path)) => Verify::CaBundle(path.clone()),
};
Self {
settings,
pool,
http: HttpClientConfig {
verify,
..http.clone()
},
auth: AwsAuthService::default(),
credentials: tokio::sync::Mutex::new(None),
}
}
pub async fn upload(&self, element: &Element) -> Upload {
let url = object_url(&self.settings, &element.key);
let headers = self.headers(element);
let mut attempt = 0;
loop {
let result = match self.put(&url, element, &headers).await {
Ok(result) => result,
Err(error) => {
tracing::error!(%error, key = %element.key, "s3 logging: cannot sign the upload");
return Upload::Retry;
}
};
let retry = match &result {
Attempt::Transport(_) => true,
Attempt::Response { status, .. } => {
RETRYABLE.contains(status)
&& !(self.settings.drop_on_terminal_error && result.is_terminal())
}
};
if retry && attempt + 1 < ATTEMPTS {
let wait = backoff(attempt);
let reason = match &result {
Attempt::Transport(error) => error.clone(),
Attempt::Response { status, .. } => status.to_string(),
};
tracing::debug!(attempt = attempt + 1, %reason, ?wait, key = %element.key, "s3 logging: upload failed, retrying");
tokio::time::sleep(wait).await;
attempt += 1;
continue;
}
return match result {
Attempt::Response { status, .. } if (200..300).contains(&status) => {
Upload::Delivered
}
Attempt::Response { status, body } => {
tracing::error!(status, body = %body.chars().take(512).collect::<String>(), key = %element.key, "s3 logging: upload failed");
if self.settings.drop_on_terminal_error && is_terminal(status, &body) {
tracing::warn!(status, key = %element.key, "s3 logging: dropping the object, S3 rejected it");
return Upload::Dropped;
}
Upload::Retry
}
Attempt::Transport(error) => {
tracing::error!(%error, key = %element.key, "s3 logging: upload failed");
Upload::Retry
}
};
}
}
/// The proxy's PUT headers, before signing.
fn headers(&self, element: &Element) -> BTreeMap<String, String> {
let body = &element.body;
let sha256: String = Sha256::digest(body)
.iter()
.map(|byte| format!("{byte:02x}"))
.collect();
let md5 =
base64::engine::general_purpose::STANDARD.encode(<Md5 as md5::Digest>::digest(body));
let mut headers = BTreeMap::from([
("Content-Type".to_owned(), element.content_type.to_owned()),
("Content-MD5".to_owned(), md5),
("x-amz-content-sha256".to_owned(), sha256),
("Content-Language".to_owned(), "en".to_owned()),
(
"Content-Disposition".to_owned(),
format!("inline; filename=\"{}\"", element.download_name),
),
(
"Cache-Control".to_owned(),
"private, immutable, max-age=31536000, s-maxage=0".to_owned(),
),
]);
if let Some(algorithm) = &self.settings.server_side_encryption {
headers.insert("x-amz-server-side-encryption".into(), algorithm.clone());
}
if let Some(key_id) = &self.settings.sse_kms_key_id {
headers.insert(
"x-amz-server-side-encryption-aws-kms-key-id".into(),
key_id.clone(),
);
}
headers
}
async fn put(
&self,
url: &str,
element: &Element,
headers: &BTreeMap<String, String>,
) -> Result<Attempt, Box<dyn std::error::Error + Send + Sync>> {
let credentials = self.credentials().await?;
let signature = sign_s3(
"PUT",
url,
&element.body,
headers,
&self.settings.region,
&credentials,
SystemTime::now(),
)?;
let client = self.pool.client(&self.http, ClientVariant::Provider)?;
let mut request = client
.put(url)
.timeout(self.settings.attempt_timeout)
.body(element.body.clone());
for (name, value) in headers.iter().chain(&signature) {
request = request.header(name, value);
}
Ok(match request.send().await {
Ok(response) => Attempt::Response {
status: response.status().as_u16(),
body: response.text().await.unwrap_or_default(),
},
Err(error) => Attempt::Transport(error.to_string()),
})
}
async fn credentials(&self) -> Result<Credentials, litellm_auth_aws::Error> {
let mut cached = self.credentials.lock().await;
if let Some((credentials, until)) = cached.as_ref()
&& Instant::now() < *until
{
return Ok(credentials.clone());
}
let env = |name: &str| std::env::var(name).ok();
let credentials = self
.auth
.resolve_credentials(self.settings.auth.clone(), &env)
.await?;
let until = credentials
.expiry()
.and_then(|expiry| expiry.duration_since(SystemTime::now()).ok())
.map_or(CREDENTIALS_TTL, |left| {
left.saturating_sub(CREDENTIALS_MARGIN)
});
*cached = Some((credentials.clone(), Instant::now() + until));
Ok(credentials)
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
#[rstest]
#[case::too_large(400, "<Error><Code>EntityTooLarge</Code></Error>", true)]
#[case::bad_digest(403, "<Error><Code>BadDigest</Code></Error>", true)]
#[case::access_denied(403, "<Error><Code>AccessDenied</Code></Error>", false)]
#[case::slow_down(503, "<Error><Code>InvalidRequest</Code></Error>", false)]
#[case::not_xml(400, "bad request", false)]
#[case::empty_code(400, "<Code></Code>", false)]
fn only_object_rejections_are_terminal(
#[case] status: u16,
#[case] body: &str,
#[case] expected: bool,
) {
assert_eq!(is_terminal(status, body), expected);
}
#[rstest]
#[case(0, 0.5, 1.5)]
#[case(1, 1.0, 3.0)]
fn backoff_doubles_with_jitter(#[case] attempt: u32, #[case] low: f64, #[case] high: f64) {
for _ in 0..50 {
let wait = backoff(attempt).as_secs_f64();
assert!((low..high).contains(&wait), "{wait}");
}
}
}

6
litellm-rust/crates/gateway/src/spend.rs
View file

@ -446,7 +446,7 @@ fn status(record: &Record) -> &'static str {
}
/// The model as the proxy logs it: the deployment's model without its provider route.
fn logged_model(record: &Record) -> String {
pub(crate) fn logged_model(record: &Record) -> String {
match &record.outcome.deployment {
Some(deployment) => match (deployment.provider(), deployment.model.split_once('/')) {
(Some(provider), Some((prefix, rest))) if prefix == provider => rest.to_owned(),
@ -597,7 +597,7 @@ fn metadata(record: &Record) -> Value {
/// `spend_logs_metadata` as the proxy merges it: the team's, then the key's, then what hooks
/// attached for the request, later keys winning.
fn spend_logs_metadata(record: &Record) -> Value {
pub(crate) fn spend_logs_metadata(record: &Record) -> Value {
let mut merged = Map::new();
if let Some(key) = record.key() {
for metadata in [&key.team_metadata, &key.metadata] {
@ -617,7 +617,7 @@ fn spend_logs_metadata(record: &Record) -> Value {
/// The proxy's usage object for a response, and its `additional_usage_values`: everything but
/// the three totals.
fn usage_values(usage: &Usage) -> (Value, Value) {
pub(crate) fn usage_values(usage: &Usage) -> (Value, Value) {
let mut additional = Map::new();
additional.insert(
"prompt_tokens_details".into(),

293
litellm-rust/crates/gateway/tests/azure_chat.rs Normal file
View file

@ -0,0 +1,293 @@
//! Azure OpenAI chat completions through the gateway: the deployment URL, `api-version` and
//! `api-key` the upstream sees, and the usage and cost the request is recorded with.
use std::{
io::Write,
sync::{Arc, Mutex, Once},
};
use axum::{
Router,
body::{Body, to_bytes},
http::{Request, header::CONTENT_TYPE},
response::IntoResponse,
};
use futures_util::future::BoxFuture;
use litellm_config::Config;
use litellm_gateway_hooks::{Facts, Observer, Record};
use rstest::rstest;
use serde_json::{Value, json};
use tokio::net::TcpListener;
use tower::ServiceExt;
const RESPONSE_MODEL: &str = "gpt-4o-2024-11-20";
const COMPLETION: &str = r#"{"id":"chatcmpl-1","object":"chat.completion","created":1,"model":"gpt-4o-2024-11-20","choices":[{"index":0,"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":5,"completion_tokens":6,"total_tokens":11}}"#;
const COMPLETION_STREAM: &str = concat!(
"data: {\"id\":\"chatcmpl-1\",\"object\":\"chat.completion.chunk\",\"model\":\"gpt-4o-2024-11-20\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hi\"}}],\"usage\":null}\n\n",
"data: {\"id\":\"chatcmpl-1\",\"object\":\"chat.completion.chunk\",\"model\":\"gpt-4o-2024-11-20\",\"choices\":[],\"usage\":{\"prompt_tokens\":7,\"completion_tokens\":3,\"total_tokens\":10}}\n\n",
"data: [DONE]\n\n",
);
/// Prices the cost map gives the model a response names, and the `base_model` one deployment
/// declares, so the two pricing routes give different costs.
const RESPONSE_MODEL_RATES: (f64, f64) = (1.0, 2.0);
const BASE_MODEL_RATES: (f64, f64) = (10.0, 20.0);
/// The proxy reads its cost map from this variable, and every test sets it before building a
/// gateway, so no test reads the environment while it changes.
fn use_local_cost_map() {
static ONCE: Once = Once::new();
ONCE.call_once(|| {
let map = json!({
format!("azure/{RESPONSE_MODEL}"): {
"litellm_provider": "azure",
"mode": "chat",
"input_cost_per_token": RESPONSE_MODEL_RATES.0,
"output_cost_per_token": RESPONSE_MODEL_RATES.1,
},
"azure/gpt-4o": {
"litellm_provider": "azure",
"mode": "chat",
"input_cost_per_token": BASE_MODEL_RATES.0,
"output_cost_per_token": BASE_MODEL_RATES.1,
},
});
let mut file = tempfile::NamedTempFile::new().unwrap();
file.write_all(map.to_string().as_bytes()).unwrap();
let (_, path) = file.keep().unwrap();
// SAFETY: set once, before any test builds a gateway or spawns a thread that reads the
// environment.
unsafe { std::env::set_var("LITELLM_MODEL_COST_MAP_PATH", path) };
});
}
#[derive(Clone, Debug)]
struct Sent {
path: String,
query: Option<String>,
api_key: Option<String>,
authorization: Option<String>,
body: Value,
}
type Calls = Arc<Mutex<Vec<Sent>>>;
async fn upstream(calls: Calls) -> String {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let app = Router::new().fallback(move |request: Request<Body>| {
let calls = calls.clone();
async move {
let (parts, body) = request.into_parts();
let header = |name: &str| {
parts
.headers
.get(name)
.and_then(|value| value.to_str().ok())
.map(str::to_owned)
};
let (api_key, authorization) = (header("api-key"), header("authorization"));
let path = parts.uri.path().to_owned();
let query = parts.uri.query().map(str::to_owned);
let body: Value =
serde_json::from_slice(&to_bytes(body, 1 << 20).await.unwrap()).unwrap();
let stream = body["stream"].as_bool().unwrap_or(false);
calls.lock().unwrap().push(Sent {
path,
query,
api_key,
authorization,
body,
});
match stream {
true => ([(CONTENT_TYPE, "text/event-stream")], COMPLETION_STREAM).into_response(),
false => ([(CONTENT_TYPE, "application/json")], COMPLETION).into_response(),
}
}
});
tokio::spawn(async move { axum::serve(listener, app).await });
format!("http://{address}")
}
#[derive(Debug)]
struct Seen {
deployment: Option<String>,
usage: Option<(u64, u64)>,
cost: Option<f64>,
}
struct Collect(Mutex<Vec<Seen>>);
impl Observer for Collect {
fn name(&self) -> &'static str {
"collect"
}
fn reads(&self) -> Facts {
Facts::DEPLOYMENT | Facts::USAGE | Facts::COST
}
fn observe<'a>(&'a self, record: &'a Record) -> BoxFuture<'a, ()> {
Box::pin(async move {
self.0.lock().unwrap().push(Seen {
deployment: record
.outcome
.deployment
.as_ref()
.map(|deployment| deployment.id.to_string()),
usage: record
.outcome
.usage
.map(|usage| (usage.prompt_tokens(), usage.output_tokens)),
cost: record.outcome.cost,
});
})
}
}
struct Harness {
app: Router,
inference: Arc<litellm_gateway_inference::Gateway>,
calls: Calls,
seen: Arc<Collect>,
}
async fn harness() -> Harness {
use_local_cost_map();
let calls = Calls::default();
let upstream = upstream(calls.clone()).await;
let config = Config::from_yaml(&format!(
"model_list:
- model_name: chat
litellm_params:
model: azure/chat-prod
api_base: {upstream}
api_key: sk-azure
api_version: 2024-10-21
model_info:
id: azure-1
- model_name: chat-base
litellm_params:
model: azure/chat-base
api_base: {upstream}/
api_key: sk-azure
api_version: 2024-10-21
base_model: gpt-4o
model_info:
id: azure-2
router_settings:
num_retries: 0
general_settings:
master_key: sk-master
"
))
.unwrap();
let inference = litellm_gateway::build_inference(&config).unwrap();
let seen = Arc::new(Collect(Mutex::new(Vec::new())));
litellm_gateway::hooks::install(&inference, &config, Vec::new(), vec![seen.clone()])
.await
.unwrap();
let app = litellm_gateway::router(inference.clone(), &config, None, None);
Harness {
app,
inference,
calls,
seen,
}
}
async fn send(app: &Router, body: Value) -> (u16, String) {
let request = Request::post("/v1/chat/completions")
.header("content-type", "application/json")
.header("authorization", "Bearer sk-master")
.body(Body::from(body.to_string()))
.unwrap();
let response = app.clone().oneshot(request).await.unwrap();
let status = response.status().as_u16();
let body = to_bytes(response.into_body(), 1 << 20).await.unwrap();
(status, String::from_utf8(body.to_vec()).unwrap())
}
fn close(left: f64, right: f64) -> bool {
(left - right).abs() < 1e-9
}
fn cost((input, output): (u64, u64), (input_rate, output_rate): (f64, f64)) -> f64 {
input as f64 * input_rate + output as f64 * output_rate
}
#[rstest]
#[case::unary_by_response_model("chat", false, "chat-prod", "azure-1", (5, 6), RESPONSE_MODEL_RATES)]
#[case::stream_by_response_model("chat", true, "chat-prod", "azure-1", (7, 3), RESPONSE_MODEL_RATES)]
#[case::unary_by_base_model("chat-base", false, "chat-base", "azure-2", (5, 6), BASE_MODEL_RATES)]
#[case::stream_by_base_model("chat-base", true, "chat-base", "azure-2", (7, 3), BASE_MODEL_RATES)]
#[tokio::test]
async fn azure_chat_reaches_the_deployment_and_is_priced(
#[case] model: &str,
#[case] stream: bool,
#[case] deployment: &str,
#[case] id: &str,
#[case] usage: (u64, u64),
#[case] rates: (f64, f64),
) {
let harness = harness().await;
let mut request = json!({
"model": model,
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 8,
"user": "end-user-1",
});
if stream {
request["stream"] = json!(true);
request["stream_options"] = json!({"include_usage": true});
}
let (status, body) = send(&harness.app, request).await;
assert_eq!(status, 200, "{body}");
match stream {
true => assert!(body.contains("\"content\":\"hi\""), "{body}"),
false => assert_eq!(
serde_json::from_str::<Value>(&body).unwrap()["model"],
RESPONSE_MODEL
),
}
let sent = harness.calls.lock().unwrap()[0].clone();
assert_eq!(
sent.path,
format!("/openai/deployments/{deployment}/chat/completions")
);
assert_eq!(sent.query.as_deref(), Some("api-version=2024-10-21"));
assert_eq!(sent.api_key.as_deref(), Some("sk-azure"));
assert_eq!(sent.authorization, None);
assert_eq!(sent.body["model"], deployment);
assert_eq!(sent.body["max_tokens"], 8);
assert_eq!(sent.body["user"], "end-user-1");
assert_eq!(sent.body["stream"].as_bool().unwrap_or(false), stream);
litellm_gateway::hooks::close(&harness.inference).await;
let seen = harness.seen.0.lock().unwrap();
let [seen] = seen.as_slice() else {
panic!("{seen:?}")
};
assert_eq!(seen.deployment.as_deref(), Some(id));
assert_eq!(seen.usage, Some(usage));
assert!(close(seen.cost.unwrap(), cost(usage, rates)), "{seen:?}");
}
#[rstest]
#[tokio::test]
async fn an_unsupported_param_is_refused_before_the_call() {
let harness = harness().await;
let (status, body) = send(
&harness.app,
json!({
"model": "chat",
"messages": [{"role": "user", "content": "hi"}],
"reasoning_effort": "high",
}),
)
.await;
assert_eq!(status, 400, "{body}");
assert!(body.contains("reasoning_effort"), "{body}");
assert!(harness.calls.lock().unwrap().is_empty());
}

508
litellm-rust/crates/gateway/tests/chat_streaming.rs Normal file
View file

@ -0,0 +1,508 @@
//! `stream: true` chat completions for deployments chat translates (Anthropic, Bedrock
//! Converse): the OpenAI chunks the client gets, closed by `[DONE]`, and the usage and cost
//! the request is recorded with, whether or not the client asked to see the usage.
use std::sync::{Arc, Mutex};
use axum::{
Router,
body::{Body, Bytes, to_bytes},
http::{Request, StatusCode, header::CONTENT_TYPE},
response::IntoResponse,
};
use futures_util::future::BoxFuture;
use litellm_config::Config;
use litellm_framer::aws_event_stream::{AwsEventStreamCodec, Header, HeaderValue, Message};
use litellm_gateway_hooks::{Facts, Observer, Record};
use rstest::rstest;
use serde_json::{Value, json};
use tokio::net::TcpListener;
use tokio_util::codec::Encoder;
use tower::ServiceExt;
/// Input, output, cache read and cache write prices of every deployment.
const RATES: (f64, f64, f64, f64) = (1.0, 2.0, 0.1, 1.25);
fn anthropic_events(events: &[Value]) -> String {
events
.iter()
.map(|event| {
format!(
"event: {}\ndata: {event}\n\n",
event["type"].as_str().unwrap()
)
})
.collect()
}
fn anthropic_stream() -> String {
anthropic_events(&[
json!({"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": "claude-x-20250101", "content": [], "stop_reason": null, "stop_sequence": null,
"usage": {"input_tokens": 10, "cache_creation_input_tokens": 6, "cache_read_input_tokens": 20, "output_tokens": 1}}}),
json!({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
json!({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Hello"}}),
json!({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": " world"}}),
json!({"type": "content_block_stop", "index": 0}),
json!({"type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence": null}, "usage": {"output_tokens": 12}}),
json!({"type": "message_stop"}),
])
}
fn bedrock_frame(kind: &str, name: &str, payload: Value, wire: &mut bytes::BytesMut) {
let type_header = match kind {
"exception" => ":exception-type",
_ => ":event-type",
};
let message = Message::new(Bytes::from(payload.to_string()))
.add_header(Header::new(
":message-type",
HeaderValue::String(kind.to_owned().into()),
))
.add_header(Header::new(
type_header,
HeaderValue::String(name.to_owned().into()),
));
AwsEventStreamCodec.encode(message, wire).unwrap();
}
fn bedrock_stream() -> Bytes {
let mut wire = bytes::BytesMut::new();
for (name, payload) in [
("messageStart", json!({"role": "assistant"})),
(
"contentBlockDelta",
json!({"contentBlockIndex": 0, "delta": {"text": "Hello"}}),
),
(
"contentBlockDelta",
json!({"contentBlockIndex": 0, "delta": {"text": " world"}}),
),
("contentBlockStop", json!({"contentBlockIndex": 0})),
("messageStop", json!({"stopReason": "end_turn"})),
(
"metadata",
json!({"usage": {"inputTokens": 10, "outputTokens": 12, "totalTokens": 48, "cacheReadInputTokens": 20, "cacheWriteInputTokens": 6}, "metrics": {"latencyMs": 9}}),
),
] {
bedrock_frame("event", name, payload, &mut wire);
}
wire.freeze()
}
fn bedrock_cut_stream() -> Bytes {
let mut wire = bytes::BytesMut::new();
bedrock_frame(
"event",
"messageStart",
json!({"role": "assistant"}),
&mut wire,
);
bedrock_frame(
"event",
"contentBlockDelta",
json!({"contentBlockIndex": 0, "delta": {"text": "Hel"}}),
&mut wire,
);
bedrock_frame(
"exception",
"throttlingException",
json!({"message": "Too many requests"}),
&mut wire,
);
wire.freeze()
}
#[derive(Clone, Debug)]
struct Sent {
path: String,
body: Value,
}
type Calls = Arc<Mutex<Vec<Sent>>>;
/// Answers by credential: each deployment's key picks what its provider does.
async fn upstream(calls: Calls) -> String {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let app = Router::new().fallback(move |request: Request<Body>| {
let calls = calls.clone();
async move {
let (parts, body) = request.into_parts();
let credential = parts
.headers
.get("x-api-key")
.or_else(|| parts.headers.get("authorization"))
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_owned();
let body: Value =
serde_json::from_slice(&to_bytes(body, 1 << 20).await.unwrap()).unwrap();
calls.lock().unwrap().push(Sent {
path: parts.uri.path().to_owned(),
body,
});
let sse = [(CONTENT_TYPE, "text/event-stream")];
let event_stream = [(CONTENT_TYPE, "application/vnd.amazon.eventstream")];
match credential.as_str() {
"sk-anthropic" => (sse, anthropic_stream()).into_response(),
"sk-anthropic-error-first" => (
sse,
anthropic_events(&[
json!({"type": "message_start", "message": {"id": "m", "type": "message", "role": "assistant", "model": "claude-x", "content": [], "stop_reason": null, "stop_sequence": null, "usage": {"input_tokens": 1}}}),
json!({"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}),
]),
)
.into_response(),
"sk-anthropic-overloaded" => (
StatusCode::from_u16(529).unwrap(),
[(CONTENT_TYPE, "application/json")],
r#"{"type":"error","error":{"type":"overloaded_error","message":"Overloaded"}}"#,
)
.into_response(),
"Bearer bedrock-token" => (event_stream, bedrock_stream()).into_response(),
"Bearer bedrock-cut" => (event_stream, bedrock_cut_stream()).into_response(),
other => (StatusCode::UNAUTHORIZED, format!("unexpected credential {other}"))
.into_response(),
}
}
});
tokio::spawn(async move { axum::serve(listener, app).await });
format!("http://{address}")
}
#[derive(Debug)]
struct Seen {
deployment: Option<String>,
status: u16,
usage: Option<litellm_gateway_hooks::Usage>,
cost: Option<f64>,
first_byte: bool,
}
struct Collect(Mutex<Vec<Seen>>);
impl Observer for Collect {
fn name(&self) -> &'static str {
"collect"
}
fn reads(&self) -> Facts {
Facts::DEPLOYMENT | Facts::USAGE | Facts::COST
}
fn observe<'a>(&'a self, record: &'a Record) -> BoxFuture<'a, ()> {
Box::pin(async move {
self.0.lock().unwrap().push(Seen {
deployment: record
.outcome
.deployment
.as_ref()
.map(|deployment| deployment.id.to_string()),
status: record.outcome.status,
usage: record.outcome.usage,
cost: record.outcome.cost,
first_byte: record.outcome.time_to_first_byte.is_some(),
});
})
}
}
struct Harness {
app: Router,
inference: Arc<litellm_gateway_inference::Gateway>,
calls: Calls,
seen: Arc<Collect>,
}
fn deployment(name: &str, id: &str, model: &str, upstream: &str, key: &str) -> String {
let (input, output, cache_read, cache_write) = RATES;
format!(
" - model_name: {name}
litellm_params:
model: {model}
api_base: {upstream}
api_key: {key}
input_cost_per_token: {input}
output_cost_per_token: {output}
cache_read_input_token_cost: {cache_read}
cache_creation_input_token_cost: {cache_write}
model_info:
id: {id}
"
)
}
async fn harness() -> Harness {
let calls = Calls::default();
let upstream = upstream(calls.clone()).await;
let bedrock = "bedrock/anthropic.claude-x-v1:0";
let models = [
deployment(
"claude",
"claude-1",
"anthropic/claude-x",
&upstream,
"sk-anthropic",
),
deployment("bedrock", "bedrock-1", bedrock, &upstream, "bedrock-token"),
deployment(
"bedrock-cut",
"bedrock-2",
bedrock,
&upstream,
"bedrock-cut",
),
deployment(
"overloaded",
"claude-2",
"anthropic/claude-x",
&upstream,
"sk-anthropic-overloaded",
),
deployment(
"fallback",
"claude-3",
"anthropic/claude-x",
&upstream,
"sk-anthropic-error-first",
),
deployment("fallback", "bedrock-3", bedrock, &upstream, "bedrock-token"),
]
.concat();
let config = Config::from_yaml(&format!(
"model_list:
{models}router_settings:
num_retries: 1
general_settings:
master_key: sk-master
"
))
.unwrap();
let inference = litellm_gateway::build_inference(&config).unwrap();
let seen = Arc::new(Collect(Mutex::new(Vec::new())));
litellm_gateway::hooks::install(&inference, &config, Vec::new(), vec![seen.clone()])
.await
.unwrap();
let app = litellm_gateway::router(inference.clone(), &config, None, None);
Harness {
app,
inference,
calls,
seen,
}
}
async fn send(app: &Router, body: Value) -> (u16, Option<String>, String) {
let request = Request::post("/v1/chat/completions")
.header("content-type", "application/json")
.header("authorization", "Bearer sk-master")
.body(Body::from(body.to_string()))
.unwrap();
let response = app.clone().oneshot(request).await.unwrap();
let status = response.status().as_u16();
let content_type = response
.headers()
.get(CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.map(str::to_owned);
let body = to_bytes(response.into_body(), 1 << 20).await.unwrap();
(
status,
content_type,
String::from_utf8(body.to_vec()).unwrap(),
)
}
fn request(model: &str, include_usage: bool) -> Value {
let mut body = json!({
"model": model,
"messages": [{"role": "user", "content": "hi"}],
"stream": true,
});
if include_usage {
body["stream_options"] = json!({"include_usage": true});
}
body
}
/// The JSON of each `data:` frame, and whether the stream closed with `[DONE]`.
fn frames(body: &str) -> (Vec<Value>, bool) {
let mut done = false;
let frames = body
.split("\n\n")
.filter(|frame| !frame.is_empty())
.filter_map(|frame| {
let data = frame.strip_prefix("data: ").expect("bare data frames");
match data {
"[DONE]" => {
done = true;
None
}
data => Some(serde_json::from_str(data).unwrap()),
}
})
.collect();
(frames, done)
}
fn usage_36_12() -> Value {
json!({"completion_tokens": 12, "prompt_tokens": 36, "total_tokens": 48,
"completion_tokens_details": {"reasoning_tokens": 0, "text_tokens": 12},
"prompt_tokens_details": {"cached_tokens": 20, "text_tokens": 10, "cache_write_tokens": 6, "cache_creation_tokens": 6},
"cache_creation_input_tokens": 6, "cache_read_input_tokens": 20})
}
fn expected_cost() -> f64 {
let (input, output, cache_read, cache_write) = RATES;
10.0 * input + 12.0 * output + 20.0 * cache_read + 6.0 * cache_write
}
async fn seen(harness: &Harness) -> Vec<Seen> {
litellm_gateway::hooks::close(&harness.inference).await;
std::mem::take(&mut *harness.seen.0.lock().unwrap())
}
#[rstest]
#[case::anthropic("claude", "claude-1", "/v1/messages")]
#[case::bedrock(
"bedrock",
"bedrock-1",
"/model/anthropic.claude-x-v1:0/converse-stream"
)]
#[tokio::test]
async fn a_translated_stream_sends_openai_chunks_and_records_its_usage(
#[case] model: &str,
#[case] id: &str,
#[case] path: &str,
#[values(true, false)] include_usage: bool,
) {
let harness = harness().await;
let (status, content_type, body) = send(&harness.app, request(model, include_usage)).await;
assert_eq!(status, 200, "{body}");
assert_eq!(content_type.as_deref(), Some("text/event-stream"));
let (frames, done) = frames(&body);
assert!(done, "{body}");
let id_of = |frame: &Value| frame["id"].clone();
assert!(frames.iter().all(|frame| id_of(frame) == id_of(&frames[0])));
assert!(
frames
.iter()
.all(|frame| frame["object"] == "chat.completion.chunk")
);
assert!(frames.iter().all(|frame| frame["model"] == model), "{body}");
let choices: Vec<Value> = frames
.iter()
.map(|frame| frame["choices"].clone())
.collect();
let mut expected = vec![
json!([{"index": 0, "delta": {"content": "Hello", "role": "assistant"}}]),
json!([{"index": 0, "delta": {"content": " world"}}]),
json!([{"index": 0, "delta": {}, "finish_reason": "stop"}]),
];
if include_usage {
expected.push(json!([{"index": 0, "delta": {}}]));
}
assert_eq!(choices, expected, "{body}");
let usages: Vec<&Value> = frames
.iter()
.filter_map(|frame| frame.get("usage"))
.collect();
match include_usage {
true => assert_eq!(usages, [&usage_36_12()]),
false => assert!(usages.is_empty(), "{body}"),
}
let sent = harness.calls.lock().unwrap()[0].clone();
assert_eq!(sent.path, path);
assert!(sent.body.get("stream_options").is_none());
let seen = seen(&harness).await;
let [seen] = seen.as_slice() else {
panic!("{seen:?}")
};
assert_eq!(seen.deployment.as_deref(), Some(id));
assert_eq!(seen.status, 200);
assert!(seen.first_byte);
let usage = seen
.usage
.expect("the usage is recorded even when the client did not ask for it");
assert_eq!(
(
usage.input_tokens,
usage.output_tokens,
usage.cache_read_tokens,
usage.cache_write_tokens
),
(10, 12, 20, 6)
);
assert!(
(seen.cost.unwrap() - expected_cost()).abs() < 1e-9,
"{seen:?}"
);
}
#[rstest]
#[tokio::test]
async fn an_anthropic_stream_gets_stream_true() {
let harness = harness().await;
let mut body = request("claude", false);
body["max_tokens"] = json!(8);
send(&harness.app, body).await;
let sent = harness.calls.lock().unwrap()[0].clone();
assert_eq!(
sent.body,
json!({"model": "claude-x", "messages": [{"role": "user", "content": [{"type": "text", "text": "hi"}]}], "max_tokens": 8, "stream": true})
);
}
#[rstest]
#[tokio::test]
async fn an_error_before_the_first_chunk_falls_over_to_another_deployment() {
let harness = harness().await;
// Whichever deployment the router picks first, the stream comes from the healthy one.
let (status, _, body) = send(&harness.app, request("fallback", true)).await;
assert_eq!(status, 200, "{body}");
let (frames, done) = frames(&body);
assert!(done);
assert_eq!(frames[0]["choices"][0]["delta"]["content"], "Hello");
let seen = seen(&harness).await;
assert_eq!(seen[0].deployment.as_deref(), Some("bedrock-3"));
assert_eq!(seen[0].status, 200);
}
#[rstest]
#[tokio::test]
async fn an_upstream_error_status_is_an_openai_error_response() {
let harness = harness().await;
let (status, content_type, body) = send(&harness.app, request("overloaded", false)).await;
assert_eq!(status, 529, "{body}");
assert_eq!(content_type.as_deref(), Some("application/json"));
let body: Value = serde_json::from_str(&body).unwrap();
assert_eq!(body["error"]["code"], 529);
let seen = seen(&harness).await;
assert_eq!(seen[0].status, 529);
assert_eq!(seen[0].usage, None);
}
#[rstest]
#[tokio::test]
async fn an_error_mid_stream_ends_with_the_openai_error_frame() {
let harness = harness().await;
let (status, _, body) = send(&harness.app, request("bedrock-cut", true)).await;
assert_eq!(status, 200, "{body}");
let (frames, done) = frames(&body);
assert!(!done, "{body}");
assert_eq!(frames[0]["choices"][0]["delta"]["content"], "Hel");
let error = &frames.last().unwrap()["error"];
assert_eq!(error["code"], 429, "{body}");
assert!(
error["message"]
.as_str()
.unwrap()
.contains("Too many requests"),
"{body}"
);
let seen = seen(&harness).await;
assert_eq!(seen[0].usage, None);
}

306
litellm-rust/crates/gateway/tests/model_capacity.rs Normal file
View file

@ -0,0 +1,306 @@
//! The model capacity hook on real requests: a group held to its deployments' summed capacity,
//! priorities held to their share, refusals and headers as the proxy sends them, and counters
//! shared through Redis. Set `LITELLM_TEST_REDIS_URL` to also run the shared case.
use std::{collections::HashMap, sync::Arc, time::Duration};
use axum::{
Router,
body::{Body, to_bytes},
http::{HeaderMap, Request, header::CONTENT_TYPE},
response::IntoResponse,
};
use litellm_auth_types::SecretValue;
use litellm_config::Config;
use litellm_gateway::{
coordination::{Coordination, RedisSettings},
model_capacity::{CapacitySettings, ModelCapacity},
};
use litellm_gateway_auth::keys::{
AllowedRoutes, KeyHash, KeyLookup, KeyModels, KeyProfile, KeyStatus, LookupFuture,
};
use rstest::rstest;
use serde_json::{Value, json};
use tokio::net::TcpListener;
use tower::ServiceExt;
/// 3 input and 4 output tokens: 7 count against the group.
const MESSAGE: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-x","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":3,"output_tokens":4}}"#;
async fn upstream() -> String {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let app = Router::new().fallback(|_: Request<Body>| async {
([(CONTENT_TYPE, "application/json")], MESSAGE).into_response()
});
tokio::spawn(async move { axum::serve(listener, app).await });
format!("http://{address}")
}
struct Keys(HashMap<KeyHash, KeyStatus>);
impl KeyLookup for Keys {
fn lookup<'a>(&'a self, hash: &'a KeyHash) -> LookupFuture<'a> {
Box::pin(async move { Ok(self.0.get(hash).cloned()) })
}
}
/// Two deployments of `claude`, each with `limits` (such as `rpm: 1`).
async fn gateway(limits: &str, settings: &str, redis: Option<Arc<Coordination>>) -> Router {
let address = upstream().await;
let config = Config::from_yaml(&format!(
"model_list:
- model_name: claude
litellm_params: {{model: anthropic/claude-x, api_base: {address}, api_key: sk-a, {limits}}}
- model_name: claude
litellm_params: {{model: anthropic/claude-y, api_base: {address}, api_key: sk-b, {limits}}}
router_settings:
num_retries: 0
general_settings:
master_key: sk-master
litellm_settings:
callbacks: [dynamic_rate_limiter_v3]
{settings}"
))
.unwrap();
let inference = litellm_gateway::build_inference(&config).unwrap();
let hook = ModelCapacity::for_gateway(&config, redis).unwrap();
litellm_gateway::hooks::install(&inference, &config, vec![Arc::new(hook)], Vec::new())
.await
.unwrap();
let mut metadata = serde_json::Map::new();
metadata.insert("priority".into(), json!("prod"));
let key = KeyStatus::Active {
expires_at: None,
models: KeyModels::default(),
allowed_routes: AllowedRoutes::default(),
team_id: None,
profile: Arc::new(KeyProfile {
metadata: Arc::new(metadata),
..KeyProfile::default()
}),
};
let keys: Arc<dyn KeyLookup> = Arc::new(Keys(HashMap::from([(
KeyHash::from_token(&SecretValue::new("sk-prod")),
key,
)])));
litellm_gateway::router_with_keys(inference, &config, None, None, Some(keys))
}
struct Reply {
status: u16,
headers: HeaderMap,
body: Value,
}
async fn send(app: &Router, path: &str, token: &str) -> Reply {
let body = json!({
"model": "claude",
"max_tokens": 8,
"messages": [{"role": "user", "content": "hi"}],
});
let response = app
.clone()
.oneshot(
Request::post(path)
.header("content-type", "application/json")
.header("authorization", format!("Bearer {token}"))
.body(Body::from(body.to_string()))
.unwrap(),
)
.await
.unwrap();
let status = response.status().as_u16();
let headers = response.headers().clone();
let body = to_bytes(response.into_body(), 1 << 20).await.unwrap();
Reply {
status,
headers,
body: serde_json::from_slice(&body).unwrap_or(Value::Null),
}
}
fn header<'a>(reply: &'a Reply, name: &str) -> Option<&'a str> {
reply.headers.get(name).map(|value| value.to_str().unwrap())
}
async fn settled() {
tokio::time::sleep(Duration::from_millis(100)).await;
}
#[rstest]
#[tokio::test]
async fn a_group_is_held_to_its_deployments_capacity() {
let app = gateway("rpm: 1", "", None).await;
let first = send(&app, "/v1/messages", "sk-master").await;
assert_eq!(first.status, 200);
assert_eq!(
header(
&first,
"x-ratelimit-model_saturation_check-remaining-requests"
),
Some("1")
);
assert_eq!(
header(&first, "x-ratelimit-model_saturation_check-limit-requests"),
Some("2")
);
assert_eq!(header(&first, "x-litellm-rate-limiter-version"), Some("v3"));
assert_eq!(header(&first, "x-litellm-priority"), Some("default"));
assert_eq!(send(&app, "/v1/messages", "sk-master").await.status, 200);
let message = "Model capacity reached for claude. Priority: None, Rate limit type: requests, Model TPM: not configured, Model RPM: 2, Remaining: 0";
let refused = send(&app, "/v1/chat/completions", "sk-master").await;
assert_eq!(refused.status, 429);
assert_eq!(
refused.body,
json!({"error": {"message": message, "type": "throttling_error", "param": null, "code": "429"}})
);
assert_eq!(header(&refused, "retry-after"), Some("60"));
assert_eq!(header(&refused, "rate_limit_type"), Some("requests"));
assert_eq!(header(&refused, "x-litellm-priority"), Some("default"));
let refused = send(&app, "/v1/messages", "sk-master").await;
assert_eq!(
refused.body,
json!({"type": "error", "error": {"type": "rate_limit_error", "message": message}})
);
assert_eq!(header(&refused, "retry-after"), None);
}
#[rstest]
#[tokio::test]
async fn tokens_count_after_the_response() {
// 2 x 5 tokens of capacity; each response uses 7.
let app = gateway("tpm: 5", "", None).await;
assert_eq!(send(&app, "/v1/messages", "sk-master").await.status, 200);
settled().await;
assert_eq!(send(&app, "/v1/messages", "sk-master").await.status, 200);
settled().await;
let refused = send(&app, "/v1/messages", "sk-master").await;
assert_eq!(refused.status, 429);
assert_eq!(
refused.body["error"]["message"],
"Model capacity reached for claude. Priority: None, Rate limit type: tokens, Model TPM: 10, Model RPM: not configured, Remaining: -4"
);
}
#[rstest]
#[tokio::test]
async fn busy_groups_hold_priorities_to_their_share() {
// 20 requests a window; `prod` gets a fifth once the group is busy, here from the start.
let app = gateway(
"rpm: 10",
" priority_reservation: {prod: 0.1}\n priority_reservation_settings: {saturation_threshold: 0.0}\n",
None,
)
.await;
assert_eq!(
send(&app, "/v1/chat/completions", "sk-prod").await.status,
200
);
let second = send(&app, "/v1/chat/completions", "sk-prod").await;
assert_eq!(second.status, 200);
assert_eq!(header(&second, "x-litellm-priority"), Some("prod"));
assert_eq!(
header(&second, "x-ratelimit-priority_model-remaining-requests"),
Some("0")
);
let refused = send(&app, "/v1/chat/completions", "sk-prod").await;
assert_eq!(refused.status, 429);
assert_eq!(
refused.body["error"]["message"],
"Priority-based rate limit exceeded. Model: claude, Priority: prod, Rate limit type: requests, Model TPM: not configured, Model RPM: 20, Remaining: 0, Model saturation: 0.0%"
);
assert_eq!(header(&refused, "x-litellm-saturation"), Some("0.00%"));
// Other callers share the default pool and the rest of the group.
assert_eq!(
send(&app, "/v1/chat/completions", "sk-master").await.status,
200
);
}
#[rstest]
fn the_hook_is_on_only_when_listed() {
let config = Config::from_yaml("model_list: []").unwrap();
assert_eq!(CapacitySettings::from_config(&config), None);
let config =
Config::from_yaml("model_list: []\nlitellm_settings:\n callbacks: model_capacity\n")
.unwrap();
assert_eq!(
CapacitySettings::from_config(&config),
Some(CapacitySettings::default())
);
}
async fn redis(url: &str, namespace: &str) -> Arc<Coordination> {
let coordination =
Coordination::new(RedisSettings::new(url, Some(namespace.to_owned()))).unwrap();
let router: Arc<dyn litellm_gateway::coordination::RouterHandle> = Arc::new(NoRouter(
litellm_gateway_inference::ModelRouter::from_config(
&Config::from_yaml("model_list: []").unwrap(),
)
.unwrap(),
));
coordination.start(router, None, tokio_util::sync::CancellationToken::new());
for _ in 0..100 {
if coordination.connected() {
return coordination;
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
panic!("Redis did not connect");
}
struct NoRouter(litellm_gateway_inference::ModelRouter);
impl litellm_gateway::coordination::RouterHandle for NoRouter {
fn router(&self) -> &litellm_gateway_inference::ModelRouter {
&self.0
}
}
#[rstest]
#[tokio::test]
async fn replicas_share_the_group_counters() {
let Ok(url) = std::env::var("LITELLM_TEST_REDIS_URL") else {
return;
};
let namespace = format!("model-capacity-{}", uuid::Uuid::new_v4());
let limits = "rpm: 2, tpm: 1000";
let first = gateway(limits, "", Some(redis(&url, &namespace).await)).await;
let second = gateway(limits, "", Some(redis(&url, &namespace).await)).await;
for app in [&first, &second, &first, &second] {
assert_eq!(send(app, "/v1/messages", "sk-master").await.status, 200);
}
settled().await;
assert_eq!(send(&second, "/v1/messages", "sk-master").await.status, 429);
let mut connection = redis::Client::open(url.as_str())
.unwrap()
.get_multiplexed_async_connection()
.await
.unwrap();
let get = async |connection: &mut redis::aio::MultiplexedConnection, key: String| {
redis::cmd("GET")
.arg(key)
.query_async::<Option<String>>(connection)
.await
.unwrap()
};
let group = format!("{namespace}:{{model_saturation_check:claude}}");
// The refused request is not counted; four responses used 7 tokens each.
assert_eq!(
get(&mut connection, format!("{group}:requests"))
.await
.as_deref(),
Some("4")
);
assert_eq!(
get(&mut connection, format!("{group}:tokens"))
.await
.as_deref(),
Some("28")
);
let pool = format!("{namespace}:{{priority_model:claude:default_pool}}:tokens");
assert_eq!(get(&mut connection, pool).await.as_deref(), Some("28"));
}

340
litellm-rust/crates/gateway/tests/request_hooks.rs Normal file
View file

@ -0,0 +1,340 @@
//! The configurable request hooks on real requests: headers forwarded and stripped, provider
//! metadata set from the caller, and spend metadata handed to observers.
use std::{
collections::HashMap,
sync::{Arc, Mutex},
};
use axum::{
Router,
body::{Body, Bytes, to_bytes},
http::{HeaderMap, Request, header::CONTENT_TYPE},
response::IntoResponse,
};
use futures_util::future::BoxFuture;
use litellm_auth_types::SecretValue;
use litellm_config::Config;
use litellm_gateway_auth::keys::{
AllowedRoutes, KeyHash, KeyLookup, KeyModels, KeyProfile, KeyStatus, LookupFuture,
};
use litellm_gateway_hooks::{
Facts, KeyFacts, Observer, Observers, Pipeline, Record, RequestFacts, Route,
};
use rstest::rstest;
use serde_json::{Map, Value, json};
use tokio::net::TcpListener;
use tower::ServiceExt;
const MESSAGE: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-x","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":3,"output_tokens":4}}"#;
type Seen = Arc<Mutex<Vec<(HeaderMap, Value)>>>;
async fn upstream() -> (String, Seen) {
let seen = Seen::default();
let recorded = seen.clone();
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let app = Router::new().fallback(move |request: Request<Body>| {
let recorded = recorded.clone();
async move {
let headers = request.headers().clone();
let body = to_bytes(request.into_body(), 1 << 20).await.unwrap();
recorded.lock().unwrap().push((
headers,
serde_json::from_slice(&body).unwrap_or(Value::Null),
));
([(CONTENT_TYPE, "application/json")], MESSAGE).into_response()
}
});
tokio::spawn(async move { axum::serve(listener, app).await });
(format!("http://{address}"), seen)
}
struct Keys(HashMap<KeyHash, KeyStatus>);
impl KeyLookup for Keys {
fn lookup<'a>(&'a self, hash: &'a KeyHash) -> LookupFuture<'a> {
Box::pin(async move { Ok(self.0.get(hash).cloned()) })
}
}
const HOOKS: &str = r#"
hooks:
- forward_headers:
headers: ["x-trace-*", "x-cost-center", "authorization"]
- strip_headers:
headers: ["x-cost-center", "anthropic-beta"]
- provider_metadata:
body:
metadata.user_id: "{end_user|key.user_id}"
headers:
x-caller-team: "team-{key.team_alias}"
"#;
async fn gateway(hooks: &str) -> (Router, Seen) {
let (address, seen) = upstream().await;
let config = Config::from_yaml(&format!(
"model_list:
- model_name: claude
litellm_params:
model: anthropic/claude-x
api_base: {address}
api_key: sk-upstream
general_settings:
master_key: sk-master
litellm_settings:{hooks}"
))
.unwrap();
let inference = litellm_gateway::build_inference(&config).unwrap();
let hooks = litellm_gateway::request_hooks::from_config(&config).unwrap();
litellm_gateway::hooks::install(&inference, &config, hooks, Vec::new())
.await
.unwrap();
let key = KeyStatus::Active {
expires_at: None,
models: KeyModels::default(),
allowed_routes: AllowedRoutes::default(),
team_id: Some("team-1".into()),
profile: Arc::new(KeyProfile {
user_id: Some("user-1".into()),
team_alias: Some("research".into()),
..KeyProfile::default()
}),
};
let keys: Arc<dyn KeyLookup> = Arc::new(Keys(HashMap::from([(
KeyHash::from_token(&SecretValue::new("sk-key")),
key,
)])));
let app = litellm_gateway::router_with_keys(inference, &config, None, None, Some(keys));
(app, seen)
}
async fn send(app: &Router, path: &str, body: Value) -> u16 {
let response = app
.clone()
.oneshot(
Request::post(path)
.header("content-type", "application/json")
.header("authorization", "Bearer sk-key")
.header("x-trace-id", "trace-1")
.header("x-cost-center", "cc-9")
.header("anthropic-beta", "some-beta")
.header("x-other", "kept-out")
.body(Body::from(body.to_string()))
.unwrap(),
)
.await
.unwrap();
response.status().as_u16()
}
fn header<'a>(headers: &'a HeaderMap, name: &str) -> Option<&'a str> {
headers.get(name).map(|value| value.to_str().unwrap())
}
#[rstest]
#[tokio::test]
async fn messages_carry_the_configured_headers_and_metadata() {
let (app, seen) = gateway(HOOKS).await;
let body = json!({
"model": "claude",
"max_tokens": 8,
"metadata": {"tag": "t"},
"messages": [{"role": "user", "content": "hi"}],
});
assert_eq!(send(&app, "/v1/messages", body).await, 200);
let (headers, body) = seen.lock().unwrap().pop().unwrap();
assert_eq!(header(&headers, "x-trace-id"), Some("trace-1"));
assert_eq!(header(&headers, "x-caller-team"), Some("team-research"));
// Stripped after forwarding, and the gateway's own pass-through is stripped too.
assert_eq!(header(&headers, "x-cost-center"), None);
assert_eq!(header(&headers, "anthropic-beta"), None);
assert_eq!(header(&headers, "x-other"), None);
// Credentials are never forwarded, even when listed.
assert_ne!(header(&headers, "authorization"), Some("Bearer sk-key"));
assert_eq!(header(&headers, "x-api-key"), Some("sk-upstream"));
// The provider keeps only the `metadata` fields it knows.
assert_eq!(body["metadata"]["user_id"], "user-1");
}
#[rstest]
#[tokio::test]
async fn the_end_user_comes_before_the_key_user() {
let (app, seen) = gateway(HOOKS).await;
let body = json!({
"model": "claude",
"max_tokens": 8,
"litellm_metadata": {"user": "end-7"},
"messages": [{"role": "user", "content": "hi"}],
});
assert_eq!(send(&app, "/v1/messages", body).await, 200);
let (_, body) = seen.lock().unwrap().pop().unwrap();
assert_eq!(body["metadata"]["user_id"], "end-7");
}
#[rstest]
#[tokio::test]
async fn chat_completions_take_the_headers_only() {
let (app, seen) = gateway(HOOKS).await;
let body = json!({
"model": "claude",
"max_tokens": 8,
"messages": [{"role": "user", "content": "hi"}],
});
assert_eq!(send(&app, "/v1/chat/completions", body).await, 200);
let (headers, body) = seen.lock().unwrap().pop().unwrap();
assert_eq!(header(&headers, "x-trace-id"), Some("trace-1"));
assert_eq!(header(&headers, "x-caller-team"), Some("team-research"));
assert_eq!(header(&headers, "x-cost-center"), None);
assert_eq!(body.get("metadata"), None);
}
#[rstest]
#[tokio::test]
async fn without_hooks_the_default_pass_through_is_unchanged() {
let (app, seen) = gateway("\n drop_params: true\n").await;
let body = json!({
"model": "claude",
"max_tokens": 8,
"messages": [{"role": "user", "content": "hi"}],
});
assert_eq!(send(&app, "/v1/messages", body).await, 200);
let (headers, _) = seen.lock().unwrap().pop().unwrap();
assert_eq!(header(&headers, "anthropic-beta"), Some("some-beta"));
assert_eq!(header(&headers, "x-trace-id"), None);
}
struct Annotations(Mutex<Vec<Option<Value>>>);
impl Observer for Annotations {
fn name(&self) -> &'static str {
"annotations"
}
fn reads(&self) -> Facts {
Facts::ANNOTATIONS
}
fn observe<'a>(&'a self, record: &'a Record) -> BoxFuture<'a, ()> {
Box::pin(async move {
self.0.lock().unwrap().push(
record
.annotations
.get(litellm_gateway::spend::SPEND_LOGS_METADATA)
.cloned(),
);
})
}
}
#[rstest]
#[tokio::test]
async fn spend_metadata_reaches_observers() {
let config = Config::from_yaml(
r#"
model_list: []
litellm_settings:
hooks:
- spend_metadata:
from_headers:
x-litellm-spend-logs-metadata: json
x-cost-center: cost_center
from_key_metadata: [owner]
"#,
)
.unwrap();
let observed = Arc::new(Annotations(Mutex::default()));
let pipeline = Pipeline::new(
litellm_gateway::request_hooks::from_config(&config).unwrap(),
Observers::start(vec![observed.clone()], 8),
Arc::default(),
);
let run = async |headers: &[(&'static str, &'static str)]| {
let mut map = HeaderMap::new();
for (name, value) in headers {
map.insert(*name, value.parse().unwrap());
}
let Value::Object(metadata) = json!({"owner": "platform"}) else {
unreachable!()
};
let key = KeyFacts {
metadata: Arc::new(metadata),
..KeyFacts::default()
};
let request = RequestFacts::new(
"r".into(),
Route::Messages,
Some(Arc::new(key)),
None,
"claude".into(),
Arc::new(map),
false,
None,
Bytes::from_static(b"{}"),
);
let admission = pipeline.admit(Arc::new(request)).await.unwrap();
pipeline.finish(
admission,
litellm_gateway_hooks::Outcome {
deployment: None,
status: 200,
usage: None,
cost: None,
time_to_first_byte: None,
upstream_started_at: None,
ended_at: std::time::SystemTime::now(),
interrupted: false,
},
None,
);
};
run(&[
(
"x-litellm-spend-logs-metadata",
r#"{"project": "p1", "owner": "caller"}"#,
),
("x-cost-center", "cc-9"),
])
.await;
run(&[("x-litellm-spend-logs-metadata", "not json")]).await;
pipeline.close().await;
let mut expected = Map::new();
expected.insert("project".into(), json!("p1"));
expected.insert("owner".into(), json!("platform"));
expected.insert("cost_center".into(), json!("cc-9"));
assert_eq!(
*observed.0.lock().unwrap(),
[
Some(Value::Object(expected)),
Some(json!({"owner": "platform"}))
]
);
}
#[rstest]
#[case::unknown_hook(" - rewrite_everything: {}\n", "unknown hook rewrite_everything")]
#[case::bad_template(
" - provider_metadata:\n headers:\n x-a: \"{key.secret}\"\n",
"provider_metadata: unknown fact key.secret"
)]
#[case::credential_header(
" - provider_metadata:\n headers:\n authorization: \"{key.user_id}\"\n",
"provider_metadata: authorization cannot be set by a hook"
)]
#[case::unknown_setting(
" - strip_headers:\n names: [x-a]\n",
"strip_headers: unknown field `names`"
)]
fn misconfigured_hooks_stop_the_gateway(#[case] hooks: &str, #[case] message: &str) {
let config = Config::from_yaml(&format!(
"model_list: []\nlitellm_settings:\n hooks:\n{hooks}"
))
.unwrap();
let error = litellm_gateway::request_hooks::from_config(&config)
.err()
.unwrap()
.to_string();
assert!(error.contains(message), "{error}");
}

313
litellm-rust/crates/gateway/tests/s3_log.rs Normal file
View file

@ -0,0 +1,313 @@
//! The `s3_v2` observer against a mock S3: a request ends up as a signed PUT under the proxy's
//! key layout, a 503 is retried, and an object S3 rejects is dropped instead of kept.
use std::{
collections::{BTreeMap, HashMap},
sync::Arc,
time::{Duration, SystemTime, UNIX_EPOCH},
};
use axum::{
Router,
body::{Body, to_bytes},
http::{Request, header::CONTENT_TYPE},
response::IntoResponse,
};
use base64::Engine;
use litellm_auth_aws::{Credentials, sign_s3};
use litellm_config::Config;
use litellm_gateway::s3_log::S3Log;
use litellm_gateway_auth::keys::{
AllowedRoutes, KeyHash, KeyLookup, KeyModels, KeyProfile, KeyStatus, LookupFuture,
};
use serde_json::{Value, json};
use sha2::{Digest, Sha256};
use tokio::net::TcpListener;
use tower::ServiceExt;
use wiremock::{
Mock, MockServer, ResponseTemplate,
matchers::{method, path_regex},
};
const COMPLETION: &str = r#"{"id":"chatcmpl-1","object":"chat.completion","created":1,"model":"gpt-x","choices":[{"index":0,"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":5,"completion_tokens":6,"total_tokens":11}}"#;
const MESSAGE: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-x","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn","usage":{"input_tokens":3,"output_tokens":4}}"#;
const OBJECT_KEY: &str =
r"^/logs/audit/research/\d{4}-\d{2}-\d{2}/time-\d{2}-\d{2}-\d{2}-\d{6}_chatcmpl-1\.json$";
/// Answers Anthropic and OpenAI calls.
async fn upstream() -> String {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let app = Router::new().fallback(|request: Request<Body>| async move {
let body = match request.uri().path().ends_with("/chat/completions") {
true => COMPLETION,
false => MESSAGE,
};
([(CONTENT_TYPE, "application/json")], body).into_response()
});
tokio::spawn(async move { axum::serve(listener, app).await });
format!("http://{address}")
}
struct Keys(HashMap<KeyHash, KeyStatus>);
impl KeyLookup for Keys {
fn lookup<'a>(&'a self, hash: &'a KeyHash) -> LookupFuture<'a> {
Box::pin(async move { Ok(self.0.get(hash).cloned()) })
}
}
fn keys() -> Arc<dyn KeyLookup> {
let status = KeyStatus::Active {
expires_at: None,
models: KeyModels::default(),
allowed_routes: AllowedRoutes::default(),
team_id: Some("team-1".into()),
profile: Arc::new(KeyProfile {
alias: Some("team-key".into()),
user_id: Some("user-1".into()),
team_alias: Some("research".into()),
..KeyProfile::default()
}),
};
let hash = KeyHash::from_token(&litellm_auth_types::SecretValue::new("sk-team"));
Arc::new(Keys(HashMap::from([(hash, status)])))
}
struct Harness {
app: Router,
inference: Arc<litellm_gateway_inference::Gateway>,
s3_log: Arc<S3Log>,
}
async fn harness(s3: &MockServer) -> Harness {
let upstream = upstream().await;
let config = Config::from_yaml(&format!(
"model_list:
- model_name: gpt
litellm_params:
model: openai/gpt-x
api_base: {upstream}/v1
api_key: sk-upstream
model_info:
id: gpt-1
- model_name: claude
litellm_params:
model: anthropic/claude-x
api_base: {upstream}
api_key: sk-upstream
router_settings:
num_retries: 0
general_settings:
master_key: sk-master
litellm_settings:
callbacks: [s3_v2]
s3_callback_params:
s3_bucket_name: logs
s3_path: audit
s3_region_name: us-east-1
s3_endpoint_url: {endpoint}
s3_aws_access_key_id: AKIDTEST
s3_aws_secret_access_key: secret
s3_use_team_prefix: true
s3_server_side_encryption: AES256
",
endpoint = s3.uri(),
))
.unwrap();
let inference = litellm_gateway::build_inference(&config).unwrap();
let s3_log = S3Log::from_config(&config, &inference)
.await
.unwrap()
.unwrap();
litellm_gateway::hooks::install(&inference, &config, Vec::new(), vec![s3_log.clone()])
.await
.unwrap();
let app =
litellm_gateway::router_with_keys(inference.clone(), &config, None, None, Some(keys()));
Harness {
app,
inference,
s3_log,
}
}
async fn send(app: &Router, path: &str, body: Value) {
let response = app
.clone()
.oneshot(
Request::post(path)
.header("content-type", "application/json")
.header("authorization", "Bearer sk-team")
.header("user-agent", "client/1.0")
.body(Body::from(body.to_string()))
.unwrap(),
)
.await
.unwrap();
let status = response.status();
let body = to_bytes(response.into_body(), 1 << 20).await.unwrap();
assert_eq!(status, 200, "{body:?}");
}
/// `YYYYMMDDTHHMMSSZ` as a time.
fn amz_date(text: &str) -> SystemTime {
let number = |range: std::ops::Range<usize>| text[range].parse::<i64>().unwrap();
let (year, month, day) = (number(0..4), number(4..6), number(6..8));
let year = if month <= 2 { year - 1 } else { year };
let era = year.div_euclid(400);
let year_of_era = year - era * 400;
let day_of_year = (153 * (month + if month > 2 { -3 } else { 9 }) + 2) / 5 + day - 1;
let day_of_era = year_of_era * 365 + year_of_era / 4 - year_of_era / 100 + day_of_year;
let days = era * 146_097 + day_of_era - 719_468;
let seconds = days * 86_400 + number(9..11) * 3600 + number(11..13) * 60 + number(13..15);
UNIX_EPOCH + Duration::from_secs(seconds as u64)
}
#[tokio::test]
async fn a_request_is_written_as_a_signed_put_and_a_503_is_retried() {
let s3 = MockServer::start().await;
Mock::given(method("PUT"))
.and(path_regex(OBJECT_KEY))
.respond_with(ResponseTemplate::new(503))
.up_to_n_times(1)
.with_priority(1)
.mount(&s3)
.await;
Mock::given(method("PUT"))
.and(path_regex(OBJECT_KEY))
.respond_with(ResponseTemplate::new(200))
.expect(1)
.mount(&s3)
.await;
let harness = harness(&s3).await;
send(
&harness.app,
"/v1/chat/completions",
json!({"model": "gpt", "temperature": 0.5, "messages": [{"role": "user", "content": "hello"}]}),
)
.await;
litellm_gateway::hooks::close(&harness.inference).await;
let requests = s3.received_requests().await.unwrap();
assert_eq!(requests.len(), 2, "the 503 is retried once");
assert_eq!(requests[0].url, requests[1].url);
let put = &requests[1];
let header = |name: &str| {
put.headers
.get(name)
.map(|value| value.to_str().unwrap().to_owned())
};
let payload: Value = serde_json::from_slice(&put.body).unwrap();
assert_eq!(payload["id"], "chatcmpl-1");
assert_eq!(payload["status"], "success");
assert_eq!(payload["model_group"], "gpt");
assert_eq!(payload["model_id"], "gpt-1");
assert_eq!(
payload["messages"],
json!([{"role": "user", "content": "hello"}])
);
assert_eq!(
payload["response"]["choices"][0]["message"]["content"],
"hi"
);
assert_eq!(payload["metadata"]["user_api_key_team_alias"], "research");
assert_eq!(payload["total_tokens"], 11);
assert_eq!(header("content-type").as_deref(), Some("application/json"));
assert_eq!(header("content-language").as_deref(), Some("en"));
assert_eq!(
header("cache-control").as_deref(),
Some("private, immutable, max-age=31536000, s-maxage=0")
);
assert_eq!(
header("x-amz-server-side-encryption").as_deref(),
Some("AES256")
);
let disposition = header("content-disposition").unwrap();
assert!(
disposition.starts_with("inline; filename=\"time-")
&& disposition.ends_with("_chatcmpl-1.json\""),
"{disposition}"
);
let sha256: String = Sha256::digest(&put.body)
.iter()
.map(|byte| format!("{byte:02x}"))
.collect();
assert_eq!(header("x-amz-content-sha256"), Some(sha256));
let md5 = base64::engine::general_purpose::STANDARD
.encode(<md5::Md5 as md5::Digest>::digest(&put.body));
assert_eq!(header("content-md5"), Some(md5));
let signed: BTreeMap<String, String> = [
"content-type",
"content-md5",
"x-amz-content-sha256",
"content-language",
"content-disposition",
"cache-control",
"x-amz-server-side-encryption",
]
.into_iter()
.map(|name| (name.to_owned(), header(name).unwrap()))
.collect();
let expected = sign_s3(
"PUT",
&format!("{}{}", s3.uri(), put.url.path()),
&put.body,
&signed,
"us-east-1",
&Credentials::new("AKIDTEST", "secret", None, None, "test"),
amz_date(&header("x-amz-date").unwrap()),
)
.unwrap();
assert_eq!(header("x-amz-date"), Some(expected["X-Amz-Date"].clone()));
assert_eq!(
header("authorization"),
Some(expected["Authorization"].clone())
);
assert!(
expected["Authorization"].contains("SignedHeaders=cache-control;content-disposition;content-language;content-md5;content-type;host;x-amz-content-sha256;x-amz-date;x-amz-server-side-encryption"),
"{}",
expected["Authorization"]
);
}
#[tokio::test]
async fn an_object_s3_rejects_is_dropped_while_the_rest_are_written() {
let s3 = MockServer::start().await;
Mock::given(method("PUT"))
.and(path_regex(r"_msg_1\.json$"))
.respond_with(
ResponseTemplate::new(400)
.set_body_string("<Error><Code>EntityTooLarge</Code></Error>"),
)
.expect(1)
.mount(&s3)
.await;
Mock::given(method("PUT"))
.and(path_regex(OBJECT_KEY))
.respond_with(ResponseTemplate::new(200))
.expect(1)
.mount(&s3)
.await;
let harness = harness(&s3).await;
send(
&harness.app,
"/v1/chat/completions",
json!({"model": "gpt", "messages": [{"role": "user", "content": "hello"}]}),
)
.await;
send(
&harness.app,
"/v1/messages",
json!({"model": "claude", "max_tokens": 8, "messages": [{"role": "user", "content": "hello"}]}),
)
.await;
litellm_gateway::hooks::close(&harness.inference).await;
harness.s3_log.flush().await;
assert_eq!(s3.received_requests().await.unwrap().len(), 2);
}

3
litellm-rust/crates/inference-chat/Cargo.toml
View file

@ -21,13 +21,16 @@ reqwest.workspace = true
serde_json = { workspace = true, features = ["preserve_order"] }
tokio = { workspace = true, features = ["time"] }
tracing.workspace = true
uuid.workspace = true
[dev-dependencies]
litellm-cache-memory.workspace = true
litellm-framer.workspace = true
litellm-host-native.workspace = true
litellm-http = { workspace = true, features = ["test-support"] }
litellm-inference-testing.workspace = true
litellm-tracing.workspace = true
rstest.workspace = true
tokio-util = { version = "0.7", features = ["codec"] }
tokio = { workspace = true, features = ["io-util", "sync"] }
wiremock.workspace = true

475
litellm-rust/crates/inference-chat/src/azure.rs Normal file
View file

@ -0,0 +1,475 @@
//! Azure OpenAI chat completions through the OpenAI passthrough: `litellm/llms/azure/azure.py`
//! with the client `BaseAzureLLM.get_azure_openai_client` builds. The deployment's settings
//! and `os.environ/` references resolve through the secret source, and an Entra ID token is
//! acquired through `litellm_auth_azure` when the deployment has no key.
use std::collections::BTreeMap;
use litellm_auth::azure::{AzureAuthInputs, SECRET_NAMES as AZURE_AUTH_SECRET_NAMES};
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
use litellm_llms::{
Error as LlmError,
azure::{
chat::transformation::{AzureChatBody, AzureChatRequest},
common_utils::{self, AZURE_API_KEY_ENV, AzureCredential},
},
base_llm::auth::resolve_auth,
};
use serde_json::{Map, Value};
use crate::{
ChatCompletionsRoute, Error,
common_utils::string_headers,
openai::{OpenAiChatCall, PreparedCall},
};
const ENV_REFERENCE: &str = "os.environ/";
/// The `litellm_params` an `os.environ/` reference may stand in for.
const REFERABLE_PARAMS: &[&str] = &[
"api_version",
"base_model",
"azure_ad_token",
"tenant_id",
"client_id",
"client_secret",
"azure_scope",
"azure_authority_host",
"azure_credential",
"azure_federated_token_file",
];
impl ChatCompletionsRoute {
pub(crate) async fn prepare_azure(
&self,
deployment: String,
call: OpenAiChatCall,
) -> Result<PreparedCall, Error> {
let OpenAiChatCall {
body,
api_key,
api_base,
extra_headers,
litellm_params,
drop_params,
model_info,
..
} = call;
let references: Vec<String> = [api_key.as_deref(), api_base.as_deref()]
.into_iter()
.flatten()
.chain(
REFERABLE_PARAMS
.iter()
.filter_map(|name| litellm_params.get(*name).and_then(Value::as_str)),
)
.filter_map(|value| value.strip_prefix(ENV_REFERENCE))
.map(str::to_owned)
.collect();
let names: Vec<&str> = common_utils::SECRET_NAMES
.iter()
.chain(AZURE_AUTH_SECRET_NAMES)
.copied()
.chain(references.iter().map(String::as_str))
.collect();
let secrets = self.secrets.resolve(&names).await?;
let lookup = |key: &str| secrets.get(key);
// Python's router reads every `os.environ/` setting from the environment when it loads
// the deployment. A reference to an unset variable counts as unset.
let resolve = |value: &str| match value.strip_prefix(ENV_REFERENCE) {
Some(name) => lookup(name),
None => Some(value.to_owned()),
};
let params: Map<String, Value> = litellm_params
.into_iter()
.map(|(name, value)| match value {
Value::String(text) if REFERABLE_PARAMS.contains(&name.as_str()) => {
let value = resolve(&text).map_or(Value::Null, Value::String);
(name, value)
}
other => (name, other),
})
.collect();
let configured = |name: &str| params.get(name).and_then(Value::as_str);
let api_version = common_utils::resolve_api_version(configured("api_version"), &lookup);
let api_base = common_utils::resolve_api_base(
api_base.as_deref().and_then(resolve).as_deref(),
&lookup,
)?;
let api_key =
common_utils::resolve_api_key(api_key.as_deref().and_then(resolve).as_deref(), &lookup);
let model_info = model_info.as_deref();
let AzureChatBody {
deployment: url_deployment,
body,
} = AzureChatRequest {
deployment: &deployment,
base_model: configured("base_model"),
api_version: &api_version,
drop_params,
model_info,
}
.transform(body)?;
let url = common_utils::complete_url(&api_base, &url_deployment, &api_version)?;
let credential = self
.azure_credential(api_key.clone(), &params, &lookup)
.await?;
let environment = common_utils::environment(
string_headers(extra_headers)?,
credential,
common_utils::is_v1_api_version(&api_version),
);
let authenticated = resolve_auth(&self.auth, environment, &lookup).await?;
Ok(PreparedCall {
model: url_deployment,
provider: <&str>::from(LlmProviders::Azure),
url,
headers: authenticated.headers,
body,
api_key,
})
}
/// The SDK sends a configured `azure_ad_token` ahead of a key. Without either, LiteLLM
/// gets a token from Entra ID: client credentials, an OIDC exchange, or managed identity
/// when `azure_credential` or `enable_azure_ad_token_refresh` asks for it.
async fn azure_credential(
&self,
api_key: Option<String>,
params: &Map<String, Value>,
lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<AzureCredential, Error> {
let configured_token = params
.get("azure_ad_token")
.and_then(Value::as_str)
.is_some_and(|token| !token.is_empty());
if let Some(key) = &api_key
&& !configured_token
{
return Ok(AzureCredential::ApiKey(litellm_auth::SecretValue::new(
key.clone(),
)));
}
let inputs = AzureAuthInputs::from_sourced_optional_params(params, &BTreeMap::new())?;
let token = self.auth.azure.get_azure_ad_token(&inputs, lookup).await?;
match (token, api_key) {
(Some(token), _) => Ok(AzureCredential::AdToken(
token.into_value().secret().clone(),
)),
(None, Some(key)) => Ok(AzureCredential::ApiKey(litellm_auth::SecretValue::new(key))),
(None, None) => Err(Error::from(LlmError::Auth(
litellm_auth::Error::MissingApiKey {
provider: "Azure",
environment_variable: AZURE_API_KEY_ENV,
},
))),
}
}
}
#[cfg(test)]
mod tests {
use std::{sync::Arc, time::Duration};
use bytes::Bytes;
use futures_util::TryStreamExt;
use litellm_host::call::CallOutput;
use litellm_http::transport;
use litellm_inference_testing::RecordingSecrets;
use rstest::rstest;
use serde_json::json;
use wiremock::{
Mock, MockServer, ResponseTemplate,
matchers::{body_partial_json, header, method, path, query_param},
};
use super::*;
use crate::test_support::{Recording, route, run};
const COMPLETION: &str = r#"{"id":"chatcmpl-1","object":"chat.completion","created":1,"model":"gpt-4o-2024-08-06","choices":[{"index":0,"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":3,"completion_tokens":2,"total_tokens":5}}"#;
const STREAM: &str = "data: {\"id\":\"c\",\"model\":\"gpt-4o-2024-08-06\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hi\"}}]}\n\ndata: {\"id\":\"c\",\"model\":\"gpt-4o-2024-08-06\",\"choices\":[],\"usage\":{\"prompt_tokens\":3,\"completion_tokens\":1,\"total_tokens\":4}}\n\ndata: [DONE]\n\n";
const DEPLOYMENT_PATH: &str = "/openai/deployments/gpt-4o/chat/completions";
fn call(api_base: &str, body: Value) -> OpenAiChatCall {
let Value::Object(body) = body else {
panic!("the body is an object");
};
OpenAiChatCall {
model: "azure/gpt-4o".into(),
body,
api_key: Some("sk-azure".into()),
api_base: Some(format!("{api_base}/")),
litellm_params: Map::from_iter([("api_version".into(), json!("2024-10-21"))]),
timeout: Some(Duration::from_secs(5)),
..OpenAiChatCall::default()
}
}
fn hi() -> Value {
json!({"model": "public", "messages": [{"role": "user", "content": "hi"}]})
}
fn with_secrets(route: ChatCompletionsRoute, values: &[(&str, &str)]) -> ChatCompletionsRoute {
ChatCompletionsRoute {
secrets: Arc::new(RecordingSecrets::new(values.iter().copied())),
..route
}
}
async fn complete(output: crate::openai::OpenAiChatOutput) -> Bytes {
match output {
CallOutput::Complete(body) => body,
CallOutput::Stream { .. } => panic!("expected a complete body"),
}
}
#[rstest]
#[tokio::test]
async fn a_completion_goes_to_the_deployment_with_its_key(route: ChatCompletionsRoute) {
let upstream = MockServer::start().await;
Mock::given(method("POST"))
.and(path(DEPLOYMENT_PATH))
.and(query_param("api-version", "2024-10-21"))
.and(header("api-key", "sk-azure"))
.and(body_partial_json(json!({"model": "gpt-4o"})))
.respond_with(ResponseTemplate::new(200).set_body_raw(COMPLETION, "application/json"))
.expect(1)
.mount(&upstream)
.await;
let interceptors = Recording::default();
let output = run(&route, call(&upstream.uri(), hi()), &interceptors)
.await
.unwrap();
assert_eq!(complete(output).await, COMPLETION.as_bytes());
let requests = upstream.received_requests().await.unwrap();
assert!(!requests[0].headers.contains_key("authorization"));
let contexts = interceptors.contexts.lock().unwrap();
assert_eq!(
(
contexts[0].model.as_str(),
contexts[0].custom_llm_provider.as_str()
),
("gpt-4o", "azure")
);
assert_eq!(*interceptors.raw.lock().unwrap(), [COMPLETION]);
assert_eq!(
interceptors.results.lock().unwrap()[0].provider.provider,
"azure"
);
}
#[rstest]
#[tokio::test]
async fn a_stream_relays_every_byte(route: ChatCompletionsRoute) {
let upstream = MockServer::start().await;
Mock::given(path(DEPLOYMENT_PATH))
.and(body_partial_json(
json!({"stream": true, "stream_options": {"include_usage": true}}),
))
.respond_with(ResponseTemplate::new(200).set_body_raw(STREAM, "text/event-stream"))
.expect(1)
.mount(&upstream)
.await;
let mut body = hi();
body["stream"] = json!(true);
body["stream_options"] = json!({"include_usage": true});
let output = run(&route, call(&upstream.uri(), body), &Recording::default())
.await
.unwrap();
let CallOutput::Stream { chunks, .. } = output else {
panic!("a streaming call streams");
};
let relayed = chunks.try_collect::<Vec<_>>().await.unwrap().concat();
assert_eq!(relayed, STREAM.as_bytes());
}
#[rstest]
#[tokio::test]
async fn the_v1_surface_takes_the_key_as_a_bearer(route: ChatCompletionsRoute) {
let upstream = MockServer::start().await;
Mock::given(path("/openai/v1/chat/completions"))
.and(header("authorization", "Bearer sk-azure"))
.respond_with(ResponseTemplate::new(200).set_body_raw(COMPLETION, "application/json"))
.expect(1)
.mount(&upstream)
.await;
let mut v1 = call(&upstream.uri(), hi());
v1.litellm_params
.insert("api_version".into(), json!("latest"));
run(&route, v1, &Recording::default()).await.unwrap();
let requests = upstream.received_requests().await.unwrap();
assert_eq!(requests[0].url.query(), None);
assert!(!requests[0].headers.contains_key("api-key"));
}
#[rstest]
#[tokio::test]
async fn base_key_and_version_fall_back_to_the_environment(route: ChatCompletionsRoute) {
let upstream = MockServer::start().await;
Mock::given(path(DEPLOYMENT_PATH))
.and(query_param("api-version", "2024-06-01"))
.and(header("api-key", "sk-env"))
.respond_with(ResponseTemplate::new(200).set_body_raw(COMPLETION, "application/json"))
.expect(1)
.mount(&upstream)
.await;
let base = upstream.uri();
let route = with_secrets(
route,
&[
("AZURE_API_BASE", &base),
("AZURE_API_KEY", "sk-env"),
("AZURE_API_VERSION", "2024-06-01"),
],
);
let bare = OpenAiChatCall {
api_key: None,
api_base: None,
litellm_params: Map::new(),
..call(&base, hi())
};
run(&route, bare, &Recording::default()).await.unwrap();
}
#[rstest]
#[tokio::test]
async fn the_default_api_version_is_pythons(route: ChatCompletionsRoute) {
let upstream = MockServer::start().await;
Mock::given(query_param(
"api-version",
common_utils::AZURE_DEFAULT_API_VERSION,
))
.respond_with(ResponseTemplate::new(200).set_body_raw(COMPLETION, "application/json"))
.expect(1)
.mount(&upstream)
.await;
let mut unversioned = call(&upstream.uri(), hi());
unversioned.litellm_params.clear();
run(&route, unversioned, &Recording::default())
.await
.unwrap();
}
#[rstest]
#[tokio::test]
async fn environment_references_resolve(route: ChatCompletionsRoute) {
let upstream = MockServer::start().await;
Mock::given(query_param("api-version", "2025-01-01-preview"))
.and(header("api-key", "sk-referenced"))
.respond_with(ResponseTemplate::new(200).set_body_raw(COMPLETION, "application/json"))
.expect(1)
.mount(&upstream)
.await;
let route = with_secrets(
route,
&[
("MY_AZURE_KEY", "sk-referenced"),
("MY_AZURE_VERSION", "2025-01-01-preview"),
],
);
let mut referenced = call(&upstream.uri(), hi());
referenced.api_key = Some("os.environ/MY_AZURE_KEY".into());
referenced
.litellm_params
.insert("api_version".into(), json!("os.environ/MY_AZURE_VERSION"));
run(&route, referenced, &Recording::default())
.await
.unwrap();
}
#[rstest]
#[case::configured_token_outranks_the_key(Some("sk-azure"), &[("azure_ad_token", "configured-token")], &[])]
#[case::environment_token(None, &[], &[("AZURE_AD_TOKEN", "configured-token")])]
#[tokio::test]
async fn an_ad_token_goes_as_a_bearer(
route: ChatCompletionsRoute,
#[case] api_key: Option<&str>,
#[case] params: &[(&str, &str)],
#[case] environment: &[(&str, &str)],
) {
let upstream = MockServer::start().await;
Mock::given(header("authorization", "Bearer configured-token"))
.respond_with(ResponseTemplate::new(200).set_body_raw(COMPLETION, "application/json"))
.expect(1)
.mount(&upstream)
.await;
let route = with_secrets(route, environment);
let mut tokened = call(&upstream.uri(), hi());
tokened.api_key = api_key.map(str::to_owned);
tokened.litellm_params.extend(
params
.iter()
.map(|(name, value)| ((*name).to_owned(), json!(value))),
);
run(&route, tokened, &Recording::default()).await.unwrap();
let requests = upstream.received_requests().await.unwrap();
assert!(!requests[0].headers.contains_key("api-key"));
}
#[rstest]
#[tokio::test]
async fn a_missing_credential_fails_before_any_request(route: ChatCompletionsRoute) {
let mut keyless = call("http://127.0.0.1:1", hi());
keyless.api_key = None;
assert_eq!(
run(&route, keyless, &Recording::default()).await.err(),
Some(Error::Auth(litellm_auth::Error::MissingApiKey {
provider: "Azure",
environment_variable: "AZURE_API_KEY",
}))
);
}
#[rstest]
#[tokio::test]
async fn an_unsupported_param_fails_before_any_request(route: ChatCompletionsRoute) {
let mut body = hi();
body["reasoning_effort"] = json!("high");
assert!(matches!(
run(
&route,
call("http://127.0.0.1:1", body),
&Recording::default()
)
.await,
Err(Error::InvalidRequest(_))
));
}
#[rstest]
#[tokio::test]
async fn a_content_filter_error_keeps_azures_body(route: ChatCompletionsRoute) {
let upstream = MockServer::start().await;
let body = r#"{"error":{"message":"The response was filtered","type":null,"param":"prompt","code":"content_filter","status":400,"innererror":{"code":"ResponsibleAIPolicyViolation","content_filter_result":{"hate":{"filtered":true,"severity":"high"}}}}}"#;
Mock::given(method("POST"))
.respond_with(ResponseTemplate::new(400).set_body_raw(body, "application/json"))
.mount(&upstream)
.await;
let error = run(&route, call(&upstream.uri(), hi()), &Recording::default())
.await
.err()
.expect("the upstream failure fails the call");
assert_eq!(
error,
Error::Transport(transport::Error::Http {
status: 400,
body: body.into()
})
);
}
}

1
litellm-rust/crates/inference-chat/src/handler.rs
View file

@ -29,6 +29,7 @@ pub(super) async fn execute(
config,
url,
body,
stream_shape: _,
optional_params,
environment,
secrets,

4
litellm-rust/crates/inference-chat/src/lib.rs
View file

@ -1,5 +1,8 @@
use litellm_host::observation::ObservationSender;
mod azure;
pub mod route;
#[cfg(test)]
mod test_support;
pub mod types;
pub use litellm_inference::RouteError as Error;
mod common_utils;
@ -7,6 +10,7 @@ pub mod constants;
pub(crate) mod handler;
pub mod openai;
mod prepare;
pub mod stream;
use litellm_llms_types::formats::chat_completions::ChatCompletionsResponse;
use prepare::{prepare_provider_request, resolve_request};

184
litellm-rust/crates/inference-chat/src/openai.rs
View file

@ -1,4 +1,4 @@
use std::{convert::Infallible, time::Duration};
use std::{convert::Infallible, sync::Arc, time::Duration};
use bytes::Bytes;
use futures_util::{StreamExt, stream::BoxStream};
@ -15,7 +15,10 @@ use litellm_host::{
use litellm_http::{outbound::OutboundRequest, request::truncate_error_body, transport};
use litellm_inference::{CallOptions, context::CallContext, provider::resolve_llm_provider};
use litellm_llms::{
base_llm::auth::resolve_auth,
base_llm::{
auth::{Headers, resolve_auth},
model_info::ModelInfo,
},
openai::chat::transformation::{OPENAI_CHAT_COMPLETIONS_CONFIG, is_stream},
};
use serde_json::{Map, Value};
@ -39,8 +42,10 @@ impl Protocol for OpenAiChatCompletions {
pub type OpenAiChatOutput = CallOutput<Bytes, (), Bytes, Error>;
/// A Chat Completions body bound for an OpenAI deployment. `body` is the caller's request;
/// `model` is the deployment's, with or without its `openai/` prefix.
/// A Chat Completions body bound for an OpenAI or Azure OpenAI deployment. `body` is the
/// caller's request; `model` is the deployment's, with or without its `openai/` or `azure/`
/// prefix.
#[derive(Default)]
pub struct OpenAiChatCall {
pub model: String,
pub body: Map<String, Value>,
@ -49,12 +54,35 @@ pub struct OpenAiChatCall {
pub custom_llm_provider: Option<String>,
pub extra_headers: Option<Map<String, Value>>,
pub timeout: Option<Duration>,
/// The deployment's other `litellm_params`: Azure reads its `api_version`, `base_model`
/// and Entra ID credentials from them.
pub litellm_params: Map<String, Value>,
/// `drop_params` from the deployment or the proxy settings.
pub drop_params: bool,
/// The model cost map, which Azure's param rules read.
pub model_info: Option<Arc<dyn ModelInfo>>,
}
/// A provider request ready to go out: where, with which headers and body, and who serves it.
pub(crate) struct PreparedCall {
pub(crate) model: String,
pub(crate) provider: &'static str,
pub(crate) url: String,
pub(crate) headers: Headers,
pub(crate) body: Map<String, Value>,
pub(crate) api_key: Option<String>,
}
impl ChatCompletionsRoute {
/// Whether the deployment is served by the OpenAI wire passthrough: OpenAI itself, or
/// Azure OpenAI.
pub fn serves_openai(model: &str, custom_llm_provider: Option<&str>) -> bool {
resolve_llm_provider(model, custom_llm_provider, "chat completions")
.is_ok_and(|resolved| resolved.provider == LlmProviders::Openai)
resolve_llm_provider(model, custom_llm_provider, "chat completions").is_ok_and(|resolved| {
matches!(
resolved.provider,
LlmProviders::Openai | LlmProviders::Azure
)
})
}
pub fn openai_machine(
@ -106,39 +134,35 @@ impl ChatCompletionsRoute {
stream,
outcome
))]
async fn run_openai(
pub(crate) async fn run_openai(
&self,
call: OpenAiChatCall,
context: CallContext<'_, impl Interceptors<Error>>,
) -> Result<OpenAiChatOutput, Error> {
let timeout = call
.timeout
.unwrap_or(Duration::from_secs(CHAT_COMPLETIONS_TIMEOUT_SECS));
litellm_inference::diagnostic::call(async {
let resolved = resolve_llm_provider(
&call.model,
call.custom_llm_provider.as_deref(),
"chat completions",
)?;
if resolved.provider != LlmProviders::Openai {
return Err(Error::InvalidProvider(
<&str>::from(resolved.provider).to_string(),
));
}
let model = resolved.model.to_string();
let provider = <&str>::from(LlmProviders::Openai);
let prepared = match resolved.provider {
LlmProviders::Openai => self.prepare_openai(model, call).await?,
LlmProviders::Azure => self.prepare_azure(model, call).await?,
other => return Err(Error::InvalidProvider(<&str>::from(other).to_string())),
};
let PreparedCall {
model,
provider,
url,
headers,
body,
api_key,
} = prepared;
litellm_inference::diagnostic::provider(&model, provider);
let config = &OPENAI_CHAT_COMPLETIONS_CONFIG;
let secrets = self
.secrets
.resolve(config.secret_names(call.api_key.as_deref(), call.api_base.as_deref()))
.await?;
let lookup = |key: &str| secrets.get(key);
let body = config.transform_request(&model, call.body)?;
let environment = config.validate_environment(
string_headers(call.extra_headers)?,
call.api_key.as_deref(),
&lookup,
)?;
let url = config.get_complete_url(call.api_base.as_deref(), &lookup);
let authenticated = resolve_auth(&self.auth, environment, &lookup).await?;
let request_context = RequestContext {
model: model.clone(),
custom_llm_provider: provider.to_string(),
@ -149,18 +173,16 @@ impl ChatCompletionsRoute {
.collect(),
),
secret_fields: Vec::new(),
api_key: call.api_key.map(litellm_auth::SecretValue::new),
api_key: api_key.map(litellm_auth::SecretValue::new),
};
let wire = WireRequest {
url,
headers,
body: Value::Object(body),
};
let wire = context
.interceptors
.before_provider_request(
WireRequest {
url,
headers: authenticated.headers,
body: Value::Object(body),
},
request_context,
)
.before_provider_request(wire, request_context)
.await?;
let stream = wire
.body
@ -168,9 +190,6 @@ impl ChatCompletionsRoute {
.map(is_stream)
.transpose()?
.unwrap_or(false);
let timeout = call
.timeout
.unwrap_or(Duration::from_secs(CHAT_COMPLETIONS_TIMEOUT_SECS));
let response = send(&self.http, wire, stream, timeout).await?;
let output = if stream {
CallOutput::Stream {
@ -200,6 +219,37 @@ impl ChatCompletionsRoute {
}
}
impl ChatCompletionsRoute {
async fn prepare_openai(
&self,
model: String,
call: OpenAiChatCall,
) -> Result<PreparedCall, Error> {
let config = &OPENAI_CHAT_COMPLETIONS_CONFIG;
let secrets = self
.secrets
.resolve(config.secret_names(call.api_key.as_deref(), call.api_base.as_deref()))
.await?;
let lookup = |key: &str| secrets.get(key);
let body = config.transform_request(&model, call.body)?;
let environment = config.validate_environment(
string_headers(call.extra_headers)?,
call.api_key.as_deref(),
&lookup,
)?;
let url = config.get_complete_url(call.api_base.as_deref(), &lookup);
let authenticated = resolve_auth(&self.auth, environment, &lookup).await?;
Ok(PreparedCall {
model,
provider: <&str>::from(LlmProviders::Openai),
url,
headers: authenticated.headers,
body,
api_key: call.api_key,
})
}
}
/// A non-streaming call keeps the usual total timeout. A stream has no total limit, so a
/// 30 minute generation is not cut off: `timeout` bounds the wait for the response head
/// here and each read after it in [`chunks`].
@ -254,11 +304,10 @@ fn silent_provider(timeout: Duration) -> Error {
#[cfg(test)]
mod tests {
use std::sync::{Arc, Mutex};
use std::sync::Arc;
use futures_util::TryStreamExt;
use litellm_host::interceptors::RawResponse;
use rstest::{fixture, rstest};
use rstest::rstest;
use serde_json::json;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use wiremock::{
@ -267,47 +316,11 @@ mod tests {
};
use super::*;
use crate::test_support::{Recording, route, run};
const COMPLETION: &str = r#"{"id":"chatcmpl-1","object":"chat.completion","created":1,"model":"gpt-6-astra","system_fingerprint":"fp","choices":[{"index":0,"message":{"role":"assistant","content":null,"tool_calls":[{"id":"call_1","type":"function","function":{"name":"f","arguments":"{}"}}]},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":3,"completion_tokens":2,"total_tokens":5}}"#;
const STREAM: &str = "data: {\"id\":\"c\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hi\"}}]}\n\ndata: {\"id\":\"c\",\"choices\":[],\"usage\":{\"prompt_tokens\":3,\"completion_tokens\":1,\"total_tokens\":4}}\n\ndata: [DONE]\n\n";
#[derive(Default)]
struct Recording {
contexts: Mutex<Vec<RequestContext>>,
raw: Mutex<Vec<String>>,
results: Mutex<Vec<ExecutionFacts>>,
}
impl Interceptors<Error> for Recording {
async fn before_provider_request(
&self,
wire: WireRequest,
context: RequestContext,
) -> Result<WireRequest, Error> {
self.contexts.lock().unwrap().push(context);
Ok(wire)
}
async fn after_provider_response(&self, raw: RawResponse) -> Result<(), Error> {
self.raw.lock().unwrap().push(raw.body);
Ok(())
}
async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), Error> {
self.results.lock().unwrap().push(facts);
Ok(())
}
}
#[fixture]
fn route() -> ChatCompletionsRoute {
ChatCompletionsRoute::new(
litellm_http::Client::plain_for_test(),
Arc::new(litellm_auth::AuthServices::default()),
litellm_inference_testing::no_secrets(),
)
}
fn call(api_base: &str, body: Value) -> OpenAiChatCall {
let Value::Object(body) = body else {
panic!("the body is an object");
@ -320,19 +333,10 @@ mod tests {
custom_llm_provider: None,
extra_headers: None,
timeout: Some(Duration::from_secs(5)),
..OpenAiChatCall::default()
}
}
async fn run(
route: &ChatCompletionsRoute,
call: OpenAiChatCall,
interceptors: &Recording,
) -> Result<OpenAiChatOutput, Error> {
route
.run_openai(call, CallContext::new(interceptors, CallOptions::default()))
.await
}
async fn drain(chunks: BoxStream<'static, Result<Bytes, Error>>) -> Result<Vec<u8>, Error> {
Ok(chunks.try_collect::<Vec<_>>().await?.concat())
}
@ -347,6 +351,8 @@ mod tests {
#[rstest]
#[case::openai_prefix("openai/gpt-6-astra", None, true)]
#[case::explicit_provider("gpt-6-astra", Some("openai"), true)]
#[case::azure("azure/gpt-4o", None, true)]
#[case::azure_ai("azure_ai/gpt-4o", None, false)]
#[case::openai_like("openai_like/gpt-6-astra", None, false)]
#[case::anthropic("anthropic/claude", None, false)]
#[case::unprefixed("gpt-6-astra", None, false)]

1
litellm-rust/crates/inference-chat/src/prepare.rs
View file

@ -121,6 +121,7 @@ pub(super) fn prepare_provider_request(
config,
url,
body: transformed.body,
stream_shape: transformed.stream_shape,
optional_params: request.optional_params,
environment,
secrets,

1246
litellm-rust/crates/inference-chat/src/stream.rs Normal file
View file

File diff suppressed because it is too large Load diff

61
litellm-rust/crates/inference-chat/src/test_support.rs Normal file
View file

@ -0,0 +1,61 @@
//! Helpers the passthrough tests of several providers share.
use std::sync::{Arc, Mutex};
use litellm_host::interceptors::{
ExecutionFacts, Interceptors, RawResponse, RequestContext, WireRequest,
};
use litellm_inference::{CallOptions, context::CallContext};
use rstest::fixture;
use crate::{
ChatCompletionsRoute, Error,
openai::{OpenAiChatCall, OpenAiChatOutput},
};
#[derive(Default)]
pub(crate) struct Recording {
pub(crate) contexts: Mutex<Vec<RequestContext>>,
pub(crate) raw: Mutex<Vec<String>>,
pub(crate) results: Mutex<Vec<ExecutionFacts>>,
}
impl Interceptors<Error> for Recording {
async fn before_provider_request(
&self,
wire: WireRequest,
context: RequestContext,
) -> Result<WireRequest, Error> {
self.contexts.lock().unwrap().push(context);
Ok(wire)
}
async fn after_provider_response(&self, raw: RawResponse) -> Result<(), Error> {
self.raw.lock().unwrap().push(raw.body);
Ok(())
}
async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), Error> {
self.results.lock().unwrap().push(facts);
Ok(())
}
}
#[fixture]
pub(crate) fn route() -> ChatCompletionsRoute {
ChatCompletionsRoute::new(
litellm_http::Client::plain_for_test(),
Arc::new(litellm_auth::AuthServices::default()),
litellm_inference_testing::no_secrets(),
)
}
pub(crate) async fn run(
route: &ChatCompletionsRoute,
call: OpenAiChatCall,
interceptors: &Recording,
) -> Result<OpenAiChatOutput, Error> {
route
.run_openai(call, CallContext::new(interceptors, CallOptions::default()))
.await
}

7
litellm-rust/crates/inference-chat/src/types.rs
View file

@ -2,7 +2,10 @@ use litellm_secrets::source::Secrets;
use std::time::Duration;
use litellm_auth::SecretValue;
use litellm_llms::base_llm::{auth::ValidatedEnvironment, chat::transformation::BaseConfig};
use litellm_llms::base_llm::{
auth::ValidatedEnvironment,
chat::{streaming::StreamShape, transformation::BaseConfig},
};
use litellm_llms_types::formats::chat_completions::ChatMessage;
use serde_json::{Map, Value};
@ -67,6 +70,8 @@ pub struct ProviderChatCompletionsRequest {
pub config: &'static dyn BaseConfig,
pub url: String,
pub body: Value,
/// How the config's stream parser reads a streamed answer to this request.
pub stream_shape: StreamShape,
/// The route's parameters before the provider transformation, reported to the host
/// beside the wire request.
pub optional_params: Map<String, Value>,

1
litellm-rust/crates/inference-chat/tests/live/openai.rs
View file

@ -47,6 +47,7 @@ fn call(extra: Value) -> OpenAiChatCall {
custom_llm_provider: None,
extra_headers: None,
timeout: Some(Duration::from_secs(60)),
..OpenAiChatCall::default()
}
}

3
litellm-rust/crates/inference/src/error.rs
View file

@ -85,6 +85,9 @@ impl From<LlmError> for RouteError {
LlmError::InvalidResponse(message) => Self::InvalidResponse(message),
LlmError::Unsupported(reason) => Self::Unsupported(reason),
LlmError::Auth(error) => Self::Auth(error),
LlmError::Provider { status, body } => {
Self::Transport(TransportError::Http { status, body })
}
}
}
}

29
litellm-rust/crates/llms-types/src/formats/chat_completions/response.rs
View file

@ -29,12 +29,17 @@ pub enum ChatCompletionThinkingBlock {
signature: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
cache_control: Option<Value>,
/// Which thinking block of the response a streamed piece belongs to.
#[serde(default, skip_serializing_if = "Option::is_none")]
index: Option<u64>,
},
RedactedThinking {
#[serde(default, skip_serializing_if = "Option::is_none")]
data: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
cache_control: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
index: Option<u64>,
},
}
@ -46,6 +51,21 @@ pub struct PromptTokensDetails {
pub cached_tokens: u64,
pub cache_creation_tokens: u64,
pub text_tokens: u64,
/// The cache writes again, under the name Python's `Usage` gives them on a stream.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cache_write_tokens: Option<u64>,
/// The cache writes split by lifetime (`ephemeral_5m_input_tokens`, `ephemeral_1h_input_tokens`).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cache_creation_token_details: Option<Value>,
}
#[macro_rules_attribute::apply(crate::wire_type)]
#[derive(Default)]
pub struct CompletionTokensDetails {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub reasoning_tokens: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub text_tokens: Option<u64>,
}
#[macro_rules_attribute::apply(crate::wire_type)]
@ -55,6 +75,15 @@ pub struct ChatCompletionsUsage {
pub completion_tokens: u64,
pub total_tokens: u64,
pub prompt_tokens_details: PromptTokensDetails,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub completion_tokens_details: Option<CompletionTokensDetails>,
/// Anthropic's server tool counts, `{"web_search_requests": n}`.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub server_tool_use: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cache_creation_input_tokens: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cache_read_input_tokens: Option<u64>,
}
#[macro_rules_attribute::apply(crate::wire_type)]

775
litellm-rust/crates/llms/src/anthropic/chat/handler.rs
View file

@ -1,150 +1,479 @@
use std::collections::HashMap;
//! Anthropic Messages stream events to Chat Completions chunks, as Python's
//! `anthropic/chat/handler.py` `ModelResponseIterator.chunk_parser` builds them.
//!
//! Each event becomes one provider chunk, empty ones included; the host's stream wrapper
//! (Python's `CustomStreamWrapper`) decides which of them the client sees.
//!
//! DIVERGENCE: not built here, since the chat route never asks for what produces them:
//! `json_mode` / `response_format` tool conversion, `web_search_calls` and
//! `code_interpreter_results` summaries, and a tool call's `caller`. The reasoning token
//! estimate Python makes from streamed thinking text (with its tokenizer) is not made, so
//! such a stream reports `completion_tokens_details` without counts unless Anthropic sends
//! `output_tokens_details.thinking_tokens`.
use litellm_core_utils::core_helpers::finish_reason_for;
use litellm_llms_types::formats::{
chat_completions::{
ChatCompletionChunk, ChatCompletionThinkingBlock, ChatCompletionToolCallChunk,
ChatCompletionsUsage,
ChatCompletionChunk, ChatCompletionDelta, ChatCompletionStreamingChoice,
ChatCompletionThinkingBlock, ChatCompletionToolCallChunk,
ChatCompletionToolCallFunctionChunk, ChatCompletionsUsage, CompletionTokensDetails,
PromptTokensDetails,
},
messages::streaming::{
MessagesContentBlock, MessagesContentBlockDelta, MessagesStreamEvent, MessagesStreamUsage,
MessagesContentBlock, MessagesContentBlockDelta, MessagesDelta, MessagesStreamEvent,
MessagesStreamUsage,
},
};
use serde_json::Value;
use serde_json::{Map, Value, json};
use crate::{
Error,
base_llm::{base_model_iterator::StreamTransformer, chat::streaming::StreamShape},
};
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum AnthropicJsonChunkType {
ValidJson,
AccumulatedJson,
}
/// Anthropic sends no status with an `error` event; Python raises it as a 500.
const STREAM_ERROR_STATUS: u16 = 500;
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum AnthropicContentBlockType {
Text,
ToolUse,
ServerToolUse,
Thinking,
RedactedThinking,
Compaction,
ToolResult(String),
Other(String),
}
#[derive(Clone, Debug, PartialEq)]
pub struct AnthropicContentBlockDeltaEvent {
pub index: u64,
pub delta: MessagesContentBlockDelta,
}
const TOOL_BLOCKS: [&str; 2] = ["tool_use", "server_tool_use"];
const THINKING_BLOCKS: [&str; 2] = ["thinking", "redacted_thinking"];
pub struct ModelResponseIterator {
pub content_blocks: Vec<AnthropicContentBlockDeltaEvent>,
pub tool_index: i64,
pub json_mode: bool,
pub speed: Option<String>,
pub tool_name_reverse_map: HashMap<String, String>,
pub response_id: String,
pub served_model: Option<String>,
pub is_response_format_tool: bool,
pub converted_response_format_tool: bool,
pub accumulated_json: String,
pub chunk_type: AnthropicJsonChunkType,
pub current_content_block_type: Option<AnthropicContentBlockType>,
pub web_search_results: Vec<Value>,
pub web_search_calls: HashMap<String, Value>,
pub compaction_blocks: Vec<Value>,
pub reasoning_content_chunks: Vec<String>,
pub server_tool_inputs: HashMap<String, Value>,
pub tool_results: Vec<Value>,
pub current_server_tool_id: Option<String>,
pub container_id: Option<String>,
shape: StreamShape,
/// The deltas of the block in progress, other than text.
content_blocks: Vec<MessagesContentBlockDelta>,
tool_index: i64,
thinking_block_index: i64,
served_model: Option<String>,
current_content_block_type: Option<String>,
/// `message_start` usage overlaid with every later `message_delta` usage.
cumulative_usage: Map<String, Value>,
streamed_reasoning: bool,
web_search_results: Vec<Value>,
tool_results: Vec<Value>,
compaction_blocks: Vec<Value>,
}
/// What one event contributes to its chunk.
#[derive(Default)]
struct Parsed {
text: String,
tool_use: Option<ChatCompletionToolCallChunk>,
finish_reason: Option<String>,
usage: Option<ChatCompletionsUsage>,
provider_specific_fields: Map<String, Value>,
reasoning_content: Option<String>,
thinking_blocks: Vec<ChatCompletionThinkingBlock>,
}
impl ModelResponseIterator {
pub fn new(_shape: StreamShape) -> Self {
todo!()
pub fn new(shape: StreamShape) -> Self {
Self {
shape,
content_blocks: Vec::new(),
tool_index: -1,
thinking_block_index: -1,
served_model: None,
current_content_block_type: None,
cumulative_usage: Map::new(),
streamed_reasoning: false,
web_search_results: Vec::new(),
tool_results: Vec::new(),
compaction_blocks: Vec::new(),
}
}
/// Whether the tool block so far streamed only empty argument fragments.
pub fn check_empty_tool_call_args(&self) -> bool {
todo!()
match self.content_blocks.first() {
None
| Some(
MessagesContentBlockDelta::TextDelta { .. }
| MessagesContentBlockDelta::ThinkingDelta { .. },
) => false,
Some(_) => self.content_blocks.iter().all(|delta| match delta {
MessagesContentBlockDelta::InputJsonDelta { partial_json } => {
partial_json.is_empty()
}
_ => true,
}),
}
}
pub fn handle_usage(&mut self, _usage: MessagesStreamUsage) -> ChatCompletionsUsage {
todo!()
fn handle_usage(&mut self, usage: &MessagesStreamUsage) -> ChatCompletionsUsage {
if let Ok(Value::Object(fields)) = serde_json::to_value(usage) {
self.cumulative_usage.extend(fields);
}
anthropic_usage(&self.cumulative_usage, self.streamed_reasoning)
}
pub fn handle_content_block_delta(
fn thinking_index(&self) -> Option<u64> {
u64::try_from(self.thinking_block_index).ok()
}
fn tool_call(
&self,
id: Option<String>,
name: Option<String>,
arguments: String,
) -> ChatCompletionToolCallChunk {
ChatCompletionToolCallChunk {
id,
tool_type: "function".into(),
function: ChatCompletionToolCallFunctionChunk {
name,
arguments,
provider_specific_fields: None,
},
index: self.tool_index,
}
}
fn content_block_delta(&mut self, delta: MessagesContentBlockDelta) -> Parsed {
let mut parsed = Parsed::default();
if let MessagesContentBlockDelta::TextDelta { text } = delta {
parsed.text = text;
return parsed;
}
self.content_blocks.push(delta.clone());
match delta {
MessagesContentBlockDelta::TextDelta { .. } => {}
MessagesContentBlockDelta::InputJsonDelta { partial_json } => {
// A `web_search_tool_result` block streams JSON too, and is not a tool call.
if self
.current_content_block_type
.as_deref()
.is_some_and(|block| TOOL_BLOCKS.contains(&block))
{
parsed.tool_use = Some(self.tool_call(None, None, partial_json));
}
}
MessagesContentBlockDelta::Citations { citation } => {
parsed
.provider_specific_fields
.insert("citation".into(), citation);
}
MessagesContentBlockDelta::ThinkingDelta { thinking } => {
if !thinking.is_empty() {
self.streamed_reasoning = true;
parsed.reasoning_content = Some(thinking.clone());
self.thinking_blocks(
&mut parsed,
ChatCompletionThinkingBlock::Thinking {
thinking: Some(thinking),
signature: None,
cache_control: None,
index: self.thinking_index(),
},
);
}
}
MessagesContentBlockDelta::SignatureDelta { signature } => {
if !signature.is_empty() {
self.thinking_blocks(
&mut parsed,
ChatCompletionThinkingBlock::Thinking {
thinking: Some(String::new()),
signature: Some(signature),
cache_control: None,
index: self.thinking_index(),
},
);
parsed.reasoning_content.get_or_insert_with(String::new);
}
}
MessagesContentBlockDelta::CompactionDelta { content } => {
parsed.provider_specific_fields.insert(
"compaction_delta".into(),
json!({"type": "compaction_delta", "content": content}),
);
}
}
parsed
}
/// A thinking block goes out both as `thinking_blocks` and under the same name in the
/// delta's `provider_specific_fields`.
fn thinking_blocks(&self, parsed: &mut Parsed, block: ChatCompletionThinkingBlock) {
parsed.thinking_blocks = vec![block];
if let Ok(blocks) = serde_json::to_value(&parsed.thinking_blocks) {
parsed
.provider_specific_fields
.insert("thinking_blocks".into(), blocks);
}
}
fn content_block_start(&mut self, block: MessagesContentBlock) -> Parsed {
let mut parsed = Parsed::default();
self.content_blocks.clear();
let block_type = block.block_type.clone();
self.current_content_block_type = Some(block_type.clone());
if THINKING_BLOCKS.contains(&block_type.as_str()) {
self.thinking_block_index += 1;
}
match block_type.as_str() {
"text" => parsed.text = block.text.unwrap_or_default(),
"tool_use" | "server_tool_use" => {
self.tool_index += 1;
let name = block.name.map(|name| {
self.shape
.tool_name_reverse_map
.get(&name)
.cloned()
.unwrap_or(name)
});
// The arguments arrive in the deltas; `input` here is always `{}`.
parsed.tool_use = Some(self.tool_call(block.id, name, String::new()));
}
"redacted_thinking" => self.thinking_blocks(
&mut parsed,
ChatCompletionThinkingBlock::RedactedThinking {
data: block.data,
cache_control: None,
index: self.thinking_index(),
},
),
"compaction" => {
let content = block.content.clone().unwrap_or_else(|| json!(""));
self.compaction_blocks.push(block_value(&block));
parsed.provider_specific_fields.insert(
"compaction_blocks".into(),
Value::Array(self.compaction_blocks.clone()),
);
parsed.provider_specific_fields.insert(
"compaction_start".into(),
json!({"type": "compaction", "content": content}),
);
}
result if result.ends_with("_tool_result") => {
if matches!(result, "web_search_tool_result" | "web_fetch_tool_result") {
self.web_search_results.push(block_value(&block));
parsed.provider_specific_fields.insert(
"web_search_results".into(),
Value::Array(self.web_search_results.clone()),
);
} else {
self.tool_results.push(block_value(&block));
parsed.provider_specific_fields.insert(
"tool_results".into(),
Value::Array(self.tool_results.clone()),
);
}
}
_ => {}
}
parsed
}
fn content_block_stop(&mut self) -> Parsed {
let mut parsed = Parsed::default();
let in_tool = self
.current_content_block_type
.as_deref()
.is_some_and(|block| TOOL_BLOCKS.contains(&block));
if in_tool && self.check_empty_tool_call_args() {
parsed.tool_use = Some(self.tool_call(None, None, "{}".into()));
}
self.current_content_block_type = None;
parsed
}
fn message_delta(
&mut self,
_index: u64,
_delta: MessagesContentBlockDelta,
) -> (
String,
Option<ChatCompletionToolCallChunk>,
Vec<ChatCompletionThinkingBlock>,
Option<Value>,
Option<String>,
) {
todo!()
}
pub fn handle_content_block_start(
&mut self,
_index: u64,
_content_block: MessagesContentBlock,
) -> Result<ChatCompletionChunk, Error> {
todo!()
}
pub fn handle_json_mode_chunk(
&mut self,
_text: String,
_tool_use: Option<ChatCompletionToolCallChunk>,
) -> (String, Option<ChatCompletionToolCallChunk>) {
todo!()
}
pub fn handle_accumulated_json_chunk(
&mut self,
_data: &str,
_is_final: bool,
) -> Result<Option<ChatCompletionChunk>, Error> {
todo!()
}
pub fn handle_redacted_thinking_content(
&mut self,
_content_block: &MessagesContentBlock,
) -> Vec<ChatCompletionThinkingBlock> {
todo!()
}
pub fn web_search_call_snapshot(&self) -> HashMap<String, Value> {
todo!()
}
pub fn complete_web_search_call(&mut self, _result: Value) {
todo!()
}
pub fn build_code_interpreter_results(&self) -> Vec<Value> {
todo!()
}
pub fn handle_message_delta(
&mut self,
_event: MessagesStreamEvent,
) -> (Option<String>, Option<ChatCompletionsUsage>, Option<Value>) {
todo!()
delta: MessagesDelta,
usage: Option<MessagesStreamUsage>,
) -> Parsed {
let mut parsed = Parsed {
finish_reason: Some(
finish_reason_for(delta.stop_reason.as_deref().unwrap_or("stop")).to_string(),
),
usage: usage.map(|usage| self.handle_usage(&usage)),
..Parsed::default()
};
// Python's `refusal_provider_specific_fields`.
if let Some(details) = delta.stop_details {
parsed
.provider_specific_fields
.insert("stop_reason".into(), json!(delta.stop_reason));
parsed
.provider_specific_fields
.insert("stop_details".into(), details);
} else if delta.stop_reason.as_deref() == Some("refusal") {
parsed
.provider_specific_fields
.insert("stop_reason".into(), json!("refusal"));
}
if let Some(container) = delta.container.filter(|container| !container.is_null()) {
parsed
.provider_specific_fields
.insert("container".into(), container);
}
parsed
}
pub fn chunk_parser(
&mut self,
_event: MessagesStreamEvent,
event: MessagesStreamEvent,
) -> Result<ChatCompletionChunk, Error> {
todo!()
let parsed = match event {
MessagesStreamEvent::ContentBlockDelta { delta, .. } => self.content_block_delta(delta),
MessagesStreamEvent::ContentBlockStart { content_block, .. } => {
self.content_block_start(content_block)
}
MessagesStreamEvent::ContentBlockStop { .. } => self.content_block_stop(),
MessagesStreamEvent::MessageDelta { delta, usage, .. } => {
self.message_delta(delta, usage)
}
MessagesStreamEvent::MessageStart { message } => {
self.served_model = Some(message.model);
Parsed {
usage: Some(self.handle_usage(&message.usage)),
..Parsed::default()
}
}
MessagesStreamEvent::MessageStop { .. } | MessagesStreamEvent::Ping => {
Parsed::default()
}
MessagesStreamEvent::Error { error } => {
let body = json!({"type": "error", "error": error});
return Err(Error::Provider {
status: STREAM_ERROR_STATUS,
body: body.to_string(),
});
}
};
Ok(self.chunk(parsed))
}
fn chunk(&self, parsed: Parsed) -> ChatCompletionChunk {
let Parsed {
text,
tool_use,
finish_reason,
usage,
provider_specific_fields,
reasoning_content,
thinking_blocks,
} = parsed;
ChatCompletionChunk {
model: self.served_model.clone(),
choices: vec![ChatCompletionStreamingChoice {
index: 0,
delta: ChatCompletionDelta {
content: Some(text),
tool_calls: tool_use.map(|tool| vec![tool]),
reasoning_content,
thinking_blocks: (!thinking_blocks.is_empty()).then_some(thinking_blocks),
provider_specific_fields: (!provider_specific_fields.is_empty())
.then_some(provider_specific_fields),
..ChatCompletionDelta::default()
},
finish_reason,
logprobs: None,
}],
usage,
..ChatCompletionChunk::default()
}
}
}
fn block_value(block: &MessagesContentBlock) -> Value {
serde_json::to_value(block).unwrap_or(Value::Null)
}
fn count(usage: &Map<String, Value>, field: &str) -> u64 {
usage.get(field).and_then(Value::as_u64).unwrap_or(0)
}
/// Python's `AnthropicConfig.calculate_usage` over a stream's cumulative usage.
pub fn anthropic_usage(
usage: &Map<String, Value>,
streamed_reasoning: bool,
) -> ChatCompletionsUsage {
let iterations = usage
.get("iterations")
.and_then(Value::as_array)
.filter(|iterations| !iterations.is_empty());
let (input, completion, cache_creation, cache_read) = match iterations {
Some(iterations) => {
let sum = |field: &str| {
iterations
.iter()
.filter_map(|iteration| iteration.get(field)?.as_u64())
.sum::<u64>()
};
(
sum("input_tokens"),
sum("output_tokens"),
sum("cache_creation_input_tokens"),
sum("cache_read_input_tokens"),
)
}
None => (
count(usage, "input_tokens"),
count(usage, "output_tokens"),
count(usage, "cache_creation_input_tokens"),
count(usage, "cache_read_input_tokens"),
),
};
let prompt = input + cache_creation + cache_read;
let thinking_tokens = usage
.get("output_tokens_details")
.and_then(|details| details.get("thinking_tokens"))
.and_then(Value::as_u64);
let completion_tokens_details = match thinking_tokens {
Some(thinking) => {
let reasoning = thinking.min(completion);
CompletionTokensDetails {
reasoning_tokens: Some(reasoning),
text_tokens: Some(completion - reasoning),
}
}
None if streamed_reasoning => CompletionTokensDetails::default(),
None => CompletionTokensDetails {
reasoning_tokens: Some(0),
text_tokens: Some(completion),
},
};
let cache_creation_token_details =
usage
.get("cache_creation")
.and_then(Value::as_object)
.map(|details| {
Value::Object(
["ephemeral_5m_input_tokens", "ephemeral_1h_input_tokens"]
.into_iter()
.filter_map(|field| Some((field.to_owned(), details.get(field)?.clone())))
.filter(|(_, value)| !value.is_null())
.collect(),
)
});
let server_tool_use = usage
.get("server_tool_use")
.and_then(Value::as_object)
.map(|tools| {
["web_search_requests", "tool_search_requests"]
.into_iter()
.filter_map(|field| Some((field.to_owned(), tools.get(field)?.clone())))
.filter(|(_, value)| !value.is_null())
.collect::<Map<String, Value>>()
})
.filter(|tools| !tools.is_empty())
.map(Value::Object);
ChatCompletionsUsage {
prompt_tokens: prompt,
completion_tokens: completion,
total_tokens: prompt + completion,
prompt_tokens_details: PromptTokensDetails {
cached_tokens: cache_read,
cache_creation_tokens: cache_creation,
text_tokens: input,
cache_write_tokens: Some(cache_creation),
cache_creation_token_details,
},
completion_tokens_details: Some(completion_tokens_details),
server_tool_use,
cache_creation_input_tokens: Some(cache_creation),
cache_read_input_tokens: Some(cache_read),
}
}
@ -153,11 +482,221 @@ impl StreamTransformer for ModelResponseIterator {
type Output = ChatCompletionChunk;
type Error = Error;
fn transform(&mut self, _input: Self::Input) -> Result<Vec<Self::Output>, Self::Error> {
todo!()
fn transform(&mut self, input: Self::Input) -> Result<Vec<Self::Output>, Self::Error> {
self.chunk_parser(input).map(|chunk| vec![chunk])
}
fn finish(&mut self) -> Result<Vec<Self::Output>, Self::Error> {
todo!()
Ok(Vec::new())
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use serde_json::json;
use super::*;
fn events(wire: Value) -> Vec<MessagesStreamEvent> {
serde_json::from_value(wire).unwrap()
}
fn parse(wire: Value) -> Vec<Value> {
let mut iterator = ModelResponseIterator::new(StreamShape::default());
events(wire)
.into_iter()
.map(|event| serde_json::to_value(iterator.chunk_parser(event).unwrap()).unwrap())
.collect()
}
fn delta(chunk: &Value) -> &Value {
&chunk["choices"][0]["delta"]
}
#[rstest]
fn text_deltas_become_content_and_message_start_names_the_model() {
let chunks = parse(json!([
{"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": "claude-x-1", "content": [], "stop_reason": null, "stop_sequence": null, "usage": {"input_tokens": 3, "output_tokens": 1}}},
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Hi"}},
{"type": "ping"}
]));
assert_eq!(chunks[0]["model"], "claude-x-1");
assert_eq!(chunks[0]["usage"]["prompt_tokens"], 3);
assert_eq!(delta(&chunks[1]), &json!({"content": ""}));
assert_eq!(delta(&chunks[2]), &json!({"content": "Hi"}));
assert_eq!(chunks[2]["model"], "claude-x-1");
assert_eq!(delta(&chunks[3]), &json!({"content": ""}));
}
#[rstest]
fn tool_calls_stream_their_name_then_partial_arguments() {
let chunks = parse(json!([
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
{"type": "content_block_stop", "index": 0},
{"type": "content_block_start", "index": 1, "content_block": {"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {}}},
{"type": "content_block_delta", "index": 1, "delta": {"type": "input_json_delta", "partial_json": "{\"city\": \"Pa"}},
{"type": "content_block_delta", "index": 1, "delta": {"type": "input_json_delta", "partial_json": "ris\"}"}},
{"type": "content_block_stop", "index": 1},
{"type": "content_block_start", "index": 2, "content_block": {"type": "tool_use", "id": "toolu_2", "name": "now", "input": {}}},
{"type": "content_block_delta", "index": 2, "delta": {"type": "input_json_delta", "partial_json": ""}},
{"type": "content_block_stop", "index": 2},
{"type": "content_block_start", "index": 3, "content_block": {"type": "tool_use", "id": "toolu_3", "name": "none", "input": {}}},
{"type": "content_block_stop", "index": 3}
]));
let tool = |chunk: &Value| delta(chunk)["tool_calls"].clone();
assert_eq!(delta(&chunks[1]), &json!({"content": ""}));
assert_eq!(
tool(&chunks[2]),
json!([{"id": "toolu_1", "type": "function", "function": {"name": "get_weather", "arguments": ""}, "index": 0}])
);
assert_eq!(
tool(&chunks[3]),
json!([{"type": "function", "function": {"arguments": "{\"city\": \"Pa"}, "index": 0}])
);
assert_eq!(tool(&chunks[4])[0]["function"]["arguments"], "ris\"}");
assert_eq!(delta(&chunks[5]), &json!({"content": ""}));
assert_eq!(tool(&chunks[6])[0]["index"], 1);
// Only empty fragments streamed: the block closes with `{}`.
assert_eq!(
tool(&chunks[8]),
json!([{"type": "function", "function": {"arguments": "{}"}, "index": 1}])
);
// No fragment at all: Python's check sees no deltas and adds nothing.
assert_eq!(delta(&chunks[10]), &json!({"content": ""}));
}
#[rstest]
fn a_reverse_mapped_tool_name_is_restored() {
let mut shape = StreamShape::default();
shape
.tool_name_reverse_map
.insert("ns_tool".into(), "ns/tool".into());
let mut iterator = ModelResponseIterator::new(shape);
let chunk = iterator
.chunk_parser(events(json!([{"type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "id": "t", "name": "ns_tool", "input": {}}}])).remove(0))
.unwrap();
let tool = &chunk.choices[0].delta.tool_calls.as_ref().unwrap()[0];
assert_eq!(tool.function.name.as_deref(), Some("ns/tool"));
}
#[rstest]
fn thinking_streams_as_reasoning_content_and_thinking_blocks() {
let chunks = parse(json!([
{"type": "content_block_start", "index": 0, "content_block": {"type": "thinking", "thinking": "", "signature": ""}},
{"type": "content_block_delta", "index": 0, "delta": {"type": "thinking_delta", "thinking": "Let me think"}},
{"type": "content_block_delta", "index": 0, "delta": {"type": "signature_delta", "signature": "sig"}},
{"type": "content_block_stop", "index": 0},
{"type": "content_block_start", "index": 1, "content_block": {"type": "redacted_thinking", "data": "enc"}}
]));
let thinking = json!([{"type": "thinking", "thinking": "Let me think", "index": 0}]);
assert_eq!(
delta(&chunks[1]),
&json!({"content": "", "reasoning_content": "Let me think", "thinking_blocks": thinking, "provider_specific_fields": {"thinking_blocks": thinking}})
);
let signature =
json!([{"type": "thinking", "thinking": "", "signature": "sig", "index": 0}]);
assert_eq!(
delta(&chunks[2]),
&json!({"content": "", "reasoning_content": "", "thinking_blocks": signature, "provider_specific_fields": {"thinking_blocks": signature}})
);
let redacted = json!([{"type": "redacted_thinking", "data": "enc", "index": 1}]);
assert_eq!(
delta(&chunks[4]),
&json!({"content": "", "thinking_blocks": redacted, "provider_specific_fields": {"thinking_blocks": redacted}})
);
}
#[rstest]
#[case::end_turn("end_turn", "stop")]
#[case::max_tokens("max_tokens", "length")]
#[case::tool_use("tool_use", "tool_calls")]
#[case::stop_sequence("stop_sequence", "stop")]
#[case::refusal("refusal", "content_filter")]
fn message_delta_maps_the_stop_reason(#[case] stop_reason: &str, #[case] expected: &str) {
let chunks = parse(
json!([{"type": "message_delta", "delta": {"stop_reason": stop_reason}, "usage": {"output_tokens": 2}}]),
);
assert_eq!(chunks[0]["choices"][0]["finish_reason"], expected);
}
#[rstest]
fn a_refusal_keeps_its_stop_reason_in_the_provider_fields() {
let chunks = parse(json!([{"type": "message_delta", "delta": {"stop_reason": "refusal"}}]));
assert_eq!(
delta(&chunks[0])["provider_specific_fields"],
json!({"stop_reason": "refusal"})
);
}
#[rstest]
fn usage_accumulates_across_start_and_delta_with_cache_tokens() {
let chunks = parse(json!([
{"type": "message_start", "message": {"id": "m", "type": "message", "role": "assistant", "model": "x", "content": [], "stop_reason": null, "stop_sequence": null,
"usage": {"input_tokens": 10, "cache_creation_input_tokens": 6, "cache_read_input_tokens": 20, "output_tokens": 1,
"cache_creation": {"ephemeral_5m_input_tokens": 2, "ephemeral_1h_input_tokens": 4}}}},
{"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 12, "server_tool_use": {"web_search_requests": 2}}}
]));
assert_eq!(
chunks[1]["usage"],
json!({
"completion_tokens": 12, "prompt_tokens": 36, "total_tokens": 48,
"completion_tokens_details": {"reasoning_tokens": 0, "text_tokens": 12},
"prompt_tokens_details": {"cached_tokens": 20, "text_tokens": 10, "cache_write_tokens": 6, "cache_creation_tokens": 6,
"cache_creation_token_details": {"ephemeral_5m_input_tokens": 2, "ephemeral_1h_input_tokens": 4}},
"server_tool_use": {"web_search_requests": 2},
"cache_creation_input_tokens": 6, "cache_read_input_tokens": 20
})
);
}
#[rstest]
fn reported_thinking_tokens_split_the_completion() {
let usage = anthropic_usage(
json!({"input_tokens": 1, "output_tokens": 10, "output_tokens_details": {"thinking_tokens": 4}})
.as_object()
.unwrap(),
true,
);
assert_eq!(
usage.completion_tokens_details,
Some(CompletionTokensDetails {
reasoning_tokens: Some(4),
text_tokens: Some(6)
})
);
}
#[rstest]
fn an_error_event_fails_with_the_providers_error() {
let mut iterator = ModelResponseIterator::new(StreamShape::default());
let error = iterator
.chunk_parser(events(json!([{"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}])).remove(0))
.unwrap_err();
let Error::Provider { status, body } = error else {
panic!("an error event is the provider's error, got {error:?}");
};
assert_eq!(status, 500);
assert_eq!(
serde_json::from_str::<Value>(&body).unwrap(),
json!({"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}})
);
}
#[rstest]
fn citations_and_server_tool_results_go_to_provider_fields() {
let chunks = parse(json!([
{"type": "content_block_delta", "index": 0, "delta": {"type": "citations_delta", "citation": {"type": "char_location"}}},
{"type": "content_block_start", "index": 1, "content_block": {"type": "web_search_tool_result", "tool_use_id": "s", "content": []}}
]));
assert_eq!(
delta(&chunks[0])["provider_specific_fields"],
json!({"citation": {"type": "char_location"}})
);
assert_eq!(
delta(&chunks[1])["provider_specific_fields"]["web_search_results"][0]["tool_use_id"],
"s"
);
}
}

11
litellm-rust/crates/llms/src/anthropic/chat/transformation.rs
View file

@ -27,7 +27,8 @@ use crate::{
streaming::{ChatStream, StreamShape},
transformation::{
BaseConfig, Headers, ProviderChatRequestData, ProviderChatResponseData,
Unsupported, ValidatedEnvironment, unsupported_message, unsupported_param,
STREAM_PARAM, Unsupported, ValidatedEnvironment, unsupported_message,
unsupported_param,
},
},
messages::streaming::event_stream as messages_event_stream,
@ -201,6 +202,14 @@ impl BaseConfig for AnthropicConfig {
))
}
/// Python sends `stream: true` in the body, after the other params.
fn streaming_request(&self, url: String, mut body: Value) -> Result<(String, Value), Error> {
if let Value::Object(fields) = &mut body {
fields.insert(STREAM_PARAM.into(), Value::Bool(true));
}
Ok((url, body))
}
fn default_headers(&self) -> &'static [(&'static str, &'static str)] {
DEFAULT_HEADERS
}

1
litellm-rust/crates/llms/src/azure/chat/mod.rs Normal file
View file

@ -0,0 +1 @@
pub mod transformation;

1432
litellm-rust/crates/llms/src/azure/chat/transformation.rs Normal file
View file

File diff suppressed because it is too large Load diff

325
litellm-rust/crates/llms/src/azure/common_utils.rs Normal file
View file

@ -0,0 +1,325 @@
//! Endpoint, API version and credential resolution for Azure OpenAI, as
//! `litellm/llms/azure/common_utils.py`, `litellm/main.py` and the `AzureOpenAI` client of the
//! OpenAI SDK resolve them.
use litellm_auth::{CredentialPlacement, SecretValue};
use litellm_core_utils::settings::resolve_non_empty;
use crate::{
Error,
base_llm::auth::{AuthScheme, Headers, ValidatedEnvironment},
};
/// `litellm.AZURE_DEFAULT_API_VERSION`.
pub const AZURE_DEFAULT_API_VERSION: &str = "2025-02-01-preview";
pub const AZURE_API_BASE_ENV: &str = "AZURE_API_BASE";
/// The SDK's own fallback when LiteLLM passes no endpoint.
pub const AZURE_OPENAI_ENDPOINT_ENV: &str = "AZURE_OPENAI_ENDPOINT";
pub const AZURE_API_VERSION_ENV: &str = "AZURE_API_VERSION";
pub const AZURE_OPENAI_API_KEY_ENV: &str = "AZURE_OPENAI_API_KEY";
pub const AZURE_API_KEY_ENV: &str = "AZURE_API_KEY";
const API_KEY_HEADER: &str = "api-key";
/// The environment variables a call may fall back to, besides the credential ones
/// `litellm_auth_azure::SECRET_NAMES` lists.
pub const SECRET_NAMES: &[&str] = &[
AZURE_API_BASE_ENV,
AZURE_OPENAI_ENDPOINT_ENV,
AZURE_API_VERSION_ENV,
AZURE_OPENAI_API_KEY_ENV,
AZURE_API_KEY_ENV,
];
/// `BaseAzureLLM._is_azure_v1_api_version`: these versions use the `/openai/v1/` surface
/// through the plain OpenAI client.
pub fn is_v1_api_version(api_version: &str) -> bool {
matches!(api_version, "preview" | "latest" | "v1")
}
/// The deployment's `api_version`, else `AZURE_API_VERSION`, else LiteLLM's default.
pub fn resolve_api_version(
api_version: Option<&str>,
lookup: &dyn Fn(&str) -> Option<String>,
) -> String {
resolve_non_empty(api_version, lookup, &[AZURE_API_VERSION_ENV])
.unwrap_or_else(|| AZURE_DEFAULT_API_VERSION.into())
}
/// The deployment's `api_base`, else `AZURE_API_BASE`, else the SDK's `AZURE_OPENAI_ENDPOINT`.
pub fn resolve_api_base(
api_base: Option<&str>,
lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
resolve_non_empty(
api_base,
lookup,
&[AZURE_API_BASE_ENV, AZURE_OPENAI_ENDPOINT_ENV],
)
.ok_or(Error::Auth(litellm_auth::Error::MissingApiBase {
provider: "Azure",
guidance: "Set `api_base` or the AZURE_API_BASE environment variable. Expected format: https://<resource-name>.openai.azure.com",
}))
}
/// The deployment's key, else `AZURE_OPENAI_API_KEY`, else `AZURE_API_KEY`, as `main.py`
/// reads them.
pub fn resolve_api_key(
api_key: Option<&str>,
lookup: &dyn Fn(&str) -> Option<String>,
) -> Option<String> {
resolve_non_empty(
api_key,
lookup,
&[AZURE_OPENAI_API_KEY_ENV, AZURE_API_KEY_ENV],
)
}
/// The Chat Completions URL the SDK builds for `deployment`:
///
/// - `v1`, `latest` and `preview` go to `{api_base}/openai/v1/chat/completions` with no
/// `api-version`, as the plain OpenAI client LiteLLM switches to for them does;
/// - an `api_base` already naming `/openai/deployments/...` is the client's `base_url`, so only
/// `/chat/completions` is added;
/// - any other `api_base` is the resource endpoint, under which the SDK puts
/// `/openai/deployments/{deployment}`.
///
/// A trailing slash on `api_base` is dropped, as `azure_endpoint.rstrip('/')` does.
pub fn complete_url(api_base: &str, deployment: &str, api_version: &str) -> Result<String, Error> {
if api_base.contains("gateway.ai.cloudflare.com") {
return Err(Error::Unsupported(
"Azure deployments behind a Cloudflare AI Gateway api_base",
));
}
let base = api_base.trim_end_matches('/');
if is_v1_api_version(api_version) {
return Ok(format!("{base}/openai/v1/chat/completions"));
}
let version: String = url::form_urlencoded::byte_serialize(api_version.as_bytes()).collect();
Ok(match base.contains("/openai/deployments") {
true => format!("{base}/chat/completions?api-version={version}"),
false => {
format!("{base}/openai/deployments/{deployment}/chat/completions?api-version={version}")
}
})
}
/// How a call authenticates once its credential is known.
#[derive(Clone, Debug)]
pub enum AzureCredential {
ApiKey(SecretValue),
/// An Entra ID access token, given or acquired.
AdToken(SecretValue),
}
/// A key goes in `api-key` and a token in `Authorization: Bearer`, as `AzureOpenAI` sends
/// them. On the v1 surface LiteLLM hands either to the plain OpenAI client as its `api_key`,
/// which sends it as a bearer.
pub fn environment(
headers: Headers,
credential: AzureCredential,
v1: bool,
) -> ValidatedEnvironment {
let (placement, secret) = match credential {
AzureCredential::ApiKey(secret) if !v1 => {
(CredentialPlacement::Header(API_KEY_HEADER), secret)
}
AzureCredential::ApiKey(secret) | AzureCredential::AdToken(secret) => {
(CredentialPlacement::Bearer, secret)
}
};
ValidatedEnvironment {
headers: litellm_http::request::with_default_headers(
headers,
&[("content-type", "application/json")],
),
auth: AuthScheme::Credential { placement, secret },
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
fn no_env(_: &str) -> Option<String> {
None
}
fn env(name: &str) -> Option<String> {
match name {
AZURE_API_BASE_ENV => Some("https://env.openai.azure.com".into()),
AZURE_API_VERSION_ENV => Some("2024-10-21".into()),
AZURE_API_KEY_ENV => Some("env-key".into()),
_ => None,
}
}
fn endpoint_only(name: &str) -> Option<String> {
(name == AZURE_OPENAI_ENDPOINT_ENV).then(|| "https://sdk.openai.azure.com".into())
}
#[rstest]
#[case::resource(
"https://res.openai.azure.com",
"2024-10-21",
"https://res.openai.azure.com/openai/deployments/gpt-4o/chat/completions?api-version=2024-10-21"
)]
#[case::trailing_slash(
"https://res.openai.azure.com/",
"2024-10-21",
"https://res.openai.azure.com/openai/deployments/gpt-4o/chat/completions?api-version=2024-10-21"
)]
#[case::preview_version(
"https://res.openai.azure.com",
"2025-02-01-preview",
"https://res.openai.azure.com/openai/deployments/gpt-4o/chat/completions?api-version=2025-02-01-preview"
)]
#[case::base_url_with_deployment(
"https://res.openai.azure.com/openai/deployments/other/",
"2024-10-21",
"https://res.openai.azure.com/openai/deployments/other/chat/completions?api-version=2024-10-21"
)]
#[case::v1(
"https://res.openai.azure.com/",
"v1",
"https://res.openai.azure.com/openai/v1/chat/completions"
)]
#[case::latest(
"https://res.openai.azure.com",
"latest",
"https://res.openai.azure.com/openai/v1/chat/completions"
)]
#[case::preview(
"https://res.openai.azure.com",
"preview",
"https://res.openai.azure.com/openai/v1/chat/completions"
)]
fn the_url_follows_the_sdk(
#[case] api_base: &str,
#[case] api_version: &str,
#[case] expected: &str,
) {
assert_eq!(
complete_url(api_base, "gpt-4o", api_version).unwrap(),
expected
);
}
#[rstest]
fn a_cloudflare_gateway_base_is_refused() {
assert!(matches!(
complete_url(
"https://gateway.ai.cloudflare.com/v1/acct/gw/azure-openai/res",
"gpt-4o",
"2024-10-21"
),
Err(Error::Unsupported(_))
));
}
#[rstest]
#[case::deployment(Some("2024-06-01"), env, "2024-06-01")]
#[case::environment(None, env, "2024-10-21")]
#[case::empty_falls_back(Some(""), env, "2024-10-21")]
#[case::default(None, no_env, AZURE_DEFAULT_API_VERSION)]
fn the_api_version_prefers_the_deployment_then_the_environment(
#[case] configured: Option<&str>,
#[case] lookup: fn(&str) -> Option<String>,
#[case] expected: &str,
) {
assert_eq!(resolve_api_version(configured, &lookup), expected);
}
#[rstest]
#[case::deployment(
Some("https://dep.openai.azure.com"),
env,
"https://dep.openai.azure.com"
)]
#[case::environment(None, env, "https://env.openai.azure.com")]
#[case::sdk_environment(None, endpoint_only, "https://sdk.openai.azure.com")]
fn the_api_base_prefers_the_deployment_then_the_environment(
#[case] configured: Option<&str>,
#[case] lookup: fn(&str) -> Option<String>,
#[case] expected: &str,
) {
assert_eq!(resolve_api_base(configured, &lookup).unwrap(), expected);
}
#[rstest]
fn a_missing_api_base_says_where_to_set_it() {
assert!(matches!(
resolve_api_base(None, &no_env),
Err(Error::Auth(litellm_auth::Error::MissingApiBase {
provider: "Azure",
..
}))
));
}
#[rstest]
#[case::deployment(Some("dep-key"), env, Some("dep-key"))]
#[case::environment(None, env, Some("env-key"))]
#[case::none(None, no_env, None)]
fn the_key_prefers_the_deployment_then_the_environment(
#[case] configured: Option<&str>,
#[case] lookup: fn(&str) -> Option<String>,
#[case] expected: Option<&str>,
) {
assert_eq!(resolve_api_key(configured, &lookup).as_deref(), expected);
}
#[rstest]
fn the_sdk_key_variable_outranks_litellms() {
let lookup = |name: &str| match name {
AZURE_OPENAI_API_KEY_ENV => Some("sdk-key".to_string()),
AZURE_API_KEY_ENV => Some("litellm-key".to_string()),
_ => None,
};
assert_eq!(resolve_api_key(None, &lookup).as_deref(), Some("sdk-key"));
}
#[rstest]
#[case::key(AzureCredential::ApiKey(SecretValue::new("k")), false, "api-key", "k")]
#[case::token(
AzureCredential::AdToken(SecretValue::new("t")),
false,
"authorization",
"Bearer t"
)]
#[case::v1_key(
AzureCredential::ApiKey(SecretValue::new("k")),
true,
"authorization",
"Bearer k"
)]
#[case::v1_token(
AzureCredential::AdToken(SecretValue::new("t")),
true,
"authorization",
"Bearer t"
)]
#[tokio::test]
async fn the_credential_lands_in_its_header(
#[case] credential: AzureCredential,
#[case] v1: bool,
#[case] header: &str,
#[case] value: &str,
) {
let authenticated = crate::base_llm::auth::resolve_auth(
&litellm_auth::AuthServices::default(),
environment(Vec::new(), credential, v1),
&no_env,
)
.await
.unwrap();
assert_eq!(
authenticated.headers,
[
("content-type".to_string(), "application/json".to_string()),
(header.to_string(), value.to_string())
]
);
}
}

5
litellm-rust/crates/llms/src/azure/mod.rs Normal file
View file

@ -0,0 +1,5 @@
//! Azure OpenAI (`azure/<deployment>`), as `litellm/llms/azure/` serves it. Azure AI Foundry's
//! `azure_ai` provider lives in `crate::azure_ai`.
pub mod chat;
pub mod common_utils;

5
litellm-rust/crates/llms/src/base_llm/chat/transformation.rs
View file

@ -71,6 +71,11 @@ pub trait BaseConfig: Sync {
None
}
/// The URL and body of the same call made as a stream.
fn streaming_request(&self, _url: String, _body: Value) -> Result<(String, Value), Error> {
Err(Error::Unsupported("streaming"))
}
/// Shapes the forwarded headers and names the credential, the way Python's
/// `validate_environment` does, without applying it: `resolve_auth` does that once
/// for every config.

1
litellm-rust/crates/llms/src/base_llm/mod.rs
View file

@ -3,6 +3,7 @@ pub mod auth;
pub mod base_model_iterator;
pub mod chat;
pub mod messages;
pub mod model_info;
pub mod ocr;
pub mod responses;
pub mod sse;

37
litellm-rust/crates/llms/src/base_llm/model_info.rs Normal file
View file

@ -0,0 +1,37 @@
//! What the model cost map says about a model, for configs whose request rules depend on it
//! the way Python's `supports_*` helpers read `litellm.model_cost`.
use serde_json::{Map, Value};
pub trait ModelInfo: Send + Sync {
/// The cost map entry `model` names, if there is one.
fn entry(&self, model: &str) -> Option<&Map<String, Value>>;
/// Whether any entry of the map sets `key`.
fn declares(&self, key: &str) -> bool;
}
/// Python's `_supports_factory`: the entry sets `key` to `true`. An unknown model supports
/// nothing.
pub fn supports(info: Option<&dyn ModelInfo>, model: &str, key: &str) -> bool {
field(info, model, key).and_then(Value::as_bool) == Some(true)
}
/// Python's `is_explicitly_disabled_factory`: the entry sets `key` to `false`. An unknown
/// model or an unset key is not disabled.
pub fn explicitly_disabled(info: Option<&dyn ModelInfo>, model: &str, key: &str) -> bool {
field(info, model, key).and_then(Value::as_bool) == Some(false)
}
/// The value the entry for `model` sets for `key`.
pub fn field<'a>(info: Option<&'a dyn ModelInfo>, model: &str, key: &str) -> Option<&'a Value> {
info?
.entry(model)?
.get(key)
.filter(|value| !value.is_null())
}
/// Whether the map has an entry for `model`.
pub fn known(info: Option<&dyn ModelInfo>, model: &str) -> bool {
info.is_some_and(|info| info.entry(model).is_some())
}

19
litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs
View file

@ -23,15 +23,19 @@ use crate::{
base_llm::{
auth::AuthScheme,
chat::{
streaming::StreamShape,
streaming::{ChatStream, StreamShape},
transformation::{
BaseConfig, Headers, ProviderChatRequestData, ProviderChatResponseData,
Unsupported, ValidatedEnvironment, unsupported_message, unsupported_param,
},
},
},
bedrock::chat::invoke_handler::converse_chat_stream,
};
/// The action a streaming Converse call goes to, in place of [`CONVERSE_PATH`].
const CONVERSE_STREAM_PATH: &str = "/converse-stream";
/// Converse parameter names, post `map_openai_params`, that the Rust path can
/// place verbatim in `inferenceConfig`.
///
@ -261,6 +265,19 @@ impl BaseConfig for AmazonConverseConfig {
&[("Content-Type", "application/json")]
}
fn model_response_iterator(&self, shape: StreamShape) -> Option<ChatStream> {
Some(converse_chat_stream(shape))
}
/// The same body, sent to the model's `converse-stream` action.
fn streaming_request(&self, url: String, body: Value) -> Result<(String, Value), Error> {
let url = match url.strip_suffix(CONVERSE_PATH) {
Some(base) => format!("{base}{CONVERSE_STREAM_PATH}"),
None => url,
};
Ok((url, body))
}
fn config_params(&self) -> &'static [&'static str] {
CONFIG_PARAMS
}

623
litellm-rust/crates/llms/src/bedrock/chat/invoke_handler.rs
View file

@ -1,18 +1,28 @@
use base64::Engine;
use bytes::Buf;
use futures_util::{Stream, StreamExt};
use futures_util::{Stream, StreamExt, stream::BoxStream};
use litellm_core_utils::core_helpers::finish_reason_for;
use litellm_framer::{
aws_event_stream::{AwsEventStreamCodec, Message},
aws_event_stream::{AwsEventStreamCodec, HeaderValue, Message},
frames,
};
use litellm_llms_types::formats::messages::streaming::MessagesStreamEvent;
use litellm_llms_types::formats::{
chat_completions::{
ChatCompletionChunk, ChatCompletionDelta, ChatCompletionStreamingChoice,
ChatCompletionThinkingBlock, ChatCompletionToolCallChunk,
ChatCompletionToolCallFunctionChunk, ChatCompletionsUsage, CompletionTokensDetails,
PromptTokensDetails,
},
messages::streaming::MessagesStreamEvent,
};
use serde::Deserialize;
use serde_json::Value;
use serde_json::{Map, Value};
use crate::{
Error,
anthropic::chat::handler::ModelResponseIterator,
base_llm::{
base_model_iterator::StreamTransformer,
chat::streaming::{ChatStream, StreamShape},
messages::streaming::{ByteStream, EventStream},
},
@ -92,6 +102,402 @@ pub fn invoke_chat_stream(invoke_provider: &str, shape: StreamShape) -> Result<C
}
}
/// The HTTP status Bedrock's ConverseStream models for each exception event; botocore
/// reports every one as a 400 on the wire.
const CONVERSE_STREAM_EXCEPTIONS: [(&str, u16); 5] = [
("internalServerException", 500),
("modelStreamErrorException", 424),
("validationException", 400),
("throttlingException", 429),
("serviceUnavailableException", 503),
];
const UNMODELED_EXCEPTION_STATUS: u16 = 400;
/// Python's status for a 200 stream whose body held no events.
const EMPTY_STREAM_STATUS: u16 = 502;
fn header<'a>(message: &'a Message, name: &str) -> Option<&'a str> {
message
.headers()
.iter()
.find(|header| header.name().as_str() == name)
.and_then(|header| match header.value() {
HeaderValue::String(value) => Some(value.as_str()),
_ => None,
})
}
/// One ConverseStream frame: its event payload, nothing for an empty one, or the
/// exception it carries as the provider's error.
pub fn decode_converse_message(message: &Message) -> Result<Option<Value>, Error> {
let body = String::from_utf8_lossy(message.payload()).into_owned();
match header(message, ":message-type") {
Some("exception") => {
let exception = header(message, ":exception-type").unwrap_or_default();
let status = CONVERSE_STREAM_EXCEPTIONS
.iter()
.find(|(name, _)| *name == exception)
.map_or(UNMODELED_EXCEPTION_STATUS, |(_, status)| *status);
Err(Error::Provider {
status,
body: format!("{exception} {body}").trim().to_owned(),
})
}
Some("error") => Err(Error::Provider {
status: UNMODELED_EXCEPTION_STATUS,
body: format!(
"{} {}",
header(message, ":error-code").unwrap_or_default(),
header(message, ":error-message").unwrap_or(&body)
)
.trim()
.to_owned(),
}),
_ if message.payload().is_empty() => Ok(None),
_ => serde_json::from_slice(message.payload())
.map(Some)
.map_err(|error| {
Error::InvalidResponse(crate::ErrorDetail::invalid("Bedrock event payload", error))
}),
}
}
/// ConverseStream's AWS event stream as the JSON payload of each event.
pub fn converse_event_stream(bytes: ByteStream) -> BoxStream<'static, Result<Value, Error>> {
let events = frames(bytes, AwsEventStreamCodec).map(|message| {
let message = message.map_err(|error| {
Error::InvalidResponse(crate::ErrorDetail::failed("stream framing", error))
})?;
decode_converse_message(&message)
});
futures_util::stream::unfold(
(Box::pin(events), false, false),
|(mut events, seen, done)| async move {
if done {
return None;
}
loop {
match events.next().await {
Some(Ok(None)) => continue,
Some(Ok(Some(event))) => return Some((Ok(event), (events, true, false))),
Some(Err(error)) => return Some((Err(error), (events, seen, true))),
None if seen => return None,
None => {
return Some((
Err(Error::Provider {
status: EMPTY_STREAM_STATUS,
body: "Bedrock answered the stream with HTTP 200 but its body decoded to no events".into(),
}),
(events, seen, true),
));
}
}
}
},
)
.boxed()
}
/// ConverseStream events to Chat Completions chunks, as Python's `AWSEventStreamDecoder`
/// `converse_chunk_parser` builds them.
///
/// DIVERGENCE: `json_mode` tool conversion is not done (the route never asks for it),
/// `redactedContent` is passed through without unwrapping, and the reasoning token estimate
/// Python makes from streamed thinking text is not made.
pub struct ConverseStreamIterator {
shape: StreamShape,
/// The tool argument fragments of the block in progress.
content_blocks: Vec<String>,
tool_calls_index: Option<i64>,
thinking_ran: bool,
thinking_block_index: i64,
thinking_content_block_index: Option<u64>,
}
#[derive(Default)]
struct ConverseParsed {
text: String,
tool_use: Option<ChatCompletionToolCallChunk>,
provider_specific_fields: Map<String, Value>,
reasoning_content: Option<String>,
thinking_blocks: Vec<ChatCompletionThinkingBlock>,
finish_reason: Option<String>,
usage: Option<ChatCompletionsUsage>,
}
impl ConverseStreamIterator {
pub fn new(shape: StreamShape) -> Self {
Self {
shape,
content_blocks: Vec::new(),
tool_calls_index: None,
thinking_ran: false,
thinking_block_index: -1,
thinking_content_block_index: None,
}
}
fn tool_call(
&self,
id: Option<String>,
name: Option<String>,
arguments: String,
index: u64,
) -> ChatCompletionToolCallChunk {
ChatCompletionToolCallChunk {
id,
tool_type: "function".into(),
function: ChatCompletionToolCallFunctionChunk {
name,
arguments,
provider_specific_fields: None,
},
index: self
.tool_calls_index
.unwrap_or_else(|| i64::try_from(index).unwrap_or(i64::MAX)),
}
}
fn thinking_blocks(
&mut self,
block: &Value,
content_block_index: u64,
) -> Vec<ChatCompletionThinkingBlock> {
if self.thinking_content_block_index != Some(content_block_index) {
self.thinking_content_block_index = Some(content_block_index);
self.thinking_block_index += 1;
}
let index = u64::try_from(self.thinking_block_index).ok();
let text = |field: &str| block.get(field).and_then(Value::as_str).map(str::to_owned);
let translated = if let Some(thinking) = text("text") {
Some(ChatCompletionThinkingBlock::Thinking {
thinking: Some(thinking),
signature: None,
cache_control: None,
index,
})
} else if let Some(signature) = text("signature") {
Some(ChatCompletionThinkingBlock::Thinking {
thinking: Some(String::new()),
signature: Some(signature),
cache_control: None,
index,
})
} else {
text("redactedContent").map(|data| ChatCompletionThinkingBlock::RedactedThinking {
data: Some(data),
cache_control: None,
index,
})
};
translated.into_iter().collect()
}
fn start(&mut self, start: &Value, index: u64) -> ConverseParsed {
let mut parsed = ConverseParsed::default();
self.content_blocks.clear();
if let Some(tool) = start.get("toolUse").filter(|tool| !tool.is_null()) {
self.tool_calls_index = Some(self.tool_calls_index.map_or(0, |index| index + 1));
let name = tool.get("name").and_then(Value::as_str).map(|name| {
self.shape
.tool_name_reverse_map
.get(name)
.cloned()
.unwrap_or_else(|| name.to_owned())
});
let id = tool
.get("toolUseId")
.and_then(Value::as_str)
.map(str::to_owned);
parsed.tool_use = Some(self.tool_call(id, name, String::new(), index));
} else if let Some(reasoning) = start.get("reasoningContent").filter(|r| !r.is_null()) {
parsed.thinking_blocks = self.thinking_blocks(reasoning, index);
parsed
.provider_specific_fields
.insert("reasoningContent".into(), reasoning.clone());
}
parsed
}
fn delta(&mut self, delta: &Value, index: u64) -> ConverseParsed {
let mut parsed = ConverseParsed::default();
if let Some(text) = delta.get("text").and_then(Value::as_str) {
parsed.text = text.to_owned();
} else if let Some(tool) = delta.get("toolUse") {
let input = tool
.get("input")
.and_then(Value::as_str)
.unwrap_or_default()
.to_owned();
self.content_blocks.push(input.clone());
parsed.tool_use = Some(self.tool_call(None, None, input, index));
} else if let Some(reasoning) = delta.get("reasoningContent") {
parsed
.provider_specific_fields
.insert("reasoningContent".into(), reasoning.clone());
parsed.reasoning_content = reasoning
.get("text")
.and_then(Value::as_str)
.map(str::to_owned);
parsed.thinking_blocks = self.thinking_blocks(reasoning, index);
if !parsed.thinking_blocks.is_empty() {
parsed.reasoning_content.get_or_insert_with(String::new);
}
} else if let Some(citations) = delta.get("citationsContent") {
parsed
.provider_specific_fields
.insert("citationsContent".into(), citations.clone());
}
parsed
}
pub fn converse_chunk_parser(&mut self, data: &Value) -> ChatCompletionChunk {
let index = data
.get("contentBlockIndex")
.and_then(Value::as_u64)
.unwrap_or(0);
let mut parsed = if let Some(start) = data.get("start") {
self.start(start, index)
} else if let Some(delta) = data.get("delta") {
self.delta(delta, index)
} else if data.get("contentBlockIndex").is_some() {
// A block's stop: a tool that streamed only empty fragments closes with `{}`.
let empty =
!self.content_blocks.is_empty() && self.content_blocks.iter().all(String::is_empty);
ConverseParsed {
tool_use: empty.then(|| self.tool_call(None, None, "{}".into(), index)),
..ConverseParsed::default()
}
} else if let Some(reason) = data.get("stopReason") {
ConverseParsed {
finish_reason: Some(
finish_reason_for(reason.as_str().unwrap_or("stop")).to_owned(),
),
..ConverseParsed::default()
}
} else if let Some(usage) = data.get("usage") {
ConverseParsed {
usage: Some(converse_usage(usage, self.thinking_ran)),
..ConverseParsed::default()
}
} else {
ConverseParsed::default()
};
if !parsed.thinking_blocks.is_empty() {
self.thinking_ran = true;
}
let trace = data.get("trace").filter(|trace| match trace {
Value::Null => false,
Value::Object(fields) => !fields.is_empty(),
Value::Array(items) => !items.is_empty(),
Value::String(text) => !text.is_empty(),
_ => true,
});
let carries_content = trace.is_some()
|| ["start", "delta", "contentBlockIndex", "stopReason"]
.iter()
.any(|key| data.get(*key).is_some());
let thinking_blocks = std::mem::take(&mut parsed.thinking_blocks);
let provider_specific_fields = std::mem::take(&mut parsed.provider_specific_fields);
ChatCompletionChunk {
choices: vec![ChatCompletionStreamingChoice {
index: 0,
delta: ChatCompletionDelta {
content: carries_content.then_some(parsed.text),
role: carries_content.then(|| "assistant".to_owned()),
tool_calls: parsed.tool_use.map(|tool| vec![tool]),
reasoning_content: parsed.reasoning_content,
thinking_blocks: (!thinking_blocks.is_empty()).then_some(thinking_blocks),
provider_specific_fields: (!provider_specific_fields.is_empty())
.then_some(provider_specific_fields),
..ChatCompletionDelta::default()
},
finish_reason: parsed.finish_reason,
logprobs: None,
}],
usage: parsed.usage,
provider_specific_fields: Some(
trace
.map(|trace| Map::from_iter([("trace".to_owned(), trace.clone())]))
.unwrap_or_default(),
),
..ChatCompletionChunk::default()
}
}
}
/// Python's `AmazonConverseConfig.transform_usage` for a stream's `metadata` usage.
pub fn converse_usage(usage: &Value, thinking_ran: bool) -> ChatCompletionsUsage {
let count = |names: &[&str]| {
names
.iter()
.find_map(|name| usage.get(*name).and_then(Value::as_u64))
.unwrap_or(0)
};
let input = count(&["inputTokens"]);
let output = count(&["outputTokens"]);
let cache_read = count(&["cacheReadInputTokens", "cacheReadInputTokenCount"]);
let cache_write = count(&["cacheWriteInputTokens", "cacheWriteInputTokenCount"]);
let prompt = input + cache_read + cache_write;
ChatCompletionsUsage {
prompt_tokens: prompt,
completion_tokens: output,
total_tokens: usage
.get("totalTokens")
.and_then(Value::as_u64)
.unwrap_or(prompt + output),
prompt_tokens_details: PromptTokensDetails {
cached_tokens: cache_read,
cache_creation_tokens: cache_write,
text_tokens: input,
cache_write_tokens: Some(cache_write),
cache_creation_token_details: None,
},
completion_tokens_details: Some(match thinking_ran {
true => CompletionTokensDetails::default(),
false => CompletionTokensDetails {
reasoning_tokens: Some(0),
text_tokens: Some(output),
},
}),
server_tool_use: None,
cache_creation_input_tokens: Some(cache_write),
cache_read_input_tokens: Some(cache_read),
}
}
impl StreamTransformer for ConverseStreamIterator {
type Input = Value;
type Output = ChatCompletionChunk;
type Error = Error;
/// Python hands only Converse events (those with a block index, a stop reason, metrics
/// or a trace) to the chunk parser; `messageStart` becomes an empty generic chunk that
/// never reaches the client.
fn transform(&mut self, input: Value) -> Result<Vec<ChatCompletionChunk>, Error> {
let converse = [
"contentBlockIndex",
"stopReason",
"metrics",
"trace",
"usage",
]
.iter()
.any(|key| input.get(*key).is_some());
Ok(match converse {
true => vec![self.converse_chunk_parser(&input)],
false => Vec::new(),
})
}
fn finish(&mut self) -> Result<Vec<ChatCompletionChunk>, Error> {
Ok(Vec::new())
}
}
pub fn converse_chat_stream(shape: StreamShape) -> ChatStream {
ChatStream::new(converse_event_stream, ConverseStreamIterator::new(shape))
}
#[cfg(test)]
mod tests {
use aws_smithy_eventstream::frame::write_message_to;
@ -163,4 +569,213 @@ mod tests {
);
assert_eq!(from_aws, from_sse);
}
fn converse_frame(kind: &str, name: &str, payload: serde_json::Value) -> Vec<u8> {
let type_header = match kind {
"exception" => ":exception-type",
_ => ":event-type",
};
let message = Message::new(Bytes::from(serde_json::to_vec(&payload).unwrap()))
.add_header(Header::new(
":message-type",
HeaderValue::String(kind.to_owned().into()),
))
.add_header(Header::new(
type_header,
HeaderValue::String(name.to_owned().into()),
));
let mut wire = Vec::new();
write_message_to(&message, &mut wire).unwrap();
wire
}
async fn converse_events(wire: Vec<u8>) -> Vec<Result<serde_json::Value, Error>> {
converse_event_stream(in_pieces(&wire)).collect().await
}
#[tokio::test]
async fn converse_frames_decode_to_their_payloads_split_anywhere() {
let wire = [
converse_frame(
"event",
"messageStart",
serde_json::json!({"role": "assistant"}),
),
converse_frame(
"event",
"contentBlockDelta",
serde_json::json!({"contentBlockIndex": 0, "delta": {"text": "hi"}}),
),
]
.concat();
let events = converse_events(wire).await;
assert_eq!(
events.into_iter().map(Result::unwrap).collect::<Vec<_>>(),
vec![
serde_json::json!({"role": "assistant"}),
serde_json::json!({"contentBlockIndex": 0, "delta": {"text": "hi"}})
]
);
}
#[rstest::rstest]
#[case::throttling("throttlingException", 429)]
#[case::validation("validationException", 400)]
#[case::internal("internalServerException", 500)]
#[case::unavailable("serviceUnavailableException", 503)]
#[case::model_stream("modelStreamErrorException", 424)]
#[case::unmodeled("somethingNewException", 400)]
#[tokio::test]
async fn an_exception_frame_is_the_providers_error_with_its_modeled_status(
#[case] exception: &str,
#[case] expected: u16,
) {
let wire = [
converse_frame(
"event",
"contentBlockDelta",
serde_json::json!({"contentBlockIndex": 0, "delta": {"text": "hi"}}),
),
converse_frame(
"exception",
exception,
serde_json::json!({"message": "slow down"}),
),
]
.concat();
let mut events = converse_events(wire).await.into_iter();
assert!(events.next().unwrap().is_ok());
let Some(Err(Error::Provider { status, body })) = events.next() else {
panic!("the exception ends the stream as the provider's error");
};
assert_eq!(status, expected);
assert_eq!(body, format!("{exception} {{\"message\":\"slow down\"}}"));
assert!(events.next().is_none());
}
#[tokio::test]
async fn a_stream_without_events_is_a_bad_gateway() {
let events = converse_events(Vec::new()).await;
assert!(matches!(
events.as_slice(),
[Err(Error::Provider { status: 502, .. })]
));
}
fn converse_chunks(events: serde_json::Value) -> Vec<serde_json::Value> {
let mut iterator = ConverseStreamIterator::new(StreamShape::default());
let serde_json::Value::Array(events) = events else {
panic!("events are a list");
};
events
.into_iter()
.flat_map(|event| iterator.transform(event).unwrap())
.map(|chunk| serde_json::to_value(chunk).unwrap())
.collect()
}
#[test]
fn converse_text_and_stop_follow_python() {
let chunks = converse_chunks(serde_json::json!([
{"role": "assistant"},
{"contentBlockIndex": 0, "delta": {"text": "Hello"}},
{"contentBlockIndex": 0},
{"stopReason": "end_turn"},
{"usage": {"inputTokens": 10, "outputTokens": 12, "totalTokens": 48, "cacheReadInputTokens": 20, "cacheWriteInputTokens": 6}, "metrics": {"latencyMs": 1}}
]));
assert_eq!(chunks.len(), 4, "messageStart is not a Converse chunk");
assert_eq!(
chunks[0],
serde_json::json!({"choices": [{"index": 0, "delta": {"content": "Hello", "role": "assistant"}}], "provider_specific_fields": {}})
);
assert_eq!(
chunks[1]["choices"][0]["delta"],
serde_json::json!({"content": "", "role": "assistant"})
);
assert_eq!(chunks[2]["choices"][0]["finish_reason"], "stop");
assert_eq!(chunks[3]["choices"][0]["delta"], serde_json::json!({}));
assert_eq!(
chunks[3]["usage"],
serde_json::json!({
"completion_tokens": 12, "prompt_tokens": 36, "total_tokens": 48,
"completion_tokens_details": {"reasoning_tokens": 0, "text_tokens": 12},
"prompt_tokens_details": {"cached_tokens": 20, "text_tokens": 10, "cache_write_tokens": 6, "cache_creation_tokens": 6},
"cache_creation_input_tokens": 6, "cache_read_input_tokens": 20
})
);
}
#[test]
fn converse_tool_calls_stream_name_then_partial_arguments() {
let chunks = converse_chunks(serde_json::json!([
{"contentBlockIndex": 1, "start": {"toolUse": {"toolUseId": "t1", "name": "get_weather"}}},
{"contentBlockIndex": 1, "delta": {"toolUse": {"input": "{\"city\": \"Pa"}}},
{"contentBlockIndex": 1, "delta": {"toolUse": {"input": "ris\"}"}}},
{"contentBlockIndex": 1},
{"contentBlockIndex": 2, "start": {"toolUse": {"toolUseId": "t2", "name": "now"}}},
{"contentBlockIndex": 2, "delta": {"toolUse": {"input": ""}}},
{"contentBlockIndex": 2},
{"stopReason": "tool_use"}
]));
let tools = |chunk: &serde_json::Value| chunk["choices"][0]["delta"]["tool_calls"].clone();
assert_eq!(
tools(&chunks[0]),
serde_json::json!([{"id": "t1", "type": "function", "function": {"name": "get_weather", "arguments": ""}, "index": 0}])
);
assert_eq!(
tools(&chunks[1]),
serde_json::json!([{"type": "function", "function": {"arguments": "{\"city\": \"Pa"}, "index": 0}])
);
assert_eq!(tools(&chunks[3]), serde_json::Value::Null);
assert_eq!(tools(&chunks[4])[0]["index"], 1);
assert_eq!(
tools(&chunks[6]),
serde_json::json!([{"type": "function", "function": {"arguments": "{}"}, "index": 1}])
);
assert_eq!(chunks[7]["choices"][0]["finish_reason"], "tool_calls");
}
#[test]
fn converse_reasoning_streams_as_thinking_blocks() {
let chunks = converse_chunks(serde_json::json!([
{"contentBlockIndex": 0, "delta": {"reasoningContent": {"text": "Let me think"}}},
{"contentBlockIndex": 0, "delta": {"reasoningContent": {"signature": "sig"}}},
{"contentBlockIndex": 1, "start": {"reasoningContent": {"redactedContent": "enc"}}},
{"stopReason": "max_tokens"},
{"usage": {"inputTokens": 1, "outputTokens": 5}, "metrics": {}}
]));
assert_eq!(
chunks[0]["choices"][0]["delta"],
serde_json::json!({"content": "", "role": "assistant", "reasoning_content": "Let me think",
"thinking_blocks": [{"type": "thinking", "thinking": "Let me think", "index": 0}],
"provider_specific_fields": {"reasoningContent": {"text": "Let me think"}}})
);
assert_eq!(chunks[1]["choices"][0]["delta"]["reasoning_content"], "");
assert_eq!(
chunks[1]["choices"][0]["delta"]["thinking_blocks"],
serde_json::json!([{"type": "thinking", "thinking": "", "signature": "sig", "index": 0}])
);
assert_eq!(
chunks[2]["choices"][0]["delta"]["thinking_blocks"],
serde_json::json!([{"type": "redacted_thinking", "data": "enc", "index": 1}])
);
assert_eq!(chunks[3]["choices"][0]["finish_reason"], "length");
assert_eq!(
chunks[4]["usage"]["completion_tokens_details"],
serde_json::json!({})
);
assert_eq!(chunks[4]["usage"]["total_tokens"], 6);
}
#[test]
fn a_guardrail_trace_rides_on_the_chunk() {
let chunks = converse_chunks(serde_json::json!([
{"stopReason": "guardrail_intervened", "trace": {"guardrail": {"action": "BLOCKED"}}}
]));
assert_eq!(chunks[0]["choices"][0]["finish_reason"], "content_filter");
assert_eq!(
chunks[0]["provider_specific_fields"],
serde_json::json!({"trace": {"guardrail": {"action": "BLOCKED"}}})
);
}
}

12
litellm-rust/crates/llms/src/error.rs
View file

@ -15,6 +15,10 @@ pub enum Error {
Unsupported(&'static str),
#[error(transparent)]
Auth(#[from] litellm_auth::Error),
/// An error the provider reported inside a response it had already started, such as an
/// `error` event on a stream, with the HTTP status the provider gives that error.
#[error("provider error {status}: {body}")]
Provider { status: u16, body: String },
}
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
@ -60,6 +64,14 @@ pub enum ErrorDetail {
value: String,
hint: String,
},
#[error(
"{provider} does not support parameters: {params:?}, for model={model}. To drop these, set `litellm.drop_params=True` or for proxy:\n\n`litellm_settings:\n drop_params: true`\n. \n If you want to use these params dynamically send allowed_openai_params={params:?} in your request."
)]
UnsupportedParams {
provider: &'static str,
params: Vec<String>,
model: String,
},
#[error(
"{provider} does not support the server-side tool '{tool}' (model '{model}'). {guidance}"
)]

1
litellm-rust/crates/llms/src/lib.rs
View file

@ -1,5 +1,6 @@
pub mod anthropic;
pub mod aws_textract;
pub mod azure;
pub mod azure_ai;
pub mod base_llm;
pub mod bedrock;

28
litellm-rust/crates/llms/src/openai/chat/transformation.rs
View file

@ -85,17 +85,7 @@ impl OpenAiChatConfig {
body: Map<String, Value>,
) -> Result<Map<String, Value>, Error> {
let mut body = body;
let extra_body = match body.remove(EXTRA_BODY) {
None | Some(Value::Null) => Map::new(),
Some(Value::Object(extra)) => extra,
Some(other) => {
return Err(Error::InvalidRequest(ErrorDetail::InvalidValue {
field: EXTRA_BODY,
expected: "an object",
actual: other,
}));
}
};
let extra_body = take_extra_body(&mut body)?;
let mut body: Map<String, Value> = body
.into_iter()
.chain(extra_body)
@ -120,7 +110,21 @@ pub fn is_stream(body: &Map<String, Value>) -> Result<bool, Error> {
}
}
fn is_litellm_owned(name: &str) -> bool {
/// Removes the body's `extra_body`, which the OpenAI SDK merges over the body it sends.
pub fn take_extra_body(body: &mut Map<String, Value>) -> Result<Map<String, Value>, Error> {
match body.remove(EXTRA_BODY) {
None | Some(Value::Null) => Ok(Map::new()),
Some(Value::Object(extra)) => Ok(extra),
Some(other) => Err(Error::InvalidRequest(ErrorDetail::InvalidValue {
field: EXTRA_BODY,
expected: "an object",
actual: other,
})),
}
}
/// Whether Python consumes `name` before the provider call, so it never reaches the wire.
pub fn is_litellm_owned(name: &str) -> bool {
LITELLM_OWNED_PARAMS.contains(&name) || name.starts_with(INTERNAL_PARAM_PREFIX)
}

2
litellm-rust/crates/llms/src/openai_like/chat/transformation.rs
View file

@ -173,7 +173,9 @@ impl BaseConfig for OpenAILikeChatConfig {
.and_then(|d| d.get("text_tokens"))
.and_then(Value::as_u64)
.unwrap_or(0),
..PromptTokensDetails::default()
},
..ChatCompletionsUsage::default()
},
})
}

28
litellm-rust/crates/router-types/src/value.rs
View file

@ -8,3 +8,31 @@ pub enum Spelled<T> {
Value(T),
Text(String),
}
impl Spelled<bool> {
/// The flag, reading a `true` word the way Python's coercion does. An `os.environ/`
/// reference left unresolved is not set.
pub fn is_set(&self) -> bool {
match self {
Self::Value(value) => *value,
Self::Text(text) => text.trim().eq_ignore_ascii_case("true"),
}
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
#[rstest]
#[case::value(Spelled::Value(true), true)]
#[case::false_value(Spelled::Value(false), false)]
#[case::word(Spelled::Text(" True ".into()), true)]
#[case::other_word(Spelled::Text("yes".into()), false)]
#[case::reference(Spelled::Text("os.environ/DROP".into()), false)]
fn flags_read_values_and_words(#[case] flag: Spelled<bool>, #[case] expected: bool) {
assert_eq!(flag.is_set(), expected);
}
}

9
litellm-rust/crates/router/src/deployment.rs
View file

@ -21,6 +21,9 @@ pub struct Deployment {
/// `model_info.id`, which names the deployment in state shared with other replicas and
/// with Python routers on the same Redis.
pub id: Option<String>,
/// The model a deployment serves under another name, from `litellm_params.base_model` or
/// `model_info.base_model`, which Azure's param rules and prices follow.
pub base_model: Option<String>,
}
impl Deployment {
@ -46,6 +49,12 @@ impl Deployment {
shaping: MessagesShaping::default(),
team_id: None,
id: None,
base_model: params
.extra
.get("base_model")
.and_then(|value| value.as_str())
.filter(|name| !name.is_empty())
.map(str::to_owned),
})
}

36
litellm-rust/crates/router/src/lib.rs
View file

@ -87,6 +87,10 @@ impl Router {
timeout: deployment.timeout.or(policy.timeout),
team_id: owner_team(model),
id: model_id(model),
base_model: deployment
.base_model
.clone()
.or_else(|| info_base_model(model)),
..deployment
},
))
@ -363,6 +367,13 @@ fn owner_team(model: &Model) -> Option<String> {
}
}
fn info_base_model(model: &Model) -> Option<String> {
match model.model_info.get("base_model")? {
litellm_config::Value::String(name) if !name.is_empty() => Some(name.clone()),
_ => None,
}
}
impl FromIterator<(String, Deployment)> for Router {
fn from_iter<I: IntoIterator<Item = (String, Deployment)>>(entries: I) -> Self {
Self::new(entries, RoutingPolicy::default())
@ -559,6 +570,31 @@ mod tests {
);
}
#[rstest]
#[case::params(
" litellm_params:\n model: azure/d\n base_model: gpt-4o\n",
Some("gpt-4o")
)]
#[case::model_info(
" litellm_params:\n model: azure/d\n model_info:\n base_model: gpt-4o\n",
Some("gpt-4o")
)]
#[case::params_win(
" litellm_params:\n model: azure/d\n base_model: gpt-4o\n model_info:\n base_model: gpt-4\n",
Some("gpt-4o")
)]
#[case::absent(" litellm_params:\n model: azure/d\n", None)]
fn reads_the_base_model(#[case] entry: &str, #[case] expected: Option<&str>) {
let config =
Config::from_yaml(&format!("model_list:\n - model_name: m\n{entry}")).unwrap();
let router = Router::from_config(&config).unwrap();
let group = router.group("m").unwrap();
assert_eq!(
group.deployments().next().unwrap().base_model.as_deref(),
expected
);
}
#[rstest]
#[tokio::test]
async fn an_unclassified_error_is_not_retried() {

1
litellm-rust/crates/router/tests/router.rs
View file

@ -117,6 +117,7 @@ fn programmatic_deployments_preserve_overrides_and_share_their_group() {
},
team_id: None,
id: None,
base_model: None,
};
let router = Router::from_iter([
("public-model".into(), Deployment::default()),

87
litellm-rust/docs/customer-proxy/azure-chat.svg Normal file
View file

@ -0,0 +1,87 @@
<svg xmlns="http://www.w3.org/2000/svg" width="980" height="760" viewBox="0 0 980 760" font-family="Helvetica, Arial, sans-serif" font-size="13">
<defs>
<marker id="a" viewBox="0 0 10 10" refX="9" refY="5" markerWidth="7" markerHeight="7" orient="auto-start-reverse">
<path d="M0,0 L10,5 L0,10 z" fill="#333"/>
</marker>
</defs>
<rect x="0" y="0" width="980" height="760" fill="#ffffff"/>
<text x="20" y="30" font-size="17" font-weight="bold" fill="#111">POST /v1/chat/completions to an azure/ deployment</text>
<!-- client -->
<rect x="20" y="60" width="150" height="56" rx="6" fill="#eef3fb" stroke="#4a6fa5"/>
<text x="95" y="84" text-anchor="middle" fill="#111" font-weight="bold">Client</text>
<text x="95" y="102" text-anchor="middle" fill="#333">OpenAI SDK body</text>
<!-- gateway handle -->
<rect x="220" y="50" width="300" height="76" rx="6" fill="#f5f5f5" stroke="#666"/>
<text x="370" y="72" text-anchor="middle" fill="#111" font-weight="bold">gateway-inference chat_completions</text>
<text x="370" y="92" text-anchor="middle" fill="#333">auth, hooks admit, route deployment,</text>
<text x="370" y="110" text-anchor="middle" fill="#333">litellm_params + base_model + drop_params</text>
<line x1="170" y1="88" x2="218" y2="88" stroke="#333" marker-end="url(#a)"/>
<!-- decision -->
<polygon points="370,150 490,190 370,230 250,190" fill="#fff8e6" stroke="#b8860b"/>
<text x="370" y="186" text-anchor="middle" fill="#111">serves_openai?</text>
<text x="370" y="202" text-anchor="middle" fill="#111">openai | azure</text>
<line x1="370" y1="126" x2="370" y2="148" stroke="#333" marker-end="url(#a)"/>
<rect x="580" y="50" width="380" height="76" rx="6" fill="#fff8e6" stroke="#b8860b"/>
<text x="770" y="72" text-anchor="middle" fill="#111" font-weight="bold">gateway startup (hooks::install)</text>
<text x="770" y="92" text-anchor="middle" fill="#333">an azure deployment loads the cost map and</text>
<text x="770" y="110" text-anchor="middle" fill="#333">installs it as ModelInfo for the param rules</text>
<!-- run_openai / prepare_azure -->
<rect x="20" y="262" width="560" height="232" rx="6" fill="#eef7ee" stroke="#3c8d3c"/>
<text x="300" y="284" text-anchor="middle" fill="#111" font-weight="bold">inference-chat run_openai, prepare_azure</text>
<text x="36" y="308" fill="#333">1. resolve os.environ/ refs and AZURE_* env through the secret source</text>
<text x="36" y="328" fill="#333">2. api_version: deployment, AZURE_API_VERSION, else 2025-02-01-preview</text>
<text x="54" y="346" fill="#333">api_base: deployment, AZURE_API_BASE, AZURE_OPENAI_ENDPOINT</text>
<text x="36" y="366" fill="#333">3. AzureChatRequest::transform: family from base_model or deployment</text>
<text x="54" y="384" fill="#333">(gpt, o-series, gpt-5), unsupported params refused or dropped</text>
<text x="36" y="404" fill="#333">4. url: {base}/openai/deployments/{dep}/chat/completions?api-version=</text>
<text x="54" y="422" fill="#333">v1, latest, preview: {base}/openai/v1/chat/completions</text>
<text x="36" y="442" fill="#333">5. credential: azure_ad_token, else api-key, else Entra ID token</text>
<text x="54" y="460" fill="#333">(client secret, OIDC, AZURE_AD_TOKEN, managed identity)</text>
<text x="36" y="482" fill="#555">then the shared interceptors, send and SSE relay of the OpenAI path</text>
<line x1="370" y1="230" x2="370" y2="260" stroke="#333" marker-end="url(#a)"/>
<text x="382" y="250" fill="#555">yes</text>
<!-- upstream -->
<rect x="680" y="300" width="280" height="110" rx="6" fill="#eef3fb" stroke="#4a6fa5"/>
<text x="820" y="326" text-anchor="middle" fill="#111" font-weight="bold">Azure OpenAI</text>
<text x="820" y="346" text-anchor="middle" fill="#333">api-key: key (deployment path)</text>
<text x="820" y="364" text-anchor="middle" fill="#333">Authorization: Bearer (AD token, v1)</text>
<text x="820" y="390" text-anchor="middle" fill="#333">application/json or text/event-stream</text>
<line x1="580" y1="340" x2="678" y2="340" stroke="#333" marker-end="url(#a)"/>
<line x1="678" y1="380" x2="580" y2="380" stroke="#333" marker-end="url(#a)"/>
<!-- refusals -->
<rect x="580" y="160" width="380" height="76" rx="6" fill="#f5f5f5" stroke="#999" stroke-dasharray="5,4"/>
<text x="770" y="182" text-anchor="middle" fill="#111">Refused before the call</text>
<text x="770" y="200" text-anchor="middle" fill="#555">unsupported param without drop_params: 400</text>
<text x="770" y="218" text-anchor="middle" fill="#555">Responses API bridge cases, Cloudflare bases: 501</text>
<!-- bottom boxes -->
<rect x="20" y="530" width="440" height="210" rx="6" fill="#ffffff" stroke="#b8860b"/>
<text x="240" y="552" text-anchor="middle" fill="#111" font-weight="bold">usage and cost (gateway-inference Finish)</text>
<text x="36" y="576" fill="#333">usage from the body or the include_usage chunk</text>
<text x="36" y="596" fill="#333">config prices on the deployment win</text>
<text x="36" y="616" fill="#333">base_model set: cost map azure/{base_model},</text>
<text x="54" y="634" fill="#333">then {base_model}</text>
<text x="36" y="654" fill="#333">otherwise: azure/{model the response names},</text>
<text x="54" y="672" fill="#333">then azure/{deployment}</text>
<text x="36" y="696" fill="#555">the bare deployment name is never priced, it names</text>
<text x="36" y="714" fill="#555">a resource, not a model</text>
<rect x="500" y="530" width="460" height="210" rx="6" fill="#ffffff" stroke="#3c8d3c"/>
<text x="730" y="552" text-anchor="middle" fill="#111" font-weight="bold">errors</text>
<text x="516" y="576" fill="#333">upstream 4xx/5xx: same status, OpenAI error envelope</text>
<text x="516" y="596" fill="#333">carrying Azure's error.message</text>
<text x="516" y="624" fill="#333">content filter (inner_error with a policy violation):</text>
<text x="516" y="644" fill="#333">provider_specific_fields keeps innererror,</text>
<text x="516" y="664" fill="#333">inner_error and the whole azure_error</text>
<text x="516" y="692" fill="#333">stream cut after headers: a final</text>
<text x="516" y="712" fill="#333">data: {"error": ...} event, then close</text>
<line x1="240" y1="494" x2="240" y2="528" stroke="#333" marker-end="url(#a)"/>
<line x1="500" y1="494" x2="730" y2="528" stroke="#333" marker-end="url(#a)"/>
</svg>
Side by side

After

Width:  |  Height:  |  Size: 6.3 KiB

104
litellm-rust/docs/customer-proxy/chat-streaming.svg Normal file
View file

@ -0,0 +1,104 @@
<svg xmlns="http://www.w3.org/2000/svg" width="980" height="800" viewBox="0 0 980 800" font-family="Helvetica, Arial, sans-serif" font-size="13">
<defs>
<marker id="a" viewBox="0 0 10 10" refX="9" refY="5" markerWidth="7" markerHeight="7" orient="auto-start-reverse">
<path d="M0,0 L10,5 L0,10 z" fill="#333"/>
</marker>
</defs>
<rect x="0" y="0" width="980" height="800" fill="#ffffff"/>
<text x="20" y="30" font-size="17" font-weight="bold" fill="#111">POST /v1/chat/completions "stream": true to an anthropic/ or bedrock/ deployment</text>
<!-- client request -->
<rect x="20" y="60" width="150" height="56" rx="6" fill="#eef3fb" stroke="#4a6fa5"/>
<text x="95" y="84" text-anchor="middle" fill="#111" font-weight="bold">Client</text>
<text x="95" y="102" text-anchor="middle" fill="#333">stream, stream_options</text>
<!-- gateway -->
<rect x="220" y="50" width="320" height="76" rx="6" fill="#f5f5f5" stroke="#666"/>
<text x="380" y="72" text-anchor="middle" fill="#111" font-weight="bold">gateway-inference chat_completions</text>
<text x="380" y="92" text-anchor="middle" fill="#333">is_stream: stream_machine, Sse encoder,</text>
<text x="380" y="110" text-anchor="middle" fill="#333">ReportedUsage response extension</text>
<line x1="170" y1="88" x2="218" y2="88" stroke="#333" marker-end="url(#a)"/>
<!-- run_stream -->
<rect x="20" y="160" width="560" height="196" rx="6" fill="#eef7ee" stroke="#3c8d3c"/>
<text x="300" y="182" text-anchor="middle" fill="#111" font-weight="bold">inference-chat run_stream (route_where retries)</text>
<text x="36" y="206" fill="#333">1. strip stream and stream_options, keep include_usage</text>
<text x="36" y="226" fill="#333">2. resolve, prepare, the same param rules as non-streaming</text>
<text x="36" y="246" fill="#333">3. streaming_request: anthropic adds "stream": true,</text>
<text x="54" y="264" fill="#333">bedrock /converse becomes /converse-stream</text>
<text x="36" y="284" fill="#333">4. auth and interceptors, send; timeout bounds the head,</text>
<text x="54" y="302" fill="#333">then an idle timeout per read, no total limit</text>
<text x="36" y="322" fill="#333">5. wait for the first OpenAI chunk before answering, so an</text>
<text x="54" y="340" fill="#333">error up to there falls over to another deployment</text>
<line x1="380" y1="126" x2="380" y2="158" stroke="#333" marker-end="url(#a)"/>
<!-- upstream -->
<rect x="660" y="160" width="300" height="110" rx="6" fill="#eef3fb" stroke="#4a6fa5"/>
<text x="810" y="186" text-anchor="middle" fill="#111" font-weight="bold">Provider</text>
<text x="810" y="208" text-anchor="middle" fill="#333">Anthropic: text/event-stream</text>
<text x="810" y="228" text-anchor="middle" fill="#333">Bedrock: AWS event stream</text>
<text x="810" y="248" text-anchor="middle" fill="#333">(application/vnd.amazon.eventstream)</text>
<line x1="580" y1="200" x2="658" y2="200" stroke="#333" marker-end="url(#a)"/>
<!-- pipeline row -->
<text x="20" y="398" font-weight="bold" fill="#111">Per chunk, pulled by the client body (dropping it drops the upstream request)</text>
<rect x="20" y="414" width="170" height="86" rx="6" fill="#eef3fb" stroke="#4a6fa5"/>
<text x="105" y="438" text-anchor="middle" fill="#111" font-weight="bold">provider stream</text>
<text x="105" y="458" text-anchor="middle" fill="#333">SSE events or</text>
<text x="105" y="476" text-anchor="middle" fill="#333">event-stream frames</text>
<rect x="220" y="414" width="190" height="86" rx="6" fill="#fff8e6" stroke="#b8860b"/>
<text x="315" y="438" text-anchor="middle" fill="#111" font-weight="bold">decoder</text>
<text x="315" y="458" text-anchor="middle" fill="#333">ModelResponseIterator</text>
<text x="315" y="476" text-anchor="middle" fill="#333">ConverseStreamIterator</text>
<rect x="440" y="414" width="220" height="86" rx="6" fill="#eef7ee" stroke="#3c8d3c"/>
<text x="550" y="438" text-anchor="middle" fill="#111" font-weight="bold">StreamWrapper</text>
<text x="550" y="458" text-anchor="middle" fill="#333">OpenAI chunks: role, one id,</text>
<text x="550" y="476" text-anchor="middle" fill="#333">finish chunk, usage chunk</text>
<rect x="690" y="414" width="130" height="86" rx="6" fill="#f5f5f5" stroke="#666"/>
<text x="755" y="438" text-anchor="middle" fill="#111" font-weight="bold">Sse encoder</text>
<text x="755" y="458" text-anchor="middle" fill="#333">data: frames,</text>
<text x="755" y="476" text-anchor="middle" fill="#333">data: [DONE]</text>
<rect x="850" y="414" width="110" height="86" rx="6" fill="#eef3fb" stroke="#4a6fa5"/>
<text x="905" y="452" text-anchor="middle" fill="#111" font-weight="bold">Client</text>
<text x="905" y="472" text-anchor="middle" fill="#333">OpenAI SSE</text>
<line x1="190" y1="457" x2="218" y2="457" stroke="#333" marker-end="url(#a)"/>
<line x1="410" y1="457" x2="438" y2="457" stroke="#333" marker-end="url(#a)"/>
<line x1="660" y1="457" x2="688" y2="457" stroke="#333" marker-end="url(#a)"/>
<line x1="820" y1="457" x2="848" y2="457" stroke="#333" marker-end="url(#a)"/>
<!-- meter -->
<rect x="440" y="530" width="380" height="76" rx="6" fill="#ffffff" stroke="#b8860b"/>
<text x="630" y="552" text-anchor="middle" fill="#111" font-weight="bold">meter and observers (gateway-inference Finish)</text>
<text x="630" y="572" text-anchor="middle" fill="#333">usage from ReportedUsage, set at stream end</text>
<text x="630" y="590" text-anchor="middle" fill="#333">even when the client did not ask for include_usage</text>
<line x1="550" y1="500" x2="550" y2="528" stroke="#333" marker-end="url(#a)"/>
<line x1="755" y1="500" x2="755" y2="528" stroke="#333" marker-end="url(#a)"/>
<!-- bottom boxes -->
<rect x="20" y="530" width="390" height="250" rx="6" fill="#ffffff" stroke="#3c8d3c"/>
<text x="215" y="552" text-anchor="middle" fill="#111" font-weight="bold">chunk shape (as the Python proxy)</text>
<text x="36" y="576" fill="#333">first chunk carries role: assistant</text>
<text x="36" y="596" fill="#333">usage stripped from every chunk; include_usage adds</text>
<text x="54" y="614" fill="#333">a final chunk with choices [{index 0, delta {}}]</text>
<text x="36" y="634" fill="#333">finish_reason: end_turn stop, max_tokens length,</text>
<text x="54" y="652" fill="#333">tool_use tool_calls; else stop or tool_calls at end</text>
<text x="36" y="672" fill="#333">tool call deltas with content "", empty args "{}"</text>
<text x="36" y="692" fill="#333">reasoning_content and thinking_blocks (with index,</text>
<text x="54" y="710" fill="#333">signature, redacted_thinking)</text>
<text x="36" y="730" fill="#333">model: the name the client asked for</text>
<text x="36" y="750" fill="#555">bedrock chunks carry provider_specific_fields {}</text>
<rect x="440" y="630" width="520" height="150" rx="6" fill="#ffffff" stroke="#999" stroke-dasharray="5,4"/>
<text x="700" y="652" text-anchor="middle" fill="#111" font-weight="bold">errors</text>
<text x="456" y="676" fill="#333">upstream status before the first chunk: retry or fallback,</text>
<text x="474" y="694" fill="#333">then an OpenAI JSON error with the same status</text>
<text x="456" y="718" fill="#333">anthropic error event, bedrock exception (throttling 429,</text>
<text x="474" y="736" fill="#333">validation 400, ...) after the first chunk: a final</text>
<text x="474" y="754" fill="#333">data: {"error": ...} frame, no [DONE], no usage recorded</text>
</svg>
Side by side

After

Width:  |  Height:  |  Size: 7.4 KiB

52
litellm-rust/docs/customer-proxy/model-capacity.svg Normal file
View file

@ -0,0 +1,52 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 1000 560" font-family="-apple-system, Segoe UI, Helvetica, Arial, sans-serif" font-size="14">
<defs>
<marker id="a" viewBox="0 0 10 10" refX="9" refY="5" markerWidth="7" markerHeight="7" orient="auto"><path d="M0 0L10 5L0 10z" fill="#475569"/></marker>
<marker id="r" viewBox="0 0 10 10" refX="9" refY="5" markerWidth="7" markerHeight="7" orient="auto"><path d="M0 0L10 5L0 10z" fill="#dc2626"/></marker>
</defs>
<rect width="1000" height="560" fill="#ffffff"/>
<text x="24" y="34" font-size="18" font-weight="600" fill="#0f172a">Model capacity hook (dynamic_rate_limiter_v3): fleet-wide group capacity and priority shares</text>
<rect x="24" y="60" width="300" height="160" rx="8" fill="#eff6ff" stroke="#3b82f6"/>
<text x="38" y="84" font-weight="600" fill="#1e3a8a">At startup</text>
<text x="38" y="106" fill="#1e293b" font-size="13">group capacity = sum of deployments'</text>
<text x="38" y="124" fill="#1e293b" font-size="13"> rpm / tpm (deployment, litellm_params,</text>
<text x="38" y="142" fill="#1e293b" font-size="13"> model_info)</text>
<text x="38" y="164" fill="#1e293b" font-size="13">priority_reservation shares: percent,</text>
<text x="38" y="182" fill="#1e293b" font-size="13"> rpm or tpm; normalized when over 1</text>
<text x="38" y="204" fill="#64748b" font-size="13">on when callbacks list the hook</text>
<rect x="360" y="60" width="616" height="160" rx="8" fill="#fffbeb" stroke="#f59e0b"/>
<text x="374" y="84" font-weight="600" fill="#92400e">Per request (one Redis script, all or nothing)</text>
<text x="374" y="106" fill="#1e293b" font-size="13">priority: team metadata, else key metadata; pool &lt;group&gt;:&lt;priority&gt; or &lt;group&gt;:default_pool</text>
<text x="374" y="128" fill="#1e293b" font-size="13">saturation = max(requests / rpm, tokens / tpm), read and reused for saturation_check_cache_ttl</text>
<text x="374" y="150" fill="#1e293b" font-size="13">{model_saturation_check:&lt;group&gt;}: +1 request, tokens checked; always enforced</text>
<text x="374" y="172" fill="#1e293b" font-size="13">{priority_model:&lt;pool&gt;}: share x capacity, enforced at saturation &gt;= threshold,</text>
<text x="374" y="190" fill="#1e293b" font-size="13"> only counted below it</text>
<line x1="324" y1="140" x2="356" y2="140" stroke="#475569" stroke-width="1.5" marker-end="url(#a)"/>
<rect x="24" y="256" width="460" height="130" rx="8" fill="#fef2f2" stroke="#dc2626"/>
<text x="38" y="280" font-weight="600" fill="#991b1b">Refusal: 429, the proxy's messages</text>
<text x="38" y="302" fill="#1e293b" font-size="13">"Model capacity reached for &lt;group&gt;. Priority: ..., Rate limit</text>
<text x="38" y="320" fill="#1e293b" font-size="13"> type: ..., Model TPM: ..., Model RPM: ..., Remaining: ..."</text>
<text x="38" y="340" fill="#1e293b" font-size="13">"Priority-based rate limit exceeded. ... Model saturation: 62.5%"</text>
<text x="38" y="362" fill="#1e293b" font-size="13">chat: retry-after, rate_limit_type, x-litellm-priority,</text>
<text x="38" y="378" fill="#1e293b" font-size="13"> x-litellm-saturation; messages: Anthropic envelope only</text>
<path d="M500 220 L500 240 L254 240 L254 252" fill="none" stroke="#dc2626" stroke-width="1.5" marker-end="url(#r)"/>
<rect x="516" y="256" width="460" height="130" rx="8" fill="#f8fafc" stroke="#64748b"/>
<text x="530" y="280" font-weight="600" fill="#0f172a">Admitted: request runs</text>
<text x="530" y="302" fill="#334155" font-size="13">complete responses carry x-litellm-priority,</text>
<text x="530" y="320" fill="#334155" font-size="13">x-litellm-rate-limiter-version: v3 and</text>
<text x="530" y="338" fill="#334155" font-size="13">x-ratelimit-model_saturation_check-* /</text>
<text x="530" y="356" fill="#334155" font-size="13">x-ratelimit-priority_model-*</text>
<path d="M760 220 L760 252" fill="none" stroke="#475569" stroke-width="1.5" marker-end="url(#a)"/>
<rect x="24" y="420" width="952" height="56" rx="8" fill="#f5f3ff" stroke="#8b5cf6"/>
<text x="38" y="444" font-weight="600" fill="#5b21b6">Settle: a successful response's tokens (token_rate_limit_type: total, input or output)</text>
<text x="38" y="464" fill="#1e293b" font-size="13">added to the group's and the admitted pool's :tokens, with the window TTL</text>
<line x1="746" y1="386" x2="746" y2="416" stroke="#475569" stroke-width="1.5" marker-end="url(#a)"/>
<rect x="24" y="496" width="952" height="52" rx="8" fill="#ecfdf5" stroke="#10b981"/>
<text x="38" y="518" font-weight="600" fill="#065f46">Same keys and scripts as the Python proxy, under its namespace: Rust and Python replicas share one count</text>
<text x="38" y="538" fill="#64748b" font-size="13">No Redis, or Redis unreachable: this replica's counters (fail_closed_rate_limit_enforcement: 503 instead)</text>
</svg>
Side by side

After

Width:  |  Height:  |  Size: 4.8 KiB

54
litellm-rust/docs/customer-proxy/prometheus.svg Normal file
View file

@ -0,0 +1,54 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 1000 470" font-family="-apple-system, Segoe UI, Helvetica, Arial, sans-serif" font-size="14">
<defs>
<marker id="a" viewBox="0 0 10 10" refX="9" refY="5" markerWidth="7" markerHeight="7" orient="auto"><path d="M0 0L10 5L0 10z" fill="#475569"/></marker>
</defs>
<rect width="1000" height="470" fill="#ffffff"/>
<text x="24" y="34" font-size="18" font-weight="600" fill="#0f172a">prometheus observer: the proxy's request metrics, off the request path</text>
<rect x="24" y="60" width="290" height="170" rx="8" fill="#eff6ff" stroke="#3b82f6"/>
<text x="38" y="84" font-weight="600" fill="#1e3a8a">At startup</text>
<text x="38" y="106" fill="#1e293b" font-size="13">on when callbacks, success_callback</text>
<text x="38" y="124" fill="#1e293b" font-size="13"> or failure_callback list prometheus</text>
<text x="38" y="146" fill="#1e293b" font-size="13">prometheus_metrics_config,</text>
<text x="38" y="164" fill="#1e293b" font-size="13"> _exclude_metrics, _exclude_labels,</text>
<text x="38" y="182" fill="#1e293b" font-size="13"> _latency_buckets as the proxy reads them</text>
<text x="38" y="204" fill="#64748b" font-size="13">a label a metric lacks stops startup</text>
<rect x="350" y="60" width="300" height="170" rx="8" fill="#fffbeb" stroke="#f59e0b"/>
<text x="364" y="84" font-weight="600" fill="#92400e">Request finishes (observer)</text>
<text x="364" y="106" fill="#1e293b" font-size="13">record: key, team, user, end user,</text>
<text x="364" y="124" fill="#1e293b" font-size="13"> deployment, usage, cost, timings</text>
<text x="364" y="146" fill="#1e293b" font-size="13">label values hashed, no allocation;</text>
<text x="364" y="164" fill="#1e293b" font-size="13"> a series is made once, on first sight</text>
<text x="364" y="186" fill="#1e293b" font-size="13">atomic adds; nothing awaits</text>
<text x="364" y="208" fill="#64748b" font-size="13">the response is already on its way</text>
<line x1="314" y1="145" x2="346" y2="145" stroke="#475569" stroke-width="1.5" marker-end="url(#a)"/>
<rect x="686" y="60" width="290" height="170" rx="8" fill="#f0fdf4" stroke="#22c55e"/>
<text x="700" y="84" font-weight="600" fill="#14532d">GET /metrics (scrape)</text>
<text x="700" y="106" fill="#1e293b" font-size="13">gateway ops metrics, then the</text>
<text x="700" y="124" fill="#1e293b" font-size="13"> observer's families in text format</text>
<text x="700" y="146" fill="#1e293b" font-size="13">shared names (total requests, total</text>
<text x="700" y="164" fill="#1e293b" font-size="13"> latency, deployment failures) come</text>
<text x="700" y="182" fill="#1e293b" font-size="13"> from the observer alone, proxy labels</text>
<line x1="650" y1="145" x2="682" y2="145" stroke="#475569" stroke-width="1.5" marker-end="url(#a)"/>
<rect x="24" y="260" width="460" height="186" rx="8" fill="#f8fafc" stroke="#94a3b8"/>
<text x="38" y="284" font-weight="600" fill="#0f172a">Success</text>
<text x="38" y="306" fill="#1e293b" font-size="13">proxy_total_requests, requests, spend</text>
<text x="38" y="326" fill="#1e293b" font-size="13">total / input / output / cached / cache creation tokens</text>
<text x="38" y="346" fill="#1e293b" font-size="13">request_total_latency, llm_api_latency (from upstream send)</text>
<text x="38" y="366" fill="#1e293b" font-size="13">llm_api_time_to_first_token (streams only)</text>
<text x="38" y="386" fill="#1e293b" font-size="13">deployment_state = 0, deployment success / total</text>
<text x="38" y="406" fill="#1e293b" font-size="13">deployment_latency_per_output_token</text>
<text x="38" y="430" fill="#64748b" font-size="13">metric and label names are the proxy's</text>
<rect x="516" y="260" width="460" height="186" rx="8" fill="#fef2f2" stroke="#ef4444"/>
<text x="530" y="284" font-weight="600" fill="#7f1d1d">Failure</text>
<text x="530" y="306" fill="#1e293b" font-size="13">proxy_failed_requests + proxy_total_requests, always</text>
<text x="530" y="326" fill="#1e293b" font-size="13">exception_class as the proxy names it</text>
<text x="530" y="346" fill="#1e293b" font-size="13"> (e.g. Openai.RateLimitError, BudgetExceededError)</text>
<text x="530" y="370" fill="#1e293b" font-size="13">after routing only: llm_api_failed_requests,</text>
<text x="530" y="390" fill="#1e293b" font-size="13"> deployment_state = 1, deployment failure / total</text>
<text x="530" y="430" fill="#64748b" font-size="13">refusals before routing carry no deployment labels</text>
</svg>
Side by side

After

Width:  |  Height:  |  Size: 4.5 KiB

56
litellm-rust/docs/customer-proxy/request-hooks.svg Normal file
View file

@ -0,0 +1,56 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 1000 560" font-family="-apple-system, Segoe UI, Helvetica, Arial, sans-serif" font-size="14">
<defs>
<marker id="a" viewBox="0 0 10 10" refX="9" refY="5" markerWidth="7" markerHeight="7" orient="auto"><path d="M0 0L10 5L0 10z" fill="#475569"/></marker>
</defs>
<rect width="1000" height="560" fill="#ffffff"/>
<text x="24" y="34" font-size="18" font-weight="600" fill="#0f172a">Configurable request hooks: litellm_settings.hooks, run top to bottom</text>
<rect x="24" y="60" width="230" height="200" rx="8" fill="#eff6ff" stroke="#3b82f6"/>
<text x="38" y="84" font-weight="600" fill="#1e3a8a">Client request</text>
<text x="38" y="108" fill="#1e293b" font-size="13">x-trace-id: trace-1</text>
<text x="38" y="126" fill="#1e293b" font-size="13">x-cost-center: cc-9</text>
<text x="38" y="144" fill="#1e293b" font-size="13">x-litellm-spend-logs-</text>
<text x="38" y="162" fill="#1e293b" font-size="13"> metadata: {...}</text>
<text x="38" y="180" fill="#1e293b" font-size="13">anthropic-beta: ...</text>
<text x="38" y="198" fill="#1e293b" font-size="13">authorization: Bearer sk-...</text>
<text x="38" y="226" fill="#64748b" font-size="13">key: user, team alias,</text>
<text x="38" y="244" fill="#64748b" font-size="13">metadata; end user</text>
<line x1="254" y1="160" x2="286" y2="160" stroke="#475569" stroke-width="1.5" marker-end="url(#a)"/>
<rect x="290" y="60" width="430" height="44" rx="8" fill="#fffbeb" stroke="#f59e0b"/>
<text x="304" y="87" fill="#1e293b" font-size="13"><tspan font-weight="600">forward_headers</tspan> copies listed client headers (exact or prefix-*)</text>
<rect x="290" y="112" width="430" height="44" rx="8" fill="#fffbeb" stroke="#f59e0b"/>
<text x="304" y="139" fill="#1e293b" font-size="13"><tspan font-weight="600">spend_metadata</tspan> headers (json or one key) + key metadata</text>
<rect x="290" y="164" width="430" height="44" rx="8" fill="#fffbeb" stroke="#f59e0b"/>
<text x="304" y="191" fill="#1e293b" font-size="13"><tspan font-weight="600">strip_headers</tspan> removes forwarded and pass-through headers</text>
<rect x="290" y="216" width="430" height="44" rx="8" fill="#fffbeb" stroke="#f59e0b"/>
<text x="304" y="243" fill="#1e293b" font-size="13"><tspan font-weight="600">provider_metadata</tspan> "{end_user|key.user_id}" templates</text>
<line x1="505" y1="104" x2="505" y2="110" stroke="#475569" stroke-width="1.5"/>
<line x1="505" y1="156" x2="505" y2="162" stroke="#475569" stroke-width="1.5"/>
<line x1="505" y1="208" x2="505" y2="214" stroke="#475569" stroke-width="1.5"/>
<rect x="756" y="60" width="220" height="200" rx="8" fill="#ecfdf5" stroke="#10b981"/>
<text x="770" y="84" font-weight="600" fill="#065f46">Provider request</text>
<text x="770" y="108" fill="#1e293b" font-size="13">x-trace-id: trace-1</text>
<text x="770" y="126" fill="#1e293b" font-size="13">x-caller-team: research</text>
<text x="770" y="144" fill="#1e293b" font-size="13">metadata.user_id: user-1</text>
<text x="770" y="162" fill="#1e293b" font-size="13"> (messages only)</text>
<text x="770" y="186" fill="#64748b" font-size="13">no x-cost-center, no</text>
<text x="770" y="204" fill="#64748b" font-size="13">anthropic-beta, never</text>
<text x="770" y="222" fill="#64748b" font-size="13">client credentials</text>
<line x1="720" y1="160" x2="752" y2="160" stroke="#475569" stroke-width="1.5" marker-end="url(#a)"/>
<rect x="290" y="296" width="430" height="70" rx="8" fill="#f5f3ff" stroke="#8b5cf6"/>
<text x="304" y="320" font-weight="600" fill="#5b21b6">Spend log</text>
<text x="304" y="342" fill="#1e293b" font-size="13">the spend_metadata annotation lands in</text>
<text x="304" y="358" fill="#1e293b" font-size="13">LiteLLM_SpendLogs.metadata.spend_logs_metadata</text>
<path d="M290 134 L270 134 L270 330 L286 330" fill="none" stroke="#8b5cf6" stroke-width="1.5" marker-end="url(#a)"/>
<rect x="24" y="400" width="952" height="140" rx="8" fill="#f8fafc" stroke="#64748b"/>
<text x="38" y="424" font-weight="600" fill="#0f172a">Rules</text>
<text x="38" y="448" fill="#1e293b" font-size="13">Each hook declares what it reads and writes; a later hook in the list wins on the same header or field.</text>
<text x="38" y="468" fill="#1e293b" font-size="13">Header edits reach the provider on both endpoints and every provider, Bedrock's signed requests included.</text>
<text x="38" y="488" fill="#1e293b" font-size="13">Credentials, signer and framing headers (authorization, x-api-key, cookie, x-amz-*, ...) are never forwarded or set.</text>
<text x="38" y="508" fill="#1e293b" font-size="13">An unknown hook, setting or template fact stops the gateway at startup; a bad JSON header is dropped with a warning.</text>
<text x="38" y="528" fill="#64748b" font-size="13">Templates name facts only: end_user, model, request_id, key.alias|user_id|team_id|team_alias, key.metadata.*, team.metadata.*, header.*</text>
</svg>
Side by side

After

Width:  |  Height:  |  Size: 4.9 KiB

61
litellm-rust/docs/customer-proxy/s3-log.svg Normal file
View file

@ -0,0 +1,61 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 1000 490" font-family="-apple-system, Segoe UI, Helvetica, Arial, sans-serif" font-size="14">
<defs>
<marker id="a" viewBox="0 0 10 10" refX="9" refY="5" markerWidth="7" markerHeight="7" orient="auto"><path d="M0 0L10 5L0 10z" fill="#475569"/></marker>
</defs>
<rect width="1000" height="490" fill="#ffffff"/>
<text x="24" y="34" font-size="18" font-weight="600" fill="#0f172a">s3_v2 observer: each request's logging payload as an S3 object, off the request path</text>
<rect x="24" y="60" width="290" height="190" rx="8" fill="#eff6ff" stroke="#3b82f6"/>
<text x="38" y="84" font-weight="600" fill="#1e3a8a">At startup</text>
<text x="38" y="106" fill="#1e293b" font-size="13">on when callbacks lists s3_v2 (or s3_log);</text>
<text x="38" y="124" fill="#1e293b" font-size="13"> success_callback / failure_callback</text>
<text x="38" y="142" fill="#1e293b" font-size="13"> log only successes / only failures</text>
<text x="38" y="164" fill="#1e293b" font-size="13">s3_callback_params, os.environ/ resolved;</text>
<text x="38" y="182" fill="#1e293b" font-size="13"> bad numbers and flags warn, use defaults</text>
<text x="38" y="204" fill="#1e293b" font-size="13">credentials: keys, profile, role, web</text>
<text x="38" y="222" fill="#1e293b" font-size="13"> identity or the default chain</text>
<text x="38" y="242" fill="#64748b" font-size="13">no s3_bucket_name stops startup</text>
<rect x="350" y="60" width="300" height="190" rx="8" fill="#fffbeb" stroke="#f59e0b"/>
<text x="364" y="84" font-weight="600" fill="#92400e">Request finishes (observe)</text>
<text x="364" y="106" fill="#1e293b" font-size="13">StandardLoggingPayload: key, team,</text>
<text x="364" y="124" fill="#1e293b" font-size="13"> deployment, usage, cost, timings,</text>
<text x="364" y="142" fill="#1e293b" font-size="13"> messages, response (streams rebuilt)</text>
<text x="364" y="164" fill="#1e293b" font-size="13">turn_off_message_logging, strip_base64,</text>
<text x="364" y="182" fill="#1e293b" font-size="13"> log_prompts_only applied here</text>
<text x="364" y="204" fill="#1e293b" font-size="13">queued; past s3_max_queue_size the</text>
<text x="364" y="222" fill="#1e293b" font-size="13"> oldest are dropped with a warning</text>
<text x="364" y="242" fill="#64748b" font-size="13">nothing awaits S3 here</text>
<line x1="314" y1="155" x2="346" y2="155" stroke="#475569" stroke-width="1.5" marker-end="url(#a)"/>
<rect x="686" y="60" width="290" height="190" rx="8" fill="#f0fdf4" stroke="#22c55e"/>
<text x="700" y="84" font-weight="600" fill="#14532d">Every 10 s, and at shutdown</text>
<text x="700" y="106" fill="#1e293b" font-size="13">the whole queue, s3_max_concurrent_uploads</text>
<text x="700" y="124" fill="#1e293b" font-size="13"> PUTs at a time (16), SigV4 signed</text>
<text x="700" y="146" fill="#1e293b" font-size="13">s3_batch_file_upload: one .jsonl per</text>
<text x="700" y="164" fill="#1e293b" font-size="13"> folder instead of one .json per request</text>
<text x="700" y="186" fill="#1e293b" font-size="13">failed for now: back at the front, sent</text>
<text x="700" y="204" fill="#1e293b" font-size="13"> after newer ones on the next flush</text>
<text x="700" y="242" fill="#64748b" font-size="13">DEFAULT_S3_FLUSH_INTERVAL_SECONDS</text>
<line x1="650" y1="155" x2="682" y2="155" stroke="#475569" stroke-width="1.5" marker-end="url(#a)"/>
<rect x="24" y="280" width="460" height="186" rx="8" fill="#f8fafc" stroke="#94a3b8"/>
<text x="38" y="304" font-weight="600" fill="#0f172a">Object key</text>
<text x="38" y="328" fill="#1e293b" font-size="13" font-family="ui-monospace, Menlo, monospace">{s3_path}/[team_alias/][key_alias/]</text>
<text x="38" y="348" fill="#1e293b" font-size="13" font-family="ui-monospace, Menlo, monospace"> YYYY-MM-DD[/HH]/time-HH-MM-SS-ffffff_{id}.json</text>
<text x="38" y="372" fill="#1e293b" font-size="13">prefixes with s3_use_team_prefix / s3_use_key_prefix</text>
<text x="38" y="392" fill="#1e293b" font-size="13">/HH with s3_partition_granularity: hour</text>
<text x="38" y="412" fill="#1e293b" font-size="13">id: the response's id, else the request's call id</text>
<text x="38" y="432" fill="#1e293b" font-size="13">past S3's 1024 / 255 byte limits: head + sha256</text>
<text x="38" y="454" fill="#64748b" font-size="13">path style, or s3_use_virtual_hosted_style</text>
<rect x="516" y="280" width="460" height="186" rx="8" fill="#fef2f2" stroke="#ef4444"/>
<text x="530" y="304" font-weight="600" fill="#7f1d1d">Upload and retries</text>
<text x="530" y="328" fill="#1e293b" font-size="13">3 attempts; transport error, 403, 429, 500, 503</text>
<text x="530" y="348" fill="#1e293b" font-size="13"> retried after 0.5-1.5 s, then 1-3 s</text>
<text x="530" y="370" fill="#1e293b" font-size="13">400 / 403 with an object error code (EntityTooLarge,</text>
<text x="530" y="390" fill="#1e293b" font-size="13"> BadDigest, ...) dropped (s3_drop_on_terminal_error)</text>
<text x="530" y="412" fill="#1e293b" font-size="13">kept until s3_max_retry_age_seconds (3600)</text>
<text x="530" y="432" fill="#1e293b" font-size="13">nothing delivered: all kept, the bucket is down</text>
<text x="530" y="454" fill="#64748b" font-size="13">Content-MD5, SSE and Content-Disposition as the proxy</text>
</svg>
Side by side

After

Width:  |  Height:  |  Size: 5.3 KiB