mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-27 06:13:07 +00:00
feat(gateway): configurable header, spend metadata and provider metadata hooks (#47054)
Adds four built-in hooks a config lists under litellm_settings.hooks and runs in that order: forward_headers copies named client headers to the provider, strip_headers removes headers from the provider request (including the gateway's default pass-through), spend_metadata adds header and key-metadata values to the spend log's spend_logs_metadata, and provider_metadata sets Messages body fields and headers from templates over request facts. Hooks' upstream header edits now reach the provider on both endpoints, and body edits can set nested fields. Credentials and signer headers are never forwarded or set, and a misconfigured hook stops the gateway at startup. Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5 Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Nate Armstrong
GitHub
parent
0919823385
commit
40d4c8b2d8
10 changed files with 1101 additions and 18 deletions
|
|
@ -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();
|
||||
|
|
|
|||
|
|
@ -93,6 +93,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,7 +132,7 @@ 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 call = ChatCompletionsCall {
|
||||
|
|
@ -130,7 +142,7 @@ 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,
|
||||
};
|
||||
Ok(litellm_host_http::serve_unary(
|
||||
|
|
@ -166,6 +178,7 @@ 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,7 +187,7 @@ 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,
|
||||
},
|
||||
None,
|
||||
|
|
|
|||
|
|
@ -154,6 +154,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) {
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
|
|
|
|||
|
|
@ -20,6 +20,8 @@ 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("coordination Redis: {0}")]
|
||||
Redis(String),
|
||||
#[error("MCP requires a configured master key")]
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ mod key_status;
|
|||
mod keys;
|
||||
pub mod metrics;
|
||||
pub mod rate_limits;
|
||||
pub mod request_hooks;
|
||||
mod secrets;
|
||||
pub mod spend;
|
||||
pub use error::Error;
|
||||
|
|
|
|||
|
|
@ -128,6 +128,7 @@ async fn main() -> Result<(), Box<dyn Error>> {
|
|||
),
|
||||
));
|
||||
}
|
||||
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();
|
||||
|
|
|
|||
556
litellm-rust/crates/gateway/src/request_hooks.rs
Normal file
556
litellm-rust/crates/gateway/src/request_hooks.rs
Normal 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);
|
||||
}
|
||||
}
|
||||
339
litellm-rust/crates/gateway/tests/request_hooks.rs
Normal file
339
litellm-rust/crates/gateway/tests/request_hooks.rs
Normal file
|
|
@ -0,0 +1,339 @@
|
|||
//! 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,
|
||||
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}");
|
||||
}
|
||||
56
litellm-rust/docs/customer-proxy/request-hooks.svg
Normal file
56
litellm-rust/docs/customer-proxy/request-hooks.svg
Normal 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>
|
||||
|
After Width: | Height: | Size: 4.9 KiB |
Loading…
Add table
Reference in a new issue