diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 8e88c0ea15b..70c2d600ba3 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -332,12 +332,12 @@ Each translation is isolated in its own file, making it easy to test and modify |--------------|----------|------------------| | `/v1/chat/completions` | Anthropic | `llms/anthropic/chat/transformation.py` | | `/v1/chat/completions` | Bedrock Converse | `llms/bedrock/chat/converse_transformation.py` | -| `/v1/chat/completions` | Bedrock Invoke | `llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py` | +| `/v1/chat/completions` | Bedrock Invoke | `llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py` | | `/v1/chat/completions` | Gemini | `llms/gemini/chat/transformation.py` | | `/v1/chat/completions` | Vertex AI | `llms/vertex_ai/gemini/transformation.py` | | `/v1/chat/completions` | OpenAI | `llms/openai/chat/gpt_transformation.py` | | `/v1/messages` (passthrough) | Anthropic | `llms/anthropic/pass_through/messages/transformation.py` | -| `/v1/messages` (passthrough) | Bedrock | `llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py` | +| `/v1/messages` (passthrough) | Bedrock | `llms/bedrock/messages/invoke_transformations/anthropic_claude_transformation.py` | | `/v1/messages` (passthrough) | Vertex AI | `llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py` | | Passthrough endpoints | All | `proxy/pass_through_endpoints/llm_provider_handlers/` | @@ -346,7 +346,7 @@ Each translation is isolated in its own file, making it easy to test and modify If `/v1/messages` → Bedrock Converse prompt caching isn't working but Bedrock Invoke works: 1. **Bedrock Converse translation**: `llms/bedrock/chat/converse_transformation.py` -2. **Bedrock Invoke translation**: `llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py` +2. **Bedrock Invoke translation**: `llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py` 3. Compare how each handles `cache_control` in `transform_request()` ### How translations work diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 32ae431f863..63137949f38 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -876,6 +876,15 @@ dependencies = [ "time", ] +[[package]] +name = "backon" +version = "1.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cffb0e931875b666fc4fcb20fee52e9bbd1ef836fd9e9e04ec21555f9f85f7ef" +dependencies = [ + "fastrand", +] + [[package]] name = "base16ct" version = "0.2.0" @@ -1248,7 +1257,11 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cfc320937d09e6de266b31b9afb480f197d7a861be86be7cb2ea7e5d1bfffc5e" dependencies = [ "bytes", + "futures-core", "memchr", + "pin-project-lite", + "tokio", + "tokio-util", ] [[package]] @@ -3642,6 +3655,14 @@ dependencies = [ "strum", ] +[[package]] +name = "litellm-circuit-breaker" +version = "0.1.0" +dependencies = [ + "proptest", + "rstest", +] + [[package]] name = "litellm-config" version = "0.1.0" @@ -3655,6 +3676,8 @@ dependencies = [ "strum", "tempfile", "thiserror 2.0.19", + "tracing", + "tracing-subscriber", ] [[package]] @@ -3743,6 +3766,7 @@ dependencies = [ "futures-util", "http-body-util", "litellm-auth-types", + "litellm-circuit-breaker", "litellm-config", "litellm-gateway-auth", "litellm-gateway-inference", @@ -3751,6 +3775,8 @@ dependencies = [ "litellm-http", "litellm-inference", "litellm-llms", + "litellm-redis-client", + "litellm-redis-testing", "litellm-secrets", "litellm-tracing", "rstest", @@ -4282,6 +4308,28 @@ dependencies = [ "thiserror 2.0.19", ] +[[package]] +name = "litellm-redis-client" +version = "0.1.0" +dependencies = [ + "futures-util", + "litellm-circuit-breaker", + "litellm-redis-testing", + "redis", + "rstest", + "thiserror 2.0.19", + "tokio", +] + +[[package]] +name = "litellm-redis-testing" +version = "0.1.0" +dependencies = [ + "testcontainers-modules", + "thiserror 2.0.19", + "tokio", +] + [[package]] name = "litellm-router" version = "0.1.0" @@ -5757,22 +5805,34 @@ dependencies = [ [[package]] name = "redis" -version = "1.7.0" +version = "1.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2acbc41a996f7652b2ddd9dfd98cc4ff602cfd742ae35382f07f608405ab50ed" +checksum = "7acd809e80afa754c63f586955699d8990efe1c4b6613183ee5f027a11fbc8b5" dependencies = [ + "arc-swap", "arcstr", + "async-lock", + "backon", + "bytes", + "cfg-if", "combine", "crc16", + "futures-channel", + "futures-util", "itoa", + "log", "num-bigint 0.5.1", "percent-encoding", + "pin-project-lite", "rand 0.10.2", "rustls 0.23.42", "rustls-native-certs", "ryu", "sha1_smol", "socket2 0.6.5", + "tokio", + "tokio-rustls 0.26.4", + "tokio-util", "url", "xxhash-rust", ] diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index fc65355c047..dd8d8704370 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -69,6 +69,9 @@ litellm-token-counter-huggingface = { path = "crates/token-counter-huggingface" litellm-token-counter-tiktoken = { path = "crates/token-counter-tiktoken" } litellm-host-python = { path = "crates/host-python" } litellm-python-compat = { path = "crates/python-compat" } +litellm-circuit-breaker = { path = "crates/circuit-breaker" } +litellm-redis-client = { path = "crates/redis-client" } +litellm-redis-testing = { path = "crates/redis-testing" } askama = { version = "0.16.1", default-features = false, features = ["derive", "std"] } tracing = "0.1" @@ -88,6 +91,9 @@ rand = "0.8" macro_rules_attribute = "0.2.3" schemars = "1" reqwest = { version = "0.12", default-features = false, features = ["json", "multipart", "rustls-tls", "http2", "stream"] } +redis = "1.8.0" +redis-test = "1.0.4" +testcontainers-modules = "0.15.0" qdrant-client = { version = "1.19.0", default-features = false } uuid = { version = "1", features = ["v4"] } rstest = "0.26.1" diff --git a/litellm-rust/crates/cache-redis-semantic/Cargo.toml b/litellm-rust/crates/cache-redis-semantic/Cargo.toml index fe5a317cec0..6a5c794bde7 100644 --- a/litellm-rust/crates/cache-redis-semantic/Cargo.toml +++ b/litellm-rust/crates/cache-redis-semantic/Cargo.toml @@ -8,12 +8,12 @@ repository.workspace = true [dependencies] litellm-cache.workspace = true litellm-cache-redis.workspace = true -redis = { version = "1.7.0", features = ["tls-rustls"] } +redis = { workspace = true, features = ["tls-rustls"] } sha2.workspace = true [dev-dependencies] litellm-cache-testing.workspace = true -redis-test = "1.0.4" +redis-test.workspace = true rstest.workspace = true serde_json.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/cache-redis/Cargo.toml b/litellm-rust/crates/cache-redis/Cargo.toml index a234286a338..bebfcafd98a 100644 --- a/litellm-rust/crates/cache-redis/Cargo.toml +++ b/litellm-rust/crates/cache-redis/Cargo.toml @@ -7,12 +7,12 @@ repository.workspace = true [dependencies] litellm-cache.workspace = true -redis = { version = "1.7.0", features = ["cluster", "tls-rustls"] } +redis = { workspace = true, features = ["cluster", "tls-rustls"] } r2d2 = "0.8.10" tokio.workspace = true [dev-dependencies] litellm-cache-testing.workspace = true -redis-test = "1.0.4" +redis-test.workspace = true rstest.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/cache-response/Cargo.toml b/litellm-rust/crates/cache-response/Cargo.toml index 869c40a12ab..8f4570edc6f 100644 --- a/litellm-rust/crates/cache-response/Cargo.toml +++ b/litellm-rust/crates/cache-response/Cargo.toml @@ -17,8 +17,8 @@ litellm-cache-gcs.workspace = true litellm-http = { workspace = true, features = ["test-support"] } litellm-cache-memory.workspace = true litellm-cache-redis.workspace = true -redis = "1.7.0" -redis-test = "1.0.4" +redis.workspace = true +redis-test.workspace = true rstest.workspace = true tokio.workspace = true wiremock.workspace = true diff --git a/litellm-rust/crates/cache-valkey-semantic/Cargo.toml b/litellm-rust/crates/cache-valkey-semantic/Cargo.toml index da9bb60c853..ec755e61b8c 100644 --- a/litellm-rust/crates/cache-valkey-semantic/Cargo.toml +++ b/litellm-rust/crates/cache-valkey-semantic/Cargo.toml @@ -8,13 +8,13 @@ repository.workspace = true [dependencies] litellm-cache.workspace = true litellm-cache-redis.workspace = true -redis = { version = "1.7.0", features = ["tls-rustls"] } +redis = { workspace = true, features = ["tls-rustls"] } sha2.workspace = true uuid = { version = "1", features = ["v4"] } [dev-dependencies] litellm-cache-testing.workspace = true -redis-test = "1.0.4" +redis-test.workspace = true rstest.workspace = true serde_json.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/circuit-breaker/AGENTS.md b/litellm-rust/crates/circuit-breaker/AGENTS.md new file mode 100644 index 00000000000..52ef0ca7435 --- /dev/null +++ b/litellm-rust/crates/circuit-breaker/AGENTS.md @@ -0,0 +1,14 @@ +- `fsm` is a generic state machine runtime and knows nothing about breakers. `breaker` is a generic circuit breaker and knows nothing about Redis. Deciding which errors count as failures belongs to the consumer +- `Machine::step` stays pure. `&self` is config and time comes in on the event, never read inside the machine. The only clock read is `Admission`'s drop, at the call-site boundary +- Transition rules are one `match (state, event)` per machine, with no macro DSL and no `_` arm that would hide a new state from the compiler +- `Shared` counts a transition, bumps the generation and calls the observer only when the state's enum variant changes, matching Python bumping `generation` only when `_state` changes. Data inside a variant (failure counters, `since`) changes without a transition +- The observer runs under the lock so transitions are seen in generation order. Keep it to logging or a metric +- A disabled breaker is `Option::None` at the consumer, not a flag inside the machine. `Config` comes in already resolved, reading `REDIS_CIRCUIT_BREAKER_*` env is the consumer's job + +## Parity with Python `RedisCircuitBreaker` (`litellm/caching/redis_cache.py`) +- The rstest table `matches_the_python_redis_circuit_breaker` was checked case by case against the Python breaker. A behavior change here needs the same check, by driving the Python class with the same steps +- Divergences, each on purpose + - monotonic `Instant` instead of `time.time()`, so a wall clock jump can't open or close the breaker + - a probe `Admission` dropped without an outcome records a hard failure and reopens. Python's probe that dies on `CancelledError` (a `BaseException`) records nothing and leaves the breaker HALF_OPEN for good, and a dropped Rust future is common + - `_failure_count` is dropped, it only fed a log line + - no `enabled` flag, see above diff --git a/litellm-rust/crates/circuit-breaker/Cargo.toml b/litellm-rust/crates/circuit-breaker/Cargo.toml new file mode 100644 index 00000000000..8eca448a2ac --- /dev/null +++ b/litellm-rust/crates/circuit-breaker/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "litellm-circuit-breaker" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] + +[dev-dependencies] +proptest.workspace = true +rstest.workspace = true diff --git a/litellm-rust/crates/circuit-breaker/src/breaker.rs b/litellm-rust/crates/circuit-breaker/src/breaker.rs new file mode 100644 index 00000000000..c9509561b37 --- /dev/null +++ b/litellm-rust/crates/circuit-breaker/src/breaker.rs @@ -0,0 +1,444 @@ +use std::{ + sync::Arc, + time::{Duration, Instant}, +}; + +use crate::{Machine, Observer, Shared}; + +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct Config { + pub failure_threshold: u32, + pub recovery_timeout: Duration, + pub timeout_min_duration: Duration, +} + +impl Default for Config { + fn default() -> Self { + Self { + failure_threshold: 5, + recovery_timeout: Duration::from_secs(60), + timeout_min_duration: Duration::from_secs(5), + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Failure { + Hard, + Timeout, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum State { + Closed { + hard: u32, + timeouts: u32, + streak_started: Option, + }, + Open { + since: Instant, + }, + HalfOpen, +} + +impl State { + const CLOSED: Self = Self::Closed { + hard: 0, + timeouts: 0, + streak_started: None, + }; +} + +#[derive(Clone, Copy, Debug)] +pub enum Event { + Admit(Instant), + Succeeded, + Failed(Failure, Instant), +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Output { + Admitted, + Probe, + Rejected, + Recorded, +} + +impl Machine for Config { + type State = State; + type Event = Event; + type Output = Output; + + fn step(&self, state: &State, event: Event) -> (State, Output) { + match (state, event) { + (State::Closed { .. }, Event::Admit(_)) => (state.clone(), Output::Admitted), + (State::Open { since }, Event::Admit(now)) + if now.duration_since(*since) > self.recovery_timeout => + { + (State::HalfOpen, Output::Probe) + } + (State::Open { .. } | State::HalfOpen, Event::Admit(_)) => { + (state.clone(), Output::Rejected) + } + (State::Closed { .. } | State::HalfOpen, Event::Succeeded) => { + (State::CLOSED, Output::Recorded) + } + (State::Open { .. }, Event::Succeeded) => (state.clone(), Output::Recorded), + (State::Closed { hard, .. }, Event::Failed(Failure::Hard, now)) => { + let hard = hard + 1; + let next = match hard >= self.failure_threshold { + true => State::Open { since: now }, + false => State::Closed { + hard, + timeouts: 0, + streak_started: None, + }, + }; + (next, Output::Recorded) + } + ( + State::Closed { + hard, + timeouts, + streak_started, + }, + Event::Failed(Failure::Timeout, now), + ) => { + let timeouts = timeouts + 1; + let started = streak_started.unwrap_or(now); + let opens = timeouts >= self.failure_threshold + && now.duration_since(started) >= self.timeout_min_duration; + let next = match opens { + true => State::Open { since: now }, + false => State::Closed { + hard: *hard, + timeouts, + streak_started: Some(started), + }, + }; + (next, Output::Recorded) + } + (State::Open { .. } | State::HalfOpen, Event::Failed(_, now)) => { + (State::Open { since: now }, Output::Recorded) + } + } + } +} + +#[derive(Clone)] +pub struct Breaker(Arc>); + +#[derive(Debug, PartialEq, Eq)] +pub struct Rejected; + +pub struct Admission { + breaker: Breaker, + generation: u64, + probe: bool, + settled: bool, +} + +impl Breaker { + pub fn new(config: Config, observer: Observer) -> Self { + Self(Arc::new(Shared::new(config, State::CLOSED, observer))) + } + + pub fn admit(&self, now: Instant) -> Result { + let (output, generation) = self.0.fire(Event::Admit(now)); + let probe = match output { + Output::Admitted => false, + Output::Probe => true, + Output::Rejected | Output::Recorded => return Err(Rejected), + }; + Ok(Admission { + breaker: self.clone(), + generation, + probe, + settled: false, + }) + } + + pub fn failed(&self, failure: Failure, now: Instant) { + self.0.fire(Event::Failed(failure, now)); + } + + pub fn state(&self) -> State { + self.0.state() + } +} + +impl Admission { + pub fn probe(&self) -> bool { + self.probe + } + + pub fn succeeded(mut self) { + self.settled = true; + self.breaker.0.fire_if(self.generation, Event::Succeeded); + } + + pub fn failed(mut self, failure: Failure, now: Instant) { + self.settled = true; + self.breaker.failed(failure, now); + } +} + +impl Drop for Admission { + fn drop(&mut self) { + if self.probe && !self.settled { + self.breaker.0.fire_if( + self.generation, + Event::Failed(Failure::Hard, Instant::now()), + ); + } + } +} + +#[cfg(test)] +mod tests { + use std::{ + collections::HashMap, + sync::{Arc, Mutex}, + time::{Duration, Instant}, + }; + + use proptest::prelude::*; + use rstest::rstest; + + use super::{Admission, Breaker, Config, Failure, State}; + use crate::Transition; + + #[derive(Clone, Copy, Debug)] + enum Step { + Fail(f64, Failure), + Admit(f64, bool), + Succeed(usize), + } + + #[derive(Clone, Copy, Debug, PartialEq, Eq)] + enum Phase { + Closed, + Open, + HalfOpen, + } + + use Phase::{Closed as C, HalfOpen as H, Open as O}; + + #[rstest] + fn dropping_an_old_probe_does_not_undo_a_newer_recovery() { + let breaker = Breaker::new( + Config { + failure_threshold: 1, + recovery_timeout: Duration::from_secs(1), + ..Config::default() + }, + Box::new(|_| {}), + ); + let now = Instant::now(); + let pending = breaker.admit(now).unwrap(); + breaker.failed(Failure::Hard, now); + let old_probe = breaker.admit(now + Duration::from_secs(2)).unwrap(); + assert!(old_probe.probe()); + pending.failed(Failure::Hard, now + Duration::from_secs(2)); + let new_probe = breaker.admit(now + Duration::from_secs(4)).unwrap(); + assert!(new_probe.probe()); + new_probe.succeeded(); + assert_eq!(breaker.state(), State::CLOSED); + + drop(old_probe); + + assert_eq!(breaker.state(), State::CLOSED); + assert!(!breaker.admit(now + Duration::from_secs(4)).unwrap().probe()); + } + + fn hard(at: f64) -> Step { + Step::Fail(at, Failure::Hard) + } + + fn timeout(at: f64) -> Step { + Step::Fail(at, Failure::Timeout) + } + + fn five_hard_at_zero() -> Vec { + vec![hard(0.0); 5] + } + + fn phase(state: &State) -> Phase { + match state { + State::Closed { .. } => Phase::Closed, + State::Open { .. } => Phase::Open, + State::HalfOpen => Phase::HalfOpen, + } + } + + fn quiet() -> Breaker { + Breaker::new(Config::default(), Box::new(|_| {})) + } + + #[rstest] + #[case::five_hard_failures_open(vec![hard(0.0), hard(1.0), hard(2.0), hard(3.0), hard(4.0)], vec![C, C, C, C, O])] + #[case::a_success_resets_the_hard_failure_count( + vec![hard(0.0), hard(1.0), hard(2.0), hard(3.0), Step::Admit(4.0, true), Step::Succeed(4), hard(5.0), hard(6.0), hard(7.0), hard(8.0)], + vec![C; 10] + )] + #[case::a_timeout_streak_opens_only_after_the_min_duration( + vec![timeout(0.0), timeout(1.0), timeout(2.0), timeout(3.0), timeout(4.9), timeout(5.0)], + vec![C, C, C, C, C, O] + )] + #[case::a_hard_failure_restarts_the_timeout_streak( + vec![timeout(0.0), timeout(1.0), timeout(2.0), timeout(3.0), hard(4.0), timeout(5.0), timeout(6.0), timeout(7.0), timeout(8.0), timeout(9.9), timeout(10.0)], + vec![C, C, C, C, C, C, C, C, C, C, O] + )] + #[case::the_probe_waits_until_strictly_past_the_recovery_timeout( + [five_hard_at_zero(), vec![Step::Admit(60.0, false), Step::Admit(60.5, true)]].concat(), + vec![C, C, C, C, O, O, H] + )] + #[case::a_probe_success_closes( + [five_hard_at_zero(), vec![Step::Admit(61.0, true), Step::Succeed(5)]].concat(), + vec![C, C, C, C, O, H, C] + )] + #[case::a_probe_failure_reopens_with_a_fresh_timer( + [five_hard_at_zero(), vec![Step::Admit(61.0, true), hard(62.0), Step::Admit(122.0, false), Step::Admit(122.5, true)]].concat(), + vec![C, C, C, C, O, H, O, O, H] + )] + #[case::only_one_probe_at_a_time( + [five_hard_at_zero(), vec![Step::Admit(61.0, true), Step::Admit(61.5, false)]].concat(), + vec![C, C, C, C, O, H, H] + )] + #[case::a_late_failure_while_open_pushes_back_the_probe( + [five_hard_at_zero(), vec![hard(30.0), Step::Admit(61.0, false), Step::Admit(90.5, true)]].concat(), + vec![C, C, C, C, O, O, O, H] + )] + #[case::a_success_from_before_the_breaker_moved_is_ignored( + vec![Step::Admit(0.0, true), hard(1.0), hard(1.0), hard(1.0), hard(1.0), hard(1.0), Step::Admit(62.0, true), Step::Succeed(0), Step::Succeed(6)], + vec![C, C, C, C, C, O, H, H, C] + )] + fn matches_the_python_redis_circuit_breaker( + #[case] steps: Vec, + #[case] expected: Vec, + ) { + let breaker = quiet(); + let t0 = Instant::now(); + let mut admissions: HashMap = HashMap::new(); + + let phases: Vec = steps + .iter() + .enumerate() + .map(|(index, step)| { + match *step { + Step::Fail(at, failure) => { + breaker.failed(failure, t0 + Duration::from_secs_f64(at)) + } + Step::Admit(at, admitted) => { + let admission = breaker.admit(t0 + Duration::from_secs_f64(at)); + assert_eq!(admission.is_ok(), admitted, "step {index}"); + admissions.extend(admission.ok().map(|admission| (index, admission))); + } + Step::Succeed(admitted_at) => { + admissions.remove(&admitted_at).unwrap().succeeded() + } + } + phase(&breaker.state()) + }) + .collect(); + + assert_eq!(phases, expected); + } + + #[rstest] + fn a_dropped_probe_reopens_instead_of_wedging_half_open() { + let breaker = quiet(); + let t0 = Instant::now(); + (0..5).for_each(|_| breaker.failed(Failure::Hard, t0)); + let probe = breaker.admit(t0 + Duration::from_secs(61)).unwrap(); + assert!(probe.probe()); + + drop(probe); + + assert!(matches!(breaker.state(), State::Open { .. })); + } + + #[rstest] + fn a_dropped_regular_call_records_nothing() { + let breaker = quiet(); + let t0 = Instant::now(); + breaker.failed(Failure::Hard, t0); + + drop(breaker.admit(t0).unwrap()); + + assert_eq!( + breaker.state(), + State::Closed { + hard: 1, + timeouts: 0, + streak_started: None + } + ); + } + + #[derive(Clone, Copy, Debug)] + enum Action { + Admit, + SucceedOldest, + FailOldest(Failure), + Fail(Failure), + } + + fn action() -> impl Strategy { + let failure = prop_oneof![Just(Failure::Hard), Just(Failure::Timeout)]; + let action = prop_oneof![ + Just(Action::Admit), + Just(Action::SucceedOldest), + failure.clone().prop_map(Action::FailOldest), + failure.prop_map(Action::Fail), + ]; + (0u64..30_000, action) + } + + proptest! { + #[test] + fn open_never_closes_directly_and_each_half_open_visit_has_one_probe( + actions in prop::collection::vec(action(), 1..200) + ) { + let transitions: Arc>>> = Arc::default(); + let seen = Arc::clone(&transitions); + let config = Config { + failure_threshold: 3, + recovery_timeout: Duration::from_secs(10), + timeout_min_duration: Duration::from_secs(2), + }; + let breaker = Breaker::new(config, Box::new(move |t| seen.lock().unwrap().push(t.clone()))); + let t0 = Instant::now(); + let mut now = t0; + let mut pending: Vec = Vec::new(); + let mut probes_per_visit: Vec = Vec::new(); + + for (advance_ms, action) in actions { + now += Duration::from_millis(advance_ms); + match action { + Action::Admit => { + if let Ok(admission) = breaker.admit(now) { + if admission.probe() { + probes_per_visit.push(1); + } + pending.push(admission); + } + } + Action::SucceedOldest if !pending.is_empty() => pending.remove(0).succeeded(), + Action::FailOldest(failure) if !pending.is_empty() => pending.remove(0).failed(failure, now), + Action::SucceedOldest | Action::FailOldest(_) => {} + Action::Fail(failure) => breaker.failed(failure, now), + } + } + + let transitions = transitions.lock().unwrap(); + let open_to_closed = transitions + .iter() + .filter(|t| matches!((&t.from, &t.to), (State::Open { .. }, State::Closed { .. }))) + .count(); + let half_open_visits = transitions.iter().filter(|t| t.to == State::HalfOpen).count(); + prop_assert_eq!(open_to_closed, 0); + prop_assert_eq!(probes_per_visit.len(), half_open_visits); + } + } +} diff --git a/litellm-rust/crates/circuit-breaker/src/fsm.rs b/litellm-rust/crates/circuit-breaker/src/fsm.rs new file mode 100644 index 00000000000..1beb0f890c8 --- /dev/null +++ b/litellm-rust/crates/circuit-breaker/src/fsm.rs @@ -0,0 +1,174 @@ +use std::{ + fmt::Debug, + mem, + sync::{Mutex, PoisonError}, +}; + +pub trait Machine { + type State: Clone + Debug; + type Event; + type Output; + + fn step(&self, state: &Self::State, event: Self::Event) -> (Self::State, Self::Output); +} + +#[derive(Clone, Debug)] +pub struct Transition { + pub from: S, + pub to: S, + pub generation: u64, +} + +pub type Observer = Box) + Send + Sync>; + +pub struct Shared { + machine: M, + current: Mutex<(M::State, u64)>, + observer: Observer, +} + +impl Shared { + pub fn new(machine: M, initial: M::State, observer: Observer) -> Self { + Self { + machine, + current: Mutex::new((initial, 0)), + observer, + } + } + + pub fn fire(&self, event: M::Event) -> (M::Output, u64) { + let mut current = self.current.lock().unwrap_or_else(PoisonError::into_inner); + self.apply(&mut current, event) + } + + pub fn fire_if(&self, generation: u64, event: M::Event) -> Option { + let mut current = self.current.lock().unwrap_or_else(PoisonError::into_inner); + (current.1 == generation).then(|| self.apply(&mut current, event).0) + } + + pub fn state(&self) -> M::State { + self.current + .lock() + .unwrap_or_else(PoisonError::into_inner) + .0 + .clone() + } + + fn apply(&self, current: &mut (M::State, u64), event: M::Event) -> (M::Output, u64) { + let (to, output) = self.machine.step(¤t.0, event); + if mem::discriminant(&to) != mem::discriminant(¤t.0) { + let from = mem::replace(&mut current.0, to.clone()); + current.1 += 1; + (self.observer)(&Transition { + from, + to, + generation: current.1, + }); + } else { + current.0 = to; + } + (output, current.1) + } +} + +#[cfg(test)] +mod tests { + use std::sync::{Arc, Mutex}; + + use rstest::rstest; + + use super::{Machine, Shared, Transition}; + + struct Wrapping(u32); + + #[derive(Clone, Debug, PartialEq, Eq)] + enum Counter { + Low(u32), + High(u32), + } + + enum Tick { + Up, + Stay, + } + + impl Machine for Wrapping { + type State = Counter; + type Event = Tick; + type Output = (); + + fn step(&self, state: &Counter, event: Tick) -> (Counter, ()) { + let value = match (state, event) { + (Counter::Low(_) | Counter::High(_), Tick::Stay) => return (state.clone(), ()), + (Counter::Low(value) | Counter::High(value), Tick::Up) => (value + 1) % self.0, + }; + match value < self.0 / 2 { + true => (Counter::Low(value), ()), + false => (Counter::High(value), ()), + } + } + } + + type Seen = Arc>>>; + + fn recorded(wrap: u32) -> (Shared, Seen) { + let transitions: Seen = Arc::default(); + let seen = Arc::clone(&transitions); + let shared = Shared::new( + Wrapping(wrap), + Counter::Low(0), + Box::new(move |transition| seen.lock().unwrap().push(transition.clone())), + ); + (shared, transitions) + } + + #[rstest] + fn concurrent_events_are_all_applied_and_transitions_are_observed_in_generation_order() { + let (shared, transitions) = recorded(4); + + std::thread::scope(|scope| { + for _ in 0..64 { + scope.spawn(|| { + for _ in 0..100 { + shared.fire(Tick::Up); + } + }); + } + }); + + assert_eq!(shared.state(), Counter::Low(6400 % 4)); + let generations: Vec = transitions + .lock() + .unwrap() + .iter() + .map(|t| t.generation) + .collect(); + assert_eq!(generations, (1..=3200).collect::>()); + } + + #[rstest] + fn a_step_that_keeps_the_variant_neither_counts_nor_notifies() { + let (shared, transitions) = recorded(4); + + let (_, after_inner_change) = shared.fire(Tick::Up); + let (_, after_self_loop) = shared.fire(Tick::Stay); + + assert_eq!(shared.state(), Counter::Low(1)); + assert_eq!((after_inner_change, after_self_loop), (0, 0)); + assert!(transitions.lock().unwrap().is_empty()); + } + + #[rstest] + fn fire_if_skips_an_event_from_an_older_generation() { + let (shared, _) = recorded(4); + shared.fire(Tick::Up); + let (_, generation) = shared.fire(Tick::Up); + + let stale = shared.fire_if(generation - 1, Tick::Up); + let current = shared.fire_if(generation, Tick::Up); + + assert!(stale.is_none()); + assert!(current.is_some()); + assert_eq!(shared.state(), Counter::High(3)); + } +} diff --git a/litellm-rust/crates/circuit-breaker/src/lib.rs b/litellm-rust/crates/circuit-breaker/src/lib.rs new file mode 100644 index 00000000000..e76228122d1 --- /dev/null +++ b/litellm-rust/crates/circuit-breaker/src/lib.rs @@ -0,0 +1,5 @@ +mod breaker; +mod fsm; + +pub use breaker::{Admission, Breaker, Config, Event, Failure, Output, Rejected, State}; +pub use fsm::{Machine, Observer, Shared, Transition}; diff --git a/litellm-rust/crates/config/Cargo.toml b/litellm-rust/crates/config/Cargo.toml index 23d53923362..9746a002bb9 100644 --- a/litellm-rust/crates/config/Cargo.toml +++ b/litellm-rust/crates/config/Cargo.toml @@ -9,11 +9,13 @@ repository.workspace = true litellm-auth-types.workspace = true litellm-router-types.workspace = true serde.workspace = true +serde_json.workspace = true serde_yaml_ng = "0.10.0" strum.workspace = true thiserror.workspace = true +tracing.workspace = true [dev-dependencies] rstest.workspace = true -serde_json.workspace = true tempfile.workspace = true +tracing-subscriber = { version = "0.3", default-features = false, features = ["registry", "std"] } diff --git a/litellm-rust/crates/config/src/error.rs b/litellm-rust/crates/config/src/error.rs index 61d19491abc..95c33967cf0 100644 --- a/litellm-rust/crates/config/src/error.rs +++ b/litellm-rust/crates/config/src/error.rs @@ -4,4 +4,11 @@ pub enum Error { Read(#[from] std::io::Error), #[error("invalid YAML config")] Parse(#[from] serde_yaml_ng::Error), + #[error("Redis host does not yet support configured setting {0}")] + RedisSetting(String), + #[error("invalid Redis setting {setting}: {reason}")] + RedisInvalid { + setting: &'static str, + reason: &'static str, + }, } diff --git a/litellm-rust/crates/config/src/lib.rs b/litellm-rust/crates/config/src/lib.rs index 8603ef3279c..b8af6f11d2e 100644 --- a/litellm-rust/crates/config/src/lib.rs +++ b/litellm-rust/crates/config/src/lib.rs @@ -2,6 +2,7 @@ mod error; mod includes; mod mcp; mod model; +mod redis; mod settings; mod value; @@ -13,6 +14,10 @@ pub use error::Error; pub use litellm_router_types::{LitellmParams, Spelled}; pub use mcp::{McpAuth, McpServer, McpTransport}; pub use model::Model; +pub use redis::{ + CacheSettings, Lookup, RedisCluster, RedisEndpoint, RedisHost, RedisNode, RedisSettings, + ResolvedRedis, environment_has_target, +}; pub use settings::{ ClickHouseStoreSettings, GeneralSettings, LiteLlmSettings, RouterSettings, TracingSettings, TracingStoreSettings, diff --git a/litellm-rust/crates/config/src/redis.rs b/litellm-rust/crates/config/src/redis.rs new file mode 100644 index 00000000000..5674f993025 --- /dev/null +++ b/litellm-rust/crates/config/src/redis.rs @@ -0,0 +1,667 @@ +use std::{collections::BTreeMap, fmt, time::Duration}; + +use litellm_auth_types::SecretValue; +use litellm_router_types::Spelled; +use serde::{Deserialize, Deserializer, de::Error as _}; + +use crate::{AdditionalFields, Error, Object, Value}; + +pub type Lookup<'a> = &'a dyn Fn(&str) -> Option; + +const SUPPORTED: &[&str] = &[ + "url", + "host", + "port", + "username", + "password", + "ssl", + "db", + "socket_timeout", + "startup_nodes", + "namespace", +]; + +const IGNORED: &[&str] = &[ + "max_connections", + "single_connection_client", + "health_check_interval", + "socket_keepalive", + "socket_keepalive_options", + "decode_responses", + "encoding", + "encoding_errors", + "charset", + "errors", + "client_name", + "lib_name", + "lib_version", + "sentinel_password", + "gcp_ssl_ca_certs", + "azure_client_id", + "azure_tenant_id", + "azure_client_secret", + "aws_iam_user_name", + "aws_iam_cache_name", + "aws_iam_region", + "aws_iam_serverless", +]; + +const WARN_IGNORED: &[&str] = &[ + "max_connections", + "health_check_interval", + "socket_keepalive", + "socket_keepalive_options", +]; + +const REDIS_PY_KWARGS: &[&str] = &[ + "aws_iam_auth", + "aws_iam_cache_name", + "aws_iam_region", + "aws_iam_serverless", + "aws_iam_user_name", + "azure_client_id", + "azure_client_secret", + "azure_redis_ad_token", + "azure_tenant_id", + "cache", + "cache_config", + "charset", + "client_name", + "db", + "decode_responses", + "encoding", + "encoding_errors", + "errors", + "event_dispatcher", + "gcp_service_account", + "gcp_ssl_ca_certs", + "health_check_interval", + "host", + "lib_name", + "lib_version", + "max_connections", + "password", + "port", + "protocol", + "redis_connect_func", + "retry_on_error", + "retry_on_timeout", + "single_connection_client", + "socket_connect_timeout", + "socket_keepalive", + "socket_keepalive_options", + "socket_timeout", + "ssl", + "ssl_ca_certs", + "ssl_ca_data", + "ssl_ca_path", + "ssl_cert_reqs", + "ssl_certfile", + "ssl_check_hostname", + "ssl_ciphers", + "ssl_keyfile", + "ssl_min_version", + "ssl_ocsp_context", + "ssl_ocsp_expected_cert", + "ssl_password", + "ssl_validate_ocsp", + "ssl_validate_ocsp_stapled", + "unix_socket_path", + "url", + "username", +]; + +const UNSUPPORTED_EXTRA: &[&str] = &[ + "credential_provider", + "sentinel_nodes", + "retry", + "connection_pool", +]; + +const CACHE_ONLY: &[&str] = &[ + "type", + "mode", + "ttl", + "default_in_memory_ttl", + "default_in_redis_ttl", + "supported_call_types", + "similarity_threshold", + "max_messages", + "_backend", +]; + +#[derive(Clone, Default)] +pub struct RedisSettings { + pub url: Option, + pub host: Option>, + pub port: Option>, + pub username: Option, + pub password: Option, + pub ssl: Option>, + pub db: Option>, + pub socket_timeout: Option>, + pub startup_nodes: Option, + pub namespace: Option>, + pub additional_fields: AdditionalFields, + unparsed: DeferredInputs, +} + +impl<'de> Deserialize<'de> for RedisSettings { + fn deserialize>(deserializer: D) -> Result { + let fields = BTreeMap::::deserialize(deserializer)?; + let scalar = |key: &str| fields.get(key).and_then(as_text); + Ok(Self { + url: scalar("url").map(SecretValue::new), + host: scalar("host").map(Spelled::Value), + port: spelled(scalar("port")), + username: scalar("username").map(SecretValue::new), + password: scalar("password").map(SecretValue::new), + ssl: spelled(scalar("ssl")), + db: spelled(scalar("db")), + socket_timeout: spelled(scalar("socket_timeout")), + startup_nodes: fields + .get("startup_nodes") + .filter(|value| !value.is_null()) + .cloned(), + namespace: scalar("namespace").map(Spelled::Value), + additional_fields: cloned(&fields, |key, _| !SUPPORTED.contains(&key)), + unparsed: deferred_inputs(&fields), + }) + } +} + +impl fmt::Debug for RedisSettings { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("RedisSettings") + .field("url", &self.url) + .field("host", &self.host) + .field("port", &self.port) + .field("username", &self.username) + .field("password", &self.password) + .field("ssl", &self.ssl) + .field("db", &self.db) + .field("socket_timeout", &self.socket_timeout) + .field( + "startup_nodes", + &self.startup_nodes.as_ref().map(|_| ""), + ) + .field("namespace", &self.namespace) + .field("additional_fields", &self.additional_fields.keys()) + .field("unparsed", &self.unparsed.keys()) + .finish() + } +} + +#[derive(Clone, Debug)] +pub enum CacheSettings { + Redis(Box), + Other(Object), +} + +impl<'de> Deserialize<'de> for CacheSettings { + fn deserialize>(deserializer: D) -> Result { + let fields = BTreeMap::::deserialize(deserializer)?; + let is_redis = matches!(fields.get("type"), Some(Value::String(kind)) if kind == "redis"); + if !is_redis { + return Ok(Self::Other(Object::from_fields(fields))); + } + let mapping = fields + .into_iter() + .map(|(key, value)| (Value::String(key), value)) + .collect(); + serde_yaml_ng::from_value(Value::Mapping(mapping)) + .map(|settings| Self::Redis(Box::new(settings))) + .map_err(D::Error::custom) + } +} + +#[derive(Clone, PartialEq)] +pub struct ResolvedRedis { + pub endpoint: RedisEndpoint, + pub namespace: Option, + pub timeout: Option, +} + +impl fmt::Debug for ResolvedRedis { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("ResolvedRedis") + .field("endpoint", &self.endpoint) + .field("namespace", &self.namespace) + .field("timeout", &self.timeout) + .finish() + } +} + +#[derive(Clone, PartialEq)] +pub enum RedisEndpoint { + Url(SecretValue), + Host(RedisHost), + Cluster(RedisCluster), +} + +impl fmt::Debug for RedisEndpoint { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Url(_) => formatter.debug_tuple("Url").field(&"").finish(), + Self::Host(host) => host.fmt(formatter), + Self::Cluster(cluster) => cluster.fmt(formatter), + } + } +} + +#[derive(Clone, Debug, PartialEq)] +pub struct RedisHost { + pub host: String, + pub port: u16, + pub username: Option, + pub password: Option, + pub tls: bool, + pub db: i64, +} + +#[derive(Clone, Debug, PartialEq)] +pub struct RedisCluster { + pub nodes: Vec, + pub username: Option, + pub password: Option, + pub tls: bool, +} + +#[derive(Clone, Debug, PartialEq)] +pub struct RedisNode { + pub host: String, + pub port: u16, +} + +#[derive(Clone)] +struct Entry { + value: Value, + source: String, +} + +type Kwargs = BTreeMap; + +impl RedisSettings { + pub fn resolve(&self, lookup: Lookup<'_>) -> Result { + self.resolve_with(lookup, false) + } + + pub fn resolve_cache(&self, lookup: Lookup<'_>) -> Result { + self.resolve_with(lookup, true) + } + + fn resolve_with(&self, lookup: Lookup<'_>, cache: bool) -> Result { + let block = self.kwargs(cache); + let block: Kwargs = block + .into_iter() + .filter_map(|(key, entry)| { + let value = environment_reference(&entry.value, lookup)?; + Some((key, Entry { value, ..entry })) + }) + .collect(); + let environment = environment_kwargs(lookup); + let names_a_target = ["host", "startup_nodes", "sentinel_nodes"] + .iter() + .any(|key| block.contains_key(*key)); + let drop_environment_url = names_a_target && !block.contains_key("url"); + let merged: Kwargs = environment + .into_iter() + .filter(|(key, _)| !(key == "url" && drop_environment_url)) + .chain(block) + .collect(); + reject_unsupported(&merged, lookup)?; + warn_ignored(&merged); + + let startup_nodes = match merged.get("startup_nodes") { + Some(entry) => Some(cluster_nodes(&entry.value)?), + None => lookup("REDIS_CLUSTER_NODES") + .map(|nodes| cluster_nodes(&Value::String(nodes))) + .transpose()?, + }; + let username = secret_value(&merged, "username"); + let password = secret_value(&merged, "password"); + let tls = bool_value(&merged, "ssl"); + let endpoint = match (startup_nodes, text_value(&merged, "url")) { + (Some(nodes), _) if !nodes.is_empty() => RedisEndpoint::Cluster(RedisCluster { + nodes, + username, + password, + tls, + }), + (_, Some(url)) => RedisEndpoint::Url(SecretValue::new(url)), + _ => RedisEndpoint::Host(RedisHost { + host: text_value(&merged, "host").ok_or(Error::RedisInvalid { + setting: "host", + reason: "either host or url must be specified for redis", + })?, + port: integer(&merged, "port")?.unwrap_or(6379), + username, + password, + tls, + db: integer(&merged, "db")?.unwrap_or(0), + }), + }; + let timeout = merged + .get("socket_timeout") + .and_then(|entry| as_text(&entry.value)) + .map(|seconds| { + seconds + .parse::() + .ok() + .and_then(|seconds| Duration::try_from_secs_f64(seconds).ok()) + .ok_or(Error::RedisInvalid { + setting: "socket_timeout", + reason: "must be a non-negative number of seconds", + }) + }) + .transpose()?; + Ok(ResolvedRedis { + endpoint, + namespace: text_value(&merged, "namespace"), + timeout, + }) + } + + fn kwargs(&self, cache: bool) -> Kwargs { + let known = [ + ( + "url", + self.url.as_ref().map(|value| value.expose().to_owned()), + ), + ("host", spelled_text(self.host.as_ref())), + ("port", spelled_text(self.port.as_ref())), + ( + "username", + self.username + .as_ref() + .map(|value| value.expose().to_owned()), + ), + ( + "password", + self.password + .as_ref() + .map(|value| value.expose().to_owned()), + ), + ("ssl", spelled_text(self.ssl.as_ref())), + ("db", spelled_text(self.db.as_ref())), + ("socket_timeout", spelled_text(self.socket_timeout.as_ref())), + ("namespace", spelled_text(self.namespace.as_ref())), + ]; + let known = known.into_iter().filter_map(|(key, value)| { + value.map(|value| { + ( + key.to_owned(), + Entry { + value: Value::String(value), + source: key.to_owned(), + }, + ) + }) + }); + let startup_nodes = self.startup_nodes.clone().map(|value| { + ( + "startup_nodes".to_owned(), + Entry { + value, + source: "startup_nodes".to_owned(), + }, + ) + }); + let extras = self.additional_fields.iter().filter_map(|(key, value)| { + if cache && CACHE_ONLY.contains(&key.as_str()) { + return None; + } + let key = match (cache, key.as_str()) { + (true, "redis_startup_nodes") => "startup_nodes".to_owned(), + _ => key.clone(), + }; + Some(( + key.clone(), + Entry { + value: value.clone(), + source: key, + }, + )) + }); + let unparsed = self.unparsed.iter().map(|(key, value)| { + ( + key.clone(), + Entry { + value: value.0.clone(), + source: key.clone(), + }, + ) + }); + extras + .chain(unparsed) + .chain(known) + .chain(startup_nodes) + .collect() + } +} + +pub fn environment_has_target(lookup: Lookup<'_>) -> bool { + [ + "REDIS_HOST", + "REDIS_URL", + "REDIS_CLUSTER_NODES", + "REDIS_SENTINEL_NODES", + ] + .iter() + .any(|name| lookup(name).is_some()) +} + +fn reject_unsupported(merged: &Kwargs, lookup: Lookup<'_>) -> Result<(), Error> { + let active_auth = [ + merged + .get("gcp_service_account") + .map(|entry| entry.source.as_str()) + .or_else(|| lookup("REDIS_GCP_SERVICE_ACCOUNT").map(|_| "REDIS_GCP_SERVICE_ACCOUNT")), + text_value(merged, "azure_redis_ad_token") + .or_else(|| lookup("REDIS_AZURE_AD_TOKEN")) + .filter(|value| value.eq_ignore_ascii_case("true")) + .map(|_| { + merged + .get("azure_redis_ad_token") + .map_or("REDIS_AZURE_AD_TOKEN", |entry| entry.source.as_str()) + }), + bool_value(merged, "aws_iam_auth").then(|| merged["aws_iam_auth"].source.as_str()), + merged + .get("sentinel_nodes") + .map(|entry| entry.source.as_str()) + .or_else(|| lookup("REDIS_SENTINEL_NODES").map(|_| "REDIS_SENTINEL_NODES")), + ]; + if let Some(setting) = active_auth.into_iter().flatten().next() { + return Err(Error::RedisSetting(setting.to_owned())); + } + match merged + .iter() + .find(|(key, _)| is_unsupported(key)) + .map(|(_, entry)| entry.source.clone()) + { + Some(setting) => Err(Error::RedisSetting(setting)), + None => Ok(()), + } +} + +fn warn_ignored(merged: &Kwargs) { + for setting in WARN_IGNORED { + if merged.contains_key(*setting) { + tracing::warn!(setting, "Redis setting is accepted but ignored"); + } + } +} + +fn is_unsupported(key: &str) -> bool { + let known = REDIS_PY_KWARGS.contains(&key) || UNSUPPORTED_EXTRA.contains(&key); + let handled = SUPPORTED.contains(&key) + || IGNORED.contains(&key) + || matches!( + key, + "gcp_service_account" | "azure_redis_ad_token" | "aws_iam_auth" + ); + known && !handled +} + +fn environment_kwargs(lookup: Lookup<'_>) -> Kwargs { + REDIS_PY_KWARGS + .iter() + .filter_map(|key| { + let source = format!("REDIS_{}", key.to_ascii_uppercase()); + let value = lookup(&source)?; + Some(( + (*key).to_owned(), + Entry { + value: Value::String(value), + source, + }, + )) + }) + .collect() +} + +fn environment_reference(value: &Value, lookup: Lookup<'_>) -> Option { + match value { + Value::String(text) => match text.strip_prefix("os.environ/") { + Some(name) => lookup(name).map(Value::String), + None => Some(value.clone()), + }, + Value::Null => None, + _ => Some(value.clone()), + } +} + +fn as_text(value: &Value) -> Option { + match value { + Value::String(text) => Some(text.clone()), + Value::Number(number) => Some(number.to_string()), + Value::Bool(flag) => Some(flag.to_string()), + _ => None, + } +} + +fn text_value(kwargs: &Kwargs, key: &str) -> Option { + kwargs.get(key).and_then(|entry| as_text(&entry.value)) +} + +fn secret_value(kwargs: &Kwargs, key: &str) -> Option { + text_value(kwargs, key).map(SecretValue::new) +} + +fn bool_value(kwargs: &Kwargs, key: &str) -> bool { + match kwargs.get(key).map(|entry| &entry.value) { + Some(Value::Bool(flag)) => *flag, + Some(value) => as_text(value) + .is_some_and(|text| matches!(text.to_ascii_lowercase().as_str(), "true" | "1" | "yes")), + None => false, + } +} + +fn integer(kwargs: &Kwargs, key: &'static str) -> Result, Error> { + kwargs + .get(key) + .and_then(|entry| as_text(&entry.value)) + .map(|text| { + text.parse().map_err(|_| Error::RedisInvalid { + setting: key, + reason: "must be an integer", + }) + }) + .transpose() +} + +const INVALID_NODES: Error = Error::RedisInvalid { + setting: "startup_nodes", + reason: "must be a list of {host, port} objects", +}; + +fn cluster_nodes(value: &Value) -> Result, Error> { + match value { + Value::String(json) => { + let nodes: Vec> = + serde_json::from_str(json).map_err(|_| INVALID_NODES)?; + nodes + .iter() + .map(|node| { + cluster_node( + node.get("host").and_then(json_text), + node.get("port").and_then(json_text), + ) + }) + .collect() + } + Value::Sequence(nodes) => nodes + .iter() + .map(|node| { + cluster_node( + node.get("host").and_then(as_text), + node.get("port").and_then(as_text), + ) + }) + .collect(), + _ => Err(INVALID_NODES), + } +} + +fn cluster_node(host: Option, port: Option) -> Result { + match (host, port.and_then(|port| port.parse().ok())) { + (Some(host), Some(port)) => Ok(RedisNode { host, port }), + _ => Err(INVALID_NODES), + } +} + +fn json_text(value: &serde_json::Value) -> Option { + match value { + serde_json::Value::String(text) => Some(text.clone()), + serde_json::Value::Number(number) => Some(number.to_string()), + _ => None, + } +} + +fn spelled_text(value: Option<&Spelled>) -> Option { + value.map(|value| match value { + Spelled::Value(value) => value.to_string(), + Spelled::Text(value) => value.clone(), + }) +} + +fn spelled(text: Option) -> Option> { + text.map(|text| match text.parse() { + Ok(value) => Spelled::Value(value), + Err(_) => Spelled::Text(text), + }) +} + +fn cloned( + fields: &BTreeMap, + keep: impl Fn(&str, &Value) -> bool, +) -> AdditionalFields { + fields + .iter() + .filter(|(key, value)| keep(key, value)) + .map(|(key, value)| (key.clone(), value.clone())) + .collect() +} + +#[derive(Clone)] +struct DeferredInput(Value); + +type DeferredInputs = BTreeMap; + +fn deferred_inputs(fields: &BTreeMap) -> DeferredInputs { + fields + .iter() + .filter(|(key, value)| is_unparsed(key, value)) + .map(|(key, value)| (key.clone(), DeferredInput(value.clone()))) + .collect() +} + +fn is_unparsed(key: &str, value: &Value) -> bool { + SUPPORTED.contains(&key) + && match key { + "startup_nodes" => value.is_null(), + _ => as_text(value).is_none(), + } +} diff --git a/litellm-rust/crates/config/src/settings.rs b/litellm-rust/crates/config/src/settings.rs index 1094de1638e..9571c4fa5de 100644 --- a/litellm-rust/crates/config/src/settings.rs +++ b/litellm-rust/crates/config/src/settings.rs @@ -3,7 +3,9 @@ use std::fmt; use litellm_auth_types::SecretValue; use serde::Deserialize; -use crate::{AdditionalFields, Object, Spelled, Value, value::one_or_many}; +use crate::{ + AdditionalFields, CacheSettings, Object, RedisSettings, Spelled, Value, value::one_or_many, +}; #[derive(Clone, Debug, Deserialize)] #[serde(rename_all = "lowercase")] @@ -76,7 +78,7 @@ pub struct GeneralSettings { pub enable_public_model_hub: bool, pub dangerously_permit_weak_or_unset_master_key: Option, pub plugins: Option>, - pub coordination_redis: Option, + pub coordination_redis: Option, pub mcp_allowed_hosts: Option>, pub mcp_allowed_origins: Box<[String]>, #[serde(flatten)] @@ -220,7 +222,7 @@ pub struct LiteLlmSettings { pub request_timeout: Option>, pub num_retries: Option, pub cache: Option, - pub cache_params: Option, + pub cache_params: Option, #[serde(deserialize_with = "one_or_many")] pub callbacks: Option>, #[serde(deserialize_with = "one_or_many")] diff --git a/litellm-rust/crates/config/src/value.rs b/litellm-rust/crates/config/src/value.rs index 5becdc87b70..667a5afa130 100644 --- a/litellm-rust/crates/config/src/value.rs +++ b/litellm-rust/crates/config/src/value.rs @@ -10,6 +10,10 @@ pub type AdditionalFields = BTreeMap; pub struct Object(BTreeMap); impl Object { + pub(crate) fn from_fields(fields: BTreeMap) -> Self { + Self(fields) + } + pub fn get(&self, key: &str) -> Option<&Value> { self.0.get(key) } diff --git a/litellm-rust/crates/config/tests/redis.rs b/litellm-rust/crates/config/tests/redis.rs new file mode 100644 index 00000000000..ab409541c8e --- /dev/null +++ b/litellm-rust/crates/config/tests/redis.rs @@ -0,0 +1,405 @@ +use std::{ + collections::HashMap, + fmt::Write, + sync::{Arc, Mutex}, +}; + +use litellm_config::{CacheSettings, Config, Error, RedisEndpoint, RedisSettings, ResolvedRedis}; +use rstest::rstest; +use tracing::{Event, Subscriber, field::Visit}; +use tracing_subscriber::{Layer, layer::Context, prelude::*}; + +fn resolve( + yaml: &str, + environment: &[(&str, &str)], +) -> Result { + let environment: HashMap = environment + .iter() + .map(|(key, value)| ((*key).to_owned(), (*value).to_owned())) + .collect(); + let config = Config::from_yaml(yaml).unwrap(); + config + .general_settings + .coordination_redis + .as_ref() + .unwrap_or(&RedisSettings::default()) + .resolve(&|name| environment.get(name).cloned()) +} + +fn resolve_cache(yaml: &str) -> Result { + let config = Config::from_yaml(yaml).unwrap(); + match config.litellm_settings.cache_params.unwrap() { + CacheSettings::Redis(settings) => settings.resolve_cache(&|_| None), + CacheSettings::Other(_) => panic!("expected a Redis cache"), + } +} + +fn cluster_ports(resolved: &ResolvedRedis) -> Vec { + match &resolved.endpoint { + RedisEndpoint::Cluster(cluster) => cluster.nodes.iter().map(|node| node.port).collect(), + endpoint => panic!("expected a cluster, got {endpoint:?}"), + } +} + +#[rstest] +#[case::host( + "general_settings:\n coordination_redis:\n host: redis\n port: 6380\n db: 2\n ssl: true\n namespace: app\n socket_timeout: 1.5", + &[], + "host:redis:6380:2:true:app:1500" +)] +#[case::environment_url( + "model_list: []", + &[("REDIS_URL", "redis://environment:6379")], + "url:redis://environment:6379" +)] +#[case::explicit_host_drops_environment_url( + "general_settings:\n coordination_redis:\n host: configured", + &[("REDIS_URL", "redis://environment:6379")], + "host:configured:6379:0:false::none" +)] +#[case::yaml_cluster( + "general_settings:\n coordination_redis:\n startup_nodes:\n - host: node-a\n port: 7000\n - host: node-b\n port: \"7001\"", + &[], + "cluster:node-a:7000,node-b:7001" +)] +#[case::json_environment_cluster( + "model_list: []", + &[("REDIS_CLUSTER_NODES", r#"[{"host":"node-a","port":"7000"},{"host":"node-b","port":7001}]"#)], + "cluster:node-a:7000,node-b:7001" +)] +#[case::missing_reference_keeps_environment_fallback( + "general_settings:\n coordination_redis:\n host: os.environ/MISSING_HOST", + &[("REDIS_HOST", "environment")], + "host:environment:6379:0:false::none" +)] +#[case::configured_values_beat_the_same_environment_values( + "general_settings:\n coordination_redis:\n host: configured\n port: 6380\n db: 3", + &[("REDIS_HOST", "environment"), ("REDIS_PORT", "6390"), ("REDIS_DB", "9")], + "host:configured:6380:3:false::none" +)] +#[case::environment_url_beats_the_environment_host( + "model_list: []", + &[("REDIS_URL", "redis://environment:6379"), ("REDIS_HOST", "ignored")], + "url:redis://environment:6379" +)] +#[case::startup_nodes_beat_a_url( + "general_settings:\n coordination_redis:\n url: redis://single:6379\n startup_nodes:\n - host: node-a\n port: 7000", + &[], + "cluster:node-a:7000" +)] +#[case::environment_references_resolve_the_host_and_port( + "general_settings:\n coordination_redis:\n host: os.environ/COORD_HOST\n port: os.environ/COORD_PORT", + &[("COORD_HOST", "referenced"), ("COORD_PORT", "6400")], + "host:referenced:6400:0:false::none" +)] +#[case::a_null_host_falls_back_to_the_environment( + "general_settings:\n coordination_redis:\n host: null", + &[("REDIS_HOST", "environment")], + "host:environment:6379:0:false::none" +)] +fn resolves_redis_through_the_parsed_config( + #[case] yaml: &str, + #[case] environment: &[(&str, &str)], + #[case] expected: &str, +) { + let resolved = resolve(yaml, environment).unwrap(); + let description = match resolved.endpoint { + RedisEndpoint::Url(url) => format!("url:{}", url.expose()), + RedisEndpoint::Host(host) => format!( + "host:{}:{}:{}:{}:{}:{}", + host.host, + host.port, + host.db, + host.tls, + resolved.namespace.unwrap_or_default(), + resolved + .timeout + .map(|timeout| timeout.as_millis().to_string()) + .unwrap_or_else(|| "none".to_owned()) + ), + RedisEndpoint::Cluster(cluster) => format!( + "cluster:{}", + cluster + .nodes + .iter() + .map(|node| format!("{}:{}", node.host, node.port)) + .collect::>() + .join(",") + ), + }; + + assert_eq!(description, expected); +} + +#[rstest] +#[case::custom_tls_file( + "general_settings:\n coordination_redis:\n host: redis\n ssl_ca_certs: /ca.pem", + &[], + "ssl_ca_certs" +)] +#[case::gcp_auth_from_config( + "general_settings:\n coordination_redis:\n host: redis\n gcp_service_account: account", + &[], + "gcp_service_account" +)] +#[case::gcp_auth_from_environment( + "model_list: []", + &[("REDIS_HOST", "redis"), ("REDIS_GCP_SERVICE_ACCOUNT", "account")], + "REDIS_GCP_SERVICE_ACCOUNT" +)] +#[case::azure_auth_from_config( + "general_settings:\n coordination_redis:\n host: redis\n azure_redis_ad_token: true", + &[], + "azure_redis_ad_token" +)] +#[case::azure_auth_from_environment( + "model_list: []", + &[("REDIS_HOST", "redis"), ("REDIS_AZURE_AD_TOKEN", "true")], + "REDIS_AZURE_AD_TOKEN" +)] +#[case::aws_auth_from_config( + "general_settings:\n coordination_redis:\n host: redis\n aws_iam_auth: true", + &[], + "aws_iam_auth" +)] +#[case::aws_auth_from_environment( + "model_list: []", + &[("REDIS_HOST", "redis"), ("REDIS_AWS_IAM_AUTH", "true")], + "REDIS_AWS_IAM_AUTH" +)] +#[case::sentinel_from_config( + "general_settings:\n coordination_redis:\n sentinel_nodes: [[redis, 26379]]", + &[], + "sentinel_nodes" +)] +#[case::sentinel_from_environment( + "model_list: []", + &[("REDIS_SENTINEL_NODES", "[[\"redis\",26379]]")], + "REDIS_SENTINEL_NODES" +)] +fn rejected_settings_name_their_source( + #[case] yaml: &str, + #[case] environment: &[(&str, &str)], + #[case] expected: &str, +) { + let error = resolve(yaml, environment).unwrap_err(); + + assert!(matches!(error, Error::RedisSetting(setting) if setting == expected)); +} + +#[rstest] +fn an_explicit_host_keeps_environment_credentials() { + let resolved = resolve( + "general_settings:\n coordination_redis:\n host: configured", + &[ + ("REDIS_PASSWORD", "environment-secret"), + ("REDIS_SSL", "True"), + ], + ) + .unwrap(); + + let RedisEndpoint::Host(host) = resolved.endpoint else { + panic!("expected a host endpoint"); + }; + assert_eq!( + host.password.as_ref().map(|password| password.expose()), + Some("environment-secret") + ); + assert!(host.tls); +} + +#[rstest] +#[case::only_a_port("general_settings:\n coordination_redis:\n port: 6379", &[])] +#[case::a_structured_host_with_an_environment_target( + "general_settings:\n coordination_redis:\n host: []", + &[("REDIS_HOST", "environment"), ("REDIS_URL", "redis://environment:6379")] +)] +fn a_block_without_a_usable_target_fails(#[case] yaml: &str, #[case] environment: &[(&str, &str)]) { + let error = resolve(yaml, environment).unwrap_err(); + + assert!(matches!( + error, + Error::RedisInvalid { + setting: "host", + .. + } + )); +} + +#[rstest] +fn the_cache_alias_names_cluster_nodes() { + let resolved = resolve_cache( + "litellm_settings:\n cache_params:\n type: redis\n redis_startup_nodes:\n - host: node-a\n port: 7000", + ) + .unwrap(); + + assert_eq!(cluster_ports(&resolved), vec![7000]); +} + +#[rstest] +fn canonical_startup_nodes_beat_the_cache_alias() { + let resolved = resolve_cache( + "litellm_settings:\n cache_params:\n type: redis\n startup_nodes:\n - host: canonical\n port: 7000\n redis_startup_nodes:\n - host: alias\n port: 7100\n - host: alias\n port: 7101", + ) + .unwrap(); + + assert_eq!(cluster_ports(&resolved), vec![7000]); +} + +#[rstest] +fn a_null_canonical_startup_nodes_setting_shadows_the_cache_alias() { + let resolved = resolve_cache( + "litellm_settings:\n cache_params:\n type: redis\n host: cache\n startup_nodes: null\n redis_startup_nodes:\n - host: alias\n port: 7100", + ) + .unwrap(); + + assert!(matches!(resolved.endpoint, RedisEndpoint::Host(_))); +} + +#[rstest] +fn a_missing_canonical_startup_nodes_reference_shadows_the_cache_alias() { + let resolved = resolve_cache( + "litellm_settings:\n cache_params:\n type: redis\n host: cache\n startup_nodes: os.environ/MISSING_NODES\n redis_startup_nodes:\n - host: alias\n port: 7100", + ) + .unwrap(); + + assert!(matches!(resolved.endpoint, RedisEndpoint::Host(_))); +} + +#[rstest] +fn rejected_environment_settings_never_print_the_value() { + let error = resolve( + "model_list: []", + &[ + ("REDIS_HOST", "redis"), + ("REDIS_SSL_CERT_REQS", "private-option-value"), + ], + ) + .unwrap_err(); + let message = error.to_string(); + + assert!(message.contains("REDIS_SSL_CERT_REQS")); + assert!(!message.contains("private-option-value")); +} + +#[rstest] +#[case::negative("-1")] +#[case::not_a_number("slow")] +fn rejects_bad_socket_timeout(#[case] timeout: &str) { + let yaml = format!( + "general_settings:\n coordination_redis:\n host: redis\n socket_timeout: {timeout}" + ); + let error = resolve(&yaml, &[]).unwrap_err(); + + assert!(matches!( + error, + Error::RedisInvalid { + setting: "socket_timeout", + .. + } + )); +} + +#[rstest] +fn defaults_socket_timeout_for_the_gateway() { + let resolved = resolve( + "general_settings:\n coordination_redis:\n host: redis", + &[], + ) + .unwrap(); + + assert_eq!(resolved.timeout, None); +} + +#[rstest] +fn debug_redacts_redis_secrets_and_url() { + let config = Config::from_yaml( + "general_settings:\n coordination_redis:\n url: redis://user:password@redis:6379\n username: credential-user\n password: credential-password", + ) + .unwrap(); + let settings = config.general_settings.coordination_redis.unwrap(); + let resolved = settings.resolve(&|_| None).unwrap(); + let debug = format!("{settings:?} {resolved:?}"); + + assert!(!debug.contains("redis://user:password")); + assert!(!debug.contains("credential-user")); + assert!(!debug.contains("credential-password")); + assert!(debug.contains("redacted") || debug.contains("REDACTED")); +} + +#[rstest] +fn inactive_redis_cache_keeps_malformed_known_inputs_deferred() { + let config = Config::from_yaml( + "litellm_settings:\n cache: false\n cache_params:\n type: redis\n host: []\n port: invalid\n socket_timeout: invalid", + ) + .unwrap(); + + assert!(matches!( + config.litellm_settings.cache_params, + Some(litellm_config::CacheSettings::Redis(_)) + )); +} + +#[rstest] +fn non_redis_cache_settings_remain_opaque() { + let config = Config::from_yaml( + "litellm_settings:\n cache_params:\n type: local\n custom_option:\n nested: true", + ) + .unwrap(); + + assert!(matches!( + config.litellm_settings.cache_params, + Some(litellm_config::CacheSettings::Other(_)) + )); +} + +#[rstest] +fn legacy_unknown_settings_do_not_change_resolution() { + let resolved = resolve( + "general_settings:\n coordination_redis:\n host: redis\n redis_flush_size: 42\n service_name: primary", + &[], + ) + .unwrap(); + + assert!(matches!(resolved.endpoint, RedisEndpoint::Host(_))); +} + +struct WarningCapture(Arc>>); + +impl Layer for WarningCapture { + fn on_event(&self, event: &Event<'_>, _context: Context<'_, S>) { + let mut fields = Fields::default(); + event.record(&mut fields); + self.0.lock().unwrap().push(fields.0); + } +} + +#[derive(Default)] +struct Fields(String); + +impl Visit for Fields { + fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) { + write!(self.0, "{}={value:?} ", field.name()).unwrap(); + } + + fn record_str(&mut self, field: &tracing::field::Field, value: &str) { + write!(self.0, "{}={value} ", field.name()).unwrap(); + } +} + +#[rstest] +fn ignored_setting_warning_names_the_setting_without_its_value() { + let events = Arc::new(Mutex::new(Vec::new())); + let subscriber = tracing_subscriber::registry().with(WarningCapture(events.clone())); + tracing::subscriber::with_default(subscriber, || { + resolve( + "general_settings:\n coordination_redis:\n host: redis\n max_connections: 47", + &[], + ) + .unwrap(); + }); + let warnings = events.lock().unwrap().join(" "); + + assert!(warnings.contains("max_connections")); + assert!(!warnings.contains("47")); +} diff --git a/litellm-rust/crates/db-testing/Cargo.toml b/litellm-rust/crates/db-testing/Cargo.toml index 3ba71fbf644..457c0f2ad87 100644 --- a/litellm-rust/crates/db-testing/Cargo.toml +++ b/litellm-rust/crates/db-testing/Cargo.toml @@ -7,7 +7,7 @@ repository.workspace = true [dependencies] sqlx = { workspace = true, features = ["migrate"] } -testcontainers-modules = { version = "0.15.0", features = ["postgres"] } +testcontainers-modules = { workspace = true, features = ["postgres"] } thiserror.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/gateway/AGENTS.md b/litellm-rust/crates/gateway/AGENTS.md index d54fa062f2b..a33bd19cc80 100644 --- a/litellm-rust/crates/gateway/AGENTS.md +++ b/litellm-rust/crates/gateway/AGENTS.md @@ -3,3 +3,10 @@ - Endpoint paths, request handling, model resolution, and response encoding belong to the mounted crates. Provider execution belongs to the `inference-` crates and `llms` - Inject shared state and infrastructure. Avoid global runtimes, duplicate client pools, and abstractions for hypothetical endpoint groups - Test mounting and server lifecycle through public HTTP behavior. Test endpoint semantics in the owning crate + +## Redis +- `build_redis` builds the Redis clients once at boot (`src/redis.rs`). Components get a clone as a field, never a per-request extractor or a pool. redis-rs multiplexes one connection per node, so there is no connection pool and pool settings like `max_connections` are accepted and ignored +- Which Redis backs which role mirrors the proxy: coordination is `general_settings.coordination_redis`, else the `cache_params` Redis when `cache: true` and `type: redis`, else `REDIS_*` env. The transaction buffer reuses coordination, else env when `use_redis_transaction_buffer` is on. The response cache gets its own client so large values never queue in front of coordination commands +- `litellm-config` follows `litellm/_redis.py`'s merge of config blocks and `REDIS_*` env. A redis-py setting we can't honor (sentinel, IAM, custom TLS files, ...) fails boot with its name, never silently dropped +- Boot never waits for a configured Redis (clients connect lazily). Only the env-only coordination fallback is pinged, for 2s, and treated as absent when unreachable, as in Python +- `/health/readiness` never touches Redis, matching the proxy diff --git a/litellm-rust/crates/gateway/Cargo.toml b/litellm-rust/crates/gateway/Cargo.toml index 387bcd43231..b2a2f09bfc7 100644 --- a/litellm-rust/crates/gateway/Cargo.toml +++ b/litellm-rust/crates/gateway/Cargo.toml @@ -19,6 +19,8 @@ litellm-gateway-auth.workspace = true litellm-gateway-ui.workspace = true litellm-auth-types.workspace = true litellm-config.workspace = true +litellm-circuit-breaker.workspace = true +litellm-redis-client.workspace = true litellm-http = { workspace = true, features = ["mcp"] } litellm-llms.workspace = true litellm-secrets.workspace = true @@ -31,6 +33,7 @@ uuid.workspace = true tower-sessions-moka-store = "0.15.0" [dev-dependencies] +litellm-redis-testing.workspace = true base64.workspace = true rstest.workspace = true tempfile.workspace = true diff --git a/litellm-rust/crates/gateway/src/error.rs b/litellm-rust/crates/gateway/src/error.rs index 75dc04d179d..170a8284c18 100644 --- a/litellm-rust/crates/gateway/src/error.rs +++ b/litellm-rust/crates/gateway/src/error.rs @@ -8,6 +8,15 @@ pub enum Error { McpSetting(String), #[error(transparent)] Mcp(#[from] litellm_gateway_mcp::ConnectError), + #[error(transparent)] + Config(#[from] litellm_config::Error), + #[error("invalid Redis setting {setting}: {reason}")] + RedisInvalid { + setting: &'static str, + reason: &'static str, + }, + #[error(transparent)] + RedisConnect(#[from] litellm_redis_client::ConnectError), #[error("MCP requires a configured master key")] Auth(#[from] litellm_gateway_auth::Error), } diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index 0b3aba81181..2ed2fc80a51 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -1,7 +1,9 @@ mod auth; mod error; +mod redis; mod secrets; pub use error::Error; +pub use redis::RedisClients; use std::{sync::Arc, time::Instant}; @@ -70,6 +72,17 @@ pub fn build_inference(config: &Config) -> Result, Error> { )?)) } +pub async fn build_redis(config: &Config) -> Result { + let environment = secrets::environment_values(&config.environment_variables)?; + let lookup = |name: &str| { + environment + .get(name) + .map(|value| value.expose().to_owned()) + .or_else(|| std::env::var(name).ok()) + }; + redis::connect(config, &lookup).await +} + pub async fn build_mcp( config: &Config, secrets: Arc, @@ -128,12 +141,17 @@ pub fn router( litellm_gateway_auth::authenticate, )) .layer(axum::middleware::from_fn(log_request)); + let health = Router::new().route("/health/readiness", axum::routing::get(readiness)); match ui { - Some(ui) => inference.merge(ui), - None => inference, + Some(ui) => inference.merge(ui).merge(health), + None => inference.merge(health), } } +async fn readiness() -> axum::Json { + axum::Json(serde_json::json!({"status": "healthy", "db": "Not connected"})) +} + async fn log_request(request: Request, next: Next) -> Response { let request_id = Uuid::new_v4().to_string(); let log_body_chunks = tracing::enabled!(tracing::Level::DEBUG); diff --git a/litellm-rust/crates/gateway/src/main.rs b/litellm-rust/crates/gateway/src/main.rs index f8150e6e860..47c1fe83f50 100644 --- a/litellm-rust/crates/gateway/src/main.rs +++ b/litellm-rust/crates/gateway/src/main.rs @@ -46,6 +46,7 @@ async fn main() -> Result<(), Box> { Logger::new(StderrSink { level }).install_global()?; let config = Config::load(settings.litellm_config)?; let inference = litellm_gateway::build_inference(&config)?; + let _redis = litellm_gateway::build_redis(&config).await?; let ui = match std::env::var_os("LITELLM_UI_PATH") { Some(directory) => { let ui_settings: UiSettings = settings::from_iter(std::env::vars_os())?; diff --git a/litellm-rust/crates/gateway/src/redis.rs b/litellm-rust/crates/gateway/src/redis.rs new file mode 100644 index 00000000000..197f40e7988 --- /dev/null +++ b/litellm-rust/crates/gateway/src/redis.rs @@ -0,0 +1,577 @@ +use std::time::Duration; + +use litellm_circuit_breaker::{Breaker, Config as BreakerConfig, State}; +use litellm_config::{ + CacheSettings, Config, Lookup, RedisEndpoint, RedisSettings, ResolvedRedis, Value, + environment_has_target, +}; +use litellm_redis_client::{ + ClusterEndpoint, ClusterNode, DEFAULT_TIMEOUT, Endpoint, HostEndpoint, Namespace, RedisClient, + RedisConfig, +}; + +use crate::Error; + +const ENVIRONMENT_PING_TIMEOUT: Duration = Duration::from_secs(2); +type Params = ResolvedRedis; + +pub struct RedisClients { + pub coordination: Option, + pub transaction_buffer: Option, + pub response_cache: Option, +} + +#[derive(Debug, PartialEq)] +pub(crate) struct Plan { + coordination: Option, + transaction_buffer_fallback: Option, + response_cache: Option, +} + +#[derive(Debug, PartialEq)] +enum Target { + Configured(Params), + Environment(Params), +} + +pub(crate) fn plan(config: &Config, lookup: Lookup<'_>) -> Result { + let response_cache = redis_cache_params(config) + .map(|settings| settings.resolve_cache(lookup)) + .transpose()?; + let coordination = match &config.general_settings.coordination_redis { + Some(settings) => Some(Target::Configured(settings.resolve(lookup)?)), + None => match (&response_cache, config.litellm_settings.cache == Some(true)) { + (Some(cache), _) => Some(Target::Configured(cache.clone())), + (None, true) => environment(lookup)?.map(Target::Configured), + (None, false) => optional_environment(lookup).map(Target::Environment), + }, + }; + let transaction_buffer_fallback = match (&coordination, use_transaction_buffer(config)) { + (Some(Target::Configured(_)), _) | (_, false) => None, + (Some(Target::Environment(_)) | None, true) => environment(lookup)?, + }; + Ok(Plan { + coordination, + transaction_buffer_fallback, + response_cache, + }) +} + +pub async fn connect(config: &Config, lookup: Lookup<'_>) -> Result { + let plan = plan(config, lookup)?; + let breaker = breaker_config(lookup)?; + let coordination = match plan.coordination { + Some(Target::Configured(params)) => Some(client(params, breaker)?), + Some(Target::Environment(params)) => match client(params, breaker) { + Ok(client) => reachable(client).await, + Err(_) => { + tracing::warn!( + "REDIS_* environment variables did not form a valid Redis client, cross-pod coordination stays in memory" + ); + None + } + }, + None => None, + }; + let transaction_buffer = match (&coordination, plan.transaction_buffer_fallback) { + (Some(coordination), _) => Some(coordination.clone()), + (None, Some(params)) => Some(client(params, breaker)?), + (None, None) => None, + }; + let response_cache = plan + .response_cache + .map(|params| client(params, breaker)) + .transpose()?; + Ok(RedisClients { + coordination, + transaction_buffer, + response_cache, + }) +} + +fn redis_cache_params(config: &Config) -> Option<&RedisSettings> { + let settings = &config.litellm_settings; + let params = settings.cache_params.as_ref()?; + match (settings.cache, params) { + (Some(true), CacheSettings::Redis(redis)) => Some(redis), + _ => None, + } +} + +fn environment(lookup: Lookup<'_>) -> Result, Error> { + match environment_has_target(lookup) { + true => RedisSettings::default() + .resolve(lookup) + .map(Some) + .map_err(Into::into), + false => Ok(None), + } +} + +fn optional_environment(lookup: Lookup<'_>) -> Option { + match environment(lookup) { + Ok(params) => params, + Err(error) => { + tracing::warn!(error = %error, "REDIS_* environment variables are invalid, cross-pod coordination stays in memory"); + None + } + } +} + +fn use_transaction_buffer(config: &Config) -> bool { + match config + .general_settings + .additional_fields + .get("use_redis_transaction_buffer") + { + Some(Value::Bool(enabled)) => *enabled, + Some(Value::String(text)) => matches!( + text.to_ascii_lowercase().as_str(), + "true" | "1" | "yes" | "on" + ), + _ => false, + } +} + +async fn reachable(client: RedisClient) -> Option { + match tokio::time::timeout(ENVIRONMENT_PING_TIMEOUT, client.ping()).await { + Ok(Ok(())) => Some(client), + Ok(Err(error)) => { + tracing::warn!(%error, "REDIS_* environment variables named a Redis that is not reachable, cross-pod coordination stays in memory"); + None + } + Err(_) => { + tracing::warn!( + "REDIS_* environment variables named a Redis that did not answer PING in 2s, cross-pod coordination stays in memory" + ); + None + } + } +} + +fn client(params: Params, breaker: Option) -> Result { + tracing::info!(endpoint = ?params.endpoint, namespace = ?params.namespace, "redis client configured"); + let endpoint = match params.endpoint { + RedisEndpoint::Url(url) => Endpoint::Url(url.expose().to_owned()), + RedisEndpoint::Host(host) => Endpoint::Host(HostEndpoint { + host: host.host, + port: host.port, + username: host.username.map(|value| value.expose().to_owned()), + password: host.password.map(|value| value.expose().to_owned()), + tls: host.tls, + db: host.db, + }), + RedisEndpoint::Cluster(cluster) => Endpoint::Cluster(ClusterEndpoint { + nodes: cluster + .nodes + .into_iter() + .map(|node| ClusterNode { + host: node.host, + port: node.port, + }) + .collect(), + username: cluster.username.map(|value| value.expose().to_owned()), + password: cluster.password.map(|value| value.expose().to_owned()), + tls: cluster.tls, + }), + }; + Ok(RedisClient::new(RedisConfig { + endpoint, + namespace: params.namespace.map(Namespace::new).unwrap_or_default(), + timeout: params.timeout.unwrap_or(DEFAULT_TIMEOUT), + breaker: breaker.map(|config| Breaker::new(config, Box::new(log_transition))), + })?) +} + +fn log_transition(transition: &litellm_circuit_breaker::Transition) { + match transition.to { + State::Open { .. } => tracing::warn!( + generation = transition.generation, + "redis circuit breaker opened, fast-failing redis calls" + ), + State::HalfOpen => tracing::info!( + generation = transition.generation, + "redis circuit breaker half open, probing redis" + ), + State::Closed { .. } => tracing::info!( + generation = transition.generation, + "redis circuit breaker closed, redis recovered" + ), + } +} + +fn breaker_config(lookup: Lookup<'_>) -> Result, Error> { + let enabled = lookup("REDIS_CIRCUIT_BREAKER_ENABLED") + .is_none_or(|value| value.eq_ignore_ascii_case("true")); + if !enabled { + return Ok(None); + } + let defaults = BreakerConfig::default(); + let seconds = |name: &'static str, default: Duration| -> Result { + lookup(name) + .map(|value| { + value + .parse::() + .ok() + .and_then(|seconds| Duration::try_from_secs_f64(seconds).ok()) + .ok_or(Error::RedisInvalid { + setting: name, + reason: "must be a non-negative number of seconds", + }) + }) + .unwrap_or(Ok(default)) + }; + let failure_threshold = lookup("REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD") + .map(|value| { + value.parse().map_err(|_| Error::RedisInvalid { + setting: "REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD", + reason: "must be an integer", + }) + }) + .unwrap_or(Ok(defaults.failure_threshold))?; + Ok(Some(BreakerConfig { + failure_threshold, + recovery_timeout: seconds( + "REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT", + defaults.recovery_timeout, + )?, + timeout_min_duration: seconds( + "REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION", + defaults.timeout_min_duration, + )?, + })) +} + +#[cfg(test)] +mod tests { + use std::{collections::HashMap, time::Duration}; + + use litellm_config::{Config, RedisEndpoint, RedisHost}; + use litellm_redis_testing::RedisServer; + use rstest::rstest; + + use super::{Params, Plan, Target, connect, plan}; + use crate::Error; + + fn host(name: &str, port: u16) -> Params { + Params { + endpoint: RedisEndpoint::Host(RedisHost { + host: name.to_owned(), + port, + username: None, + password: None, + tls: false, + db: 0, + }), + namespace: None, + timeout: None, + } + } + + fn resolved(yaml: &str, environment: &[(&str, &str)]) -> Result { + let environment: HashMap = environment + .iter() + .map(|(key, value)| ((*key).to_owned(), (*value).to_owned())) + .collect(); + plan(&Config::from_yaml(yaml).unwrap(), &|name| { + environment.get(name).cloned() + }) + } + + fn configured(coordination: Params, response_cache: Option) -> Plan { + Plan { + coordination: Some(Target::Configured(coordination)), + transaction_buffer_fallback: None, + response_cache, + } + } + + const NO_REDIS: Plan = Plan { + coordination: None, + transaction_buffer_fallback: None, + response_cache: None, + }; + + #[rstest] + #[case::nothing_configured("model_list: []", &[], NO_REDIS)] + #[case::coordination_block_wins_over_the_cache_redis( + "general_settings:\n coordination_redis:\n host: coord\nlitellm_settings:\n cache: true\n cache_params:\n type: redis\n host: cache", + &[], + configured(host("coord", 6379), Some(host("cache", 6379))) + )] + #[case::cache_params_without_cache_true_are_unused( + "litellm_settings:\n cache_params:\n type: redis\n host: cache", + &[], + NO_REDIS + )] + #[case::non_redis_cache_takes_coordination_from_the_environment_without_a_ping( + "litellm_settings:\n cache: true\n cache_params:\n type: local", + &[("REDIS_HOST", "env")], + configured(host("env", 6379), None) + )] + #[case::environment_only_coordination_is_probed( + "model_list: []", + &[("REDIS_HOST", "env")], + Plan { coordination: Some(Target::Environment(host("env", 6379))), transaction_buffer_fallback: None, response_cache: None } + )] + #[case::the_transaction_buffer_falls_back_to_the_environment_only_when_enabled( + "general_settings:\n use_redis_transaction_buffer: \"true\"", + &[("REDIS_HOST", "env")], + Plan { coordination: Some(Target::Environment(host("env", 6379))), transaction_buffer_fallback: Some(host("env", 6379)), response_cache: None } + )] + fn resolves_each_redis_role_like_the_proxy( + #[case] yaml: &str, + #[case] environment: &[(&str, &str)], + #[case] expected: Plan, + ) { + assert_eq!(resolved(yaml, environment).unwrap(), expected); + } + + #[rstest] + fn redis_migration_regression_preserves_borrowed_cache_namespace() { + let yaml = "litellm_settings:\n cache: true\n cache_params:\n type: redis\n host: cache\n namespace: responses\n ttl: 600"; + let cache = Params { + namespace: Some("responses".to_owned()), + ..host("cache", 6379) + }; + + assert_eq!( + resolved(yaml, &[]).unwrap(), + configured(cache.clone(), Some(cache)) + ); + } + + #[rstest] + #[tokio::test] + async fn borrowed_coordination_and_transaction_buffer_share_the_cache_namespace() { + let server = RedisServer::start().await.unwrap(); + let config = Config::from_yaml(&format!( + "litellm_settings:\n cache: true\n cache_params:\n type: redis\n host: 127.0.0.1\n port: {}\n namespace: responses", + port_of(&server) + )) + .unwrap(); + let clients = connect(&config, &|_| None).await.unwrap(); + let coordination = clients.coordination.unwrap(); + let buffer = clients.transaction_buffer.unwrap(); + let cache = clients.response_cache.unwrap(); + let unnamespaced = connect(&coordination_config(&server), &|_| None) + .await + .unwrap() + .coordination + .unwrap(); + + coordination + .set("rate-limit", "1", Duration::from_secs(60)) + .await + .unwrap(); + buffer + .set("spend", "2", Duration::from_secs(60)) + .await + .unwrap(); + + assert_eq!( + cache.get::("rate-limit").await.unwrap().as_deref(), + Some("1") + ); + assert_eq!( + cache.get::("spend").await.unwrap().as_deref(), + Some("2") + ); + assert_eq!( + unnamespaced + .get::("responses:rate-limit") + .await + .unwrap() + .as_deref(), + Some("1") + ); + assert_eq!( + unnamespaced + .get::("responses:spend") + .await + .unwrap() + .as_deref(), + Some("2") + ); + assert!( + unnamespaced + .get::("rate-limit") + .await + .unwrap() + .is_none() + ); + assert!(unnamespaced.get::("spend").await.unwrap().is_none()); + } + + #[rstest] + #[case::malformed_port(&[("REDIS_HOST", "env"), ("REDIS_PORT", "not-a-port")])] + #[case::malformed_cluster_nodes(&[("REDIS_CLUSTER_NODES", "[{\"host\": ")])] + fn redis_migration_regression_malformed_optional_environment_is_absent( + #[case] environment: &[(&str, &str)], + ) { + assert_eq!(resolved("model_list: []", environment).unwrap(), NO_REDIS); + } + + #[rstest] + fn redis_migration_regression_explicit_malformed_setting_still_fails() { + let error = resolved( + "general_settings:\n coordination_redis:\n host: coord\n port: not-a-port", + &[], + ) + .unwrap_err(); + + assert!(matches!( + error, + Error::Config(litellm_config::Error::RedisInvalid { + setting: "port", + .. + }) + )); + } + + #[rstest] + fn redis_migration_regression_enabled_transaction_buffer_rejects_invalid_environment() { + let error = resolved( + "general_settings:\n use_redis_transaction_buffer: true", + &[("REDIS_HOST", "env"), ("REDIS_PORT", "not-a-port")], + ) + .unwrap_err(); + + assert!(matches!( + error, + Error::Config(litellm_config::Error::RedisInvalid { + setting: "port", + .. + }) + )); + } + + #[rstest] + fn redis_migration_regression_invalid_response_cache_still_fails() { + let error = resolved( + "litellm_settings:\n cache: true\n cache_params:\n type: redis\n host: cache\n port: not-a-port", + &[], + ) + .unwrap_err(); + + assert!(matches!( + error, + Error::Config(litellm_config::Error::RedisInvalid { + setting: "port", + .. + }) + )); + } + + #[rstest] + #[tokio::test] + async fn redis_migration_regression_invalid_optional_environment_url_is_absent() { + let config = Config::from_yaml("model_list: []").unwrap(); + + let clients = connect(&config, &|name| { + (name == "REDIS_URL").then(|| "not a redis URL".to_owned()) + }) + .await + .unwrap(); + + assert!(clients.coordination.is_none()); + } + + fn port_of(server: &RedisServer) -> String { + server.url().rsplit(':').next().unwrap().to_owned() + } + + fn coordination_config(server: &RedisServer) -> Config { + Config::from_yaml(&format!( + "general_settings:\n coordination_redis:\n host: 127.0.0.1\n port: {}\n socket_timeout: 0.5", + port_of(server) + )) + .unwrap() + } + + async fn eventually_pings(client: &litellm_redis_client::RedisClient) -> bool { + let deadline = tokio::time::Instant::now() + Duration::from_secs(30); + while tokio::time::Instant::now() < deadline { + if client.ping().await.is_ok() { + return true; + } + tokio::time::sleep(Duration::from_millis(200)).await; + } + false + } + + #[rstest] + #[tokio::test] + async fn boot_does_not_wait_for_a_configured_redis_that_is_down() { + let server = RedisServer::start().await.unwrap(); + server.stop().await.unwrap(); + + let clients = connect(&coordination_config(&server), &|_| None) + .await + .unwrap(); + let coordination = clients.coordination.unwrap(); + let while_down = coordination.ping().await; + server.resume().await.unwrap(); + + assert!(while_down.is_err()); + assert!(eventually_pings(&coordination).await); + } + + #[rstest] + #[tokio::test] + async fn the_coordination_client_survives_a_redis_restart_and_backs_the_transaction_buffer() { + let server = RedisServer::start().await.unwrap(); + let clients = connect(&coordination_config(&server), &|_| None) + .await + .unwrap(); + let coordination = clients.coordination.unwrap(); + coordination + .set("written", "1", Duration::from_secs(60)) + .await + .unwrap(); + + server.restart().await.unwrap(); + + assert!(eventually_pings(&coordination).await); + let buffer = clients.transaction_buffer.unwrap(); + assert!(buffer.ping().await.is_ok()); + } + + #[rstest] + #[tokio::test] + async fn an_unreachable_environment_only_redis_leaves_coordination_in_memory() { + let server = RedisServer::start().await.unwrap(); + let port = port_of(&server); + server.stop().await.unwrap(); + let config = Config::from_yaml("model_list: []").unwrap(); + + let clients = connect(&config, &|name| match name { + "REDIS_HOST" => Some("127.0.0.1".to_owned()), + "REDIS_PORT" => Some(port.clone()), + _ => None, + }) + .await + .unwrap(); + + assert!(clients.coordination.is_none()); + assert!(clients.transaction_buffer.is_none()); + } + + #[rstest] + #[tokio::test] + async fn a_reachable_environment_only_redis_becomes_coordination() { + let server = RedisServer::start().await.unwrap(); + let port = port_of(&server); + let config = Config::from_yaml("model_list: []").unwrap(); + + let clients = connect(&config, &|name| match name { + "REDIS_HOST" => Some("127.0.0.1".to_owned()), + "REDIS_PORT" => Some(port.clone()), + _ => None, + }) + .await + .unwrap(); + + assert!(clients.coordination.unwrap().ping().await.is_ok()); + } +} diff --git a/litellm-rust/crates/gateway/tests/server.rs b/litellm-rust/crates/gateway/tests/server.rs index 71edd3c3650..53a97e7098c 100644 --- a/litellm-rust/crates/gateway/tests/server.rs +++ b/litellm-rust/crates/gateway/tests/server.rs @@ -287,3 +287,33 @@ async fn ui_routes_are_absent_when_not_mounted( .unwrap(); assert_eq!(response.status().as_u16(), 404); } + +#[rstest] +#[case::with_master_key("model_list: []\ngeneral_settings:\n master_key: gateway-key\n")] +#[case::without_master_key("model_list: []\n")] +#[tokio::test] +async fn readiness_answers_without_credentials_like_the_proxy_with_no_database( + inference: Arc, + #[case] yaml: &str, +) { + let config = Config::from_yaml(yaml).unwrap(); + let app = litellm_gateway::router(inference, &config, None, None); + + let response = app + .oneshot( + Request::get("/health/readiness") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status().as_u16(), 200); + let body = axum::body::to_bytes(response.into_body(), 1024) + .await + .unwrap(); + assert_eq!( + serde_json::from_slice::(&body).unwrap(), + json!({"status": "healthy", "db": "Not connected"}) + ); +} diff --git a/litellm-rust/crates/inference-messages/src/common_utils.rs b/litellm-rust/crates/inference-messages/src/common_utils.rs index 7394d7ae21e..9be7ee6ffc3 100644 --- a/litellm-rust/crates/inference-messages/src/common_utils.rs +++ b/litellm-rust/crates/inference-messages/src/common_utils.rs @@ -6,7 +6,7 @@ use litellm_llms::{ anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG, azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG, base_llm::messages::transformation::BaseMessagesConfig, - bedrock::messages::invoke_transformations::anthropic_claude3_transformation::BEDROCK_ANTHROPIC_MESSAGES_CONFIG, + bedrock::messages::invoke_transformations::anthropic_claude_transformation::BEDROCK_ANTHROPIC_MESSAGES_CONFIG, deepseek::messages::transformation::DEEPSEEK_ANTHROPIC_MESSAGES_CONFIG, edenai::messages::transformation::EDENAI_MESSAGES_CONFIG, minimax::messages::transformation::MINIMAX_MESSAGES_CONFIG, diff --git a/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.rs b/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs similarity index 100% rename from litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.rs rename to litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs diff --git a/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/mod.rs b/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/mod.rs index 4d67a0c0696..68ceaac3444 100644 --- a/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/mod.rs +++ b/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/mod.rs @@ -1 +1 @@ -pub mod anthropic_claude3_transformation; +pub mod anthropic_claude_transformation; diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 6fc05c43146..8afde8362f6 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -53,7 +53,7 @@ pyo3.workspace = true prost.workspace = true pyo3-async-runtimes.workspace = true reqwest.workspace = true -redis = { version = "1.7.0", features = ["tls-rustls"] } +redis = { workspace = true, features = ["tls-rustls"] } serde_json.workspace = true strum.workspace = true veil.workspace = true diff --git a/litellm-rust/crates/redis-client/AGENTS.md b/litellm-rust/crates/redis-client/AGENTS.md new file mode 100644 index 00000000000..c73a4ed8c4d --- /dev/null +++ b/litellm-rust/crates/redis-client/AGENTS.md @@ -0,0 +1,43 @@ +- The one async Redis client for the gateway and the Rust cache crates. It owns connecting, timeouts, error classes and the key namespace +- It knows nothing about domains, so key builders, TTL policy, value formats and fallback decisions stay in the crate that owns the state. Every write takes its TTL from the caller, and there is no default +- Async only. Build one `RedisClient` per Redis instance at the process boundary and clone it. Never add a pool, a sync API, `spawn_blocking`, or a client built inside a route +- Set every redis-rs timeout and retry explicitly. Its defaults (500ms response, 1s connect, 6 reconnect retries) are not the Python proxy's (5s socket timeout from `REDIS_SOCKET_TIMEOUT`) +- Resolving config from `config.yaml` / `REDIS_*` env the way `litellm/_redis.py` does belongs to the consumer that reads that config, not here. `RedisConfig` takes already-resolved values + +## Errors +- `Transport`, `Server`, `Decode`, each wrapping the redis-rs error as `#[source]`. `ConnectError` is separate because only `connect` returns it +- Classify in `From for Error` by `ErrorKind`. LOADING, CLUSTERDOWN, MASTERDOWN, TRYAGAIN and READONLY are typed server replies, but they mean the node can't serve right now, so they are `Transport`. Codes redis-rs doesn't type (WRONGTYPE) come back as `ErrorKind::Extension` and are `Server`. Each code gets its own rstest case built from a parsed RESP error +- A `Server` error never drops the connection +- `Transport` includes timeouts, where the command may still have applied. Callers of counters and locks treat it as an unknown outcome +- Never swallow an error or turn it into a miss. The caller picks the fallback +- Reads return `Option`: a missing key is `None`, and a value that doesn't decode to `T` (`"15.0"` as `i64`) is `Decode` + +## Circuit breaker +- The `litellm-circuit-breaker` `Breaker` is injected through `RedisConfig::breaker` and shared by every clone. `None` turns it off +- Every public call (each command, `eval`, one `run` of a pipeline) takes one admission through `guarded` / `admit` in `client.rs`. Never call the connection around it +- `Transport` counts as a failure (`Timeout` when `is_timeout()`, else `Hard`). `Server` and `Decode` count as nothing, since Redis answered and the request was bad. A probe that gets one still reopens, through the dropped-probe rule in `circuit-breaker`. A pipeline counts one failure if any op came back `Transport`, else one success +- An open breaker returns `Error::BreakerOpen` without touching the socket +- Python doesn't count MASTERDOWN, TRYAGAIN or READONLY as health failures. Here they are `Transport` and do count + +## Namespace +- Applied once, here, to every key. A key that already starts with `{namespace}:` is left alone, matching `RedisCache.check_and_fix_namespace` +- The rule never changes, because Python and Rust pods share keys in the same Redis + +## Commands +- One typed method per command, built with `redis::cmd`. Empty batches (MGET, DEL) return without a round trip, since Redis rejects an empty MGET +- On cluster, redis-rs splits MGET and DEL by slot and returns MGET values in input order. Other multi-key commands are not split, so a new one needs a cross-slot test +- TTLs are `Duration` and are sent in milliseconds (`PX`, `PEXPIRE`, `PTTL`) + +## Scripts and pipelines +- A `Script` declares every key it touches in KEYS. On cluster an undeclared key works or fails depending on which node owns its slot, so it can pass tests and break in production +- `eval` goes through redis-rs `invoke_async`, which loads the script on NOSCRIPT. Inside a pipeline nothing reloads, so a NOSCRIPT op reruns alone with `invoke_async` +- A pipeline is never atomic. Every op gets its own `Result`, from `.ignore_errors()`, so one rejected op leaves the others alone. A transport failure fails every op of that round trip +- redis-rs rejects an async cluster pipeline whose ops span slots. `Pipeline` groups ops by the slot of their first key, runs the groups concurrently and returns results in input order +- Pipeline ops take one key each. A multi-key op would span slots inside one group +- Command builders live in `commands.rs` and are shared by `RedisClient` and `Pipeline`, so a command is written once + +## Tests +- Against real Redis from `litellm-redis-testing`, never a mocked wire or asserted command bytes +- Every command test runs on each topology through `#[values(Topology::Standalone, Topology::Cluster)]`. A test is single-topology only when the behavior exists on one (a `db` number has no cluster meaning) +- Assert values and Redis state. Outages come from the container (`restart`, `freeze` / `thaw`), not `DEBUG`, which the official image disables +- Use a short client timeout in outage tests so they stay fast. Never sleep to wait for a state; poll with a deadline diff --git a/litellm-rust/crates/redis-client/Cargo.toml b/litellm-rust/crates/redis-client/Cargo.toml new file mode 100644 index 00000000000..eab57b36423 --- /dev/null +++ b/litellm-rust/crates/redis-client/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "litellm-redis-client" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-circuit-breaker.workspace = true +futures-util.workspace = true +redis = { workspace = true, features = ["cluster-async", "connection-manager", "tokio-rustls-comp"] } +thiserror.workspace = true + +[dev-dependencies] +litellm-redis-testing.workspace = true +rstest.workspace = true +tokio.workspace = true diff --git a/litellm-rust/crates/redis-client/src/client.rs b/litellm-rust/crates/redis-client/src/client.rs new file mode 100644 index 00000000000..a02b0847663 --- /dev/null +++ b/litellm-rust/crates/redis-client/src/client.rs @@ -0,0 +1,647 @@ +use std::{ + future::Future, + time::{Duration, Instant}, +}; + +use litellm_circuit_breaker::{Admission, Breaker, Failure}; +use redis::{Cmd, FromRedisValue, IntoConnectionInfo, ToRedisArgs, ToSingleRedisArg, Value}; + +use crate::{ + ConnectError, Endpoint, Error, Namespace, Pipeline, RedisConfig, Script, commands, + connection::Connection, +}; + +#[derive(Clone)] +pub struct RedisClient { + pub(crate) connection: Connection, + namespace: Namespace, + breaker: Option, +} + +#[derive(Debug, PartialEq, Eq)] +pub enum Ttl { + Missing, + Persistent, + Expires(Duration), +} + +impl RedisClient { + pub fn new(config: RedisConfig) -> Result { + let connection = match &config.endpoint { + Endpoint::Url(url) => { + Connection::standalone(url.as_str().into_connection_info()?, config.timeout)? + } + Endpoint::Host(host) => { + Connection::standalone(host.connection_info()?, config.timeout)? + } + Endpoint::Cluster(cluster) => Connection::cluster(cluster, config.timeout)?, + }; + Ok(Self { + connection, + namespace: config.namespace, + breaker: config.breaker, + }) + } + + pub async fn ping(&self) -> Result<(), Error> { + self.query(&commands::ping()).await + } + + pub async fn get(&self, key: &str) -> Result, Error> { + self.query(&commands::get(self.key(key))).await + } + + pub async fn mget(&self, keys: &[&str]) -> Result>, Error> { + if keys.is_empty() { + return Ok(Vec::new()); + } + self.query(&commands::mget(self.keys(keys))).await + } + + pub async fn set( + &self, + key: &str, + value: impl ToSingleRedisArg, + ttl: Duration, + ) -> Result<(), Error> { + self.query(&commands::set(self.key(key), value, ttl)).await + } + + pub async fn set_nx( + &self, + key: &str, + value: impl ToSingleRedisArg, + ttl: Duration, + ) -> Result { + let reply: Option = self + .query(&commands::set_nx(self.key(key), value, ttl)) + .await?; + Ok(reply.is_some()) + } + + pub async fn del(&self, keys: &[&str]) -> Result { + if keys.is_empty() { + return Ok(0); + } + self.query(&commands::del(self.keys(keys))).await + } + + pub async fn expire(&self, key: &str, ttl: Duration) -> Result { + self.query(&commands::expire(self.key(key), ttl)).await + } + + pub async fn ttl(&self, key: &str) -> Result { + let reply: i64 = self.query(&commands::ttl(self.key(key))).await?; + Ok(match reply { + -2 => Ttl::Missing, + -1 => Ttl::Persistent, + millis => Ttl::Expires(Duration::from_millis(millis.unsigned_abs())), + }) + } + + pub async fn incr_by(&self, key: &str, delta: i64) -> Result { + self.query(&commands::incr_by(self.key(key), delta)).await + } + + pub async fn incr_by_float(&self, key: &str, delta: f64) -> Result { + self.query(&commands::incr_by_float(self.key(key), delta)) + .await + } + + pub async fn rpush(&self, key: &str, values: impl ToRedisArgs) -> Result { + self.query(&commands::rpush(self.key(key), values)).await + } + + pub async fn eval( + &self, + script: &Script, + keys: &[&str], + args: impl ToRedisArgs, + ) -> Result { + let keys = self.keys(keys); + let args = args.to_redis_args(); + self.guarded(async { + Ok(script + .invocation(&keys, &args) + .invoke_async(&mut self.connection.clone()) + .await?) + }) + .await + } + + pub fn pipeline(&self) -> Pipeline { + Pipeline::new(self.namespace.clone()) + } + + pub async fn run(&self, pipeline: &Pipeline) -> Vec> { + let admission = match self.admit() { + Ok(admission) => admission, + Err(_) => { + return (0..pipeline.len()) + .map(|_| Err(Error::BreakerOpen)) + .collect(); + } + }; + let results = pipeline.run(&self.connection).await; + let transport = results.iter().find_map(|result| { + result + .as_ref() + .err() + .filter(|error| matches!(error, Error::Transport(_))) + }); + settle(admission, transport); + results + } + + fn key(&self, key: &str) -> String { + self.namespace.apply(key) + } + + fn keys(&self, keys: &[&str]) -> Vec { + keys.iter().map(|key| self.key(key)).collect() + } + + async fn query(&self, command: &Cmd) -> Result { + self.guarded(async { Ok(command.query_async(&mut self.connection.clone()).await?) }) + .await + } + + async fn guarded(&self, call: impl Future>) -> Result { + let admission = self.admit()?; + let result = call.await; + settle(admission, result.as_ref().err()); + result + } + + fn admit(&self) -> Result, Error> { + self.breaker + .as_ref() + .map(|breaker| { + breaker + .admit(Instant::now()) + .map_err(|_| Error::BreakerOpen) + }) + .transpose() + } +} + +fn settle(admission: Option, error: Option<&Error>) { + let Some(admission) = admission else { + return; + }; + match error { + None => admission.succeeded(), + Some(Error::Transport(error)) if error.is_timeout() => { + admission.failed(Failure::Timeout, Instant::now()) + } + Some(Error::Transport(_)) => admission.failed(Failure::Hard, Instant::now()), + Some(Error::Server(_) | Error::Decode(_) | Error::BreakerOpen) => drop(admission), + } +} + +#[cfg(test)] +mod tests { + use std::time::{Duration, Instant}; + + use litellm_redis_testing::RedisServer; + use rstest::rstest; + + use litellm_circuit_breaker::{Breaker, Config, State}; + + use super::{RedisClient, Ttl}; + use crate::{ + DEFAULT_TIMEOUT, Endpoint, Error, HostEndpoint, Namespace, RedisConfig, + test_support::{Server, Topology, client_id}, + }; + + #[rstest] + #[tokio::test] + async fn namespaces_isolate_keys_on_the_server( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + ) { + let server = Server::start(topology).await; + let a = server.client("a", DEFAULT_TIMEOUT).await; + let b = server.client("b", DEFAULT_TIMEOUT).await; + let bare = server.client("", DEFAULT_TIMEOUT).await; + + a.set("x", "1", Duration::from_secs(60)).await.unwrap(); + + assert_eq!(a.get::("x").await.unwrap().as_deref(), Some("1")); + assert_eq!(a.get::("a:x").await.unwrap().as_deref(), Some("1")); + assert_eq!( + bare.get::("a:x").await.unwrap().as_deref(), + Some("1") + ); + assert_eq!(b.get::("x").await.unwrap(), None); + } + + #[rstest] + #[tokio::test] + async fn a_rejected_command_keeps_the_connection( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + ) { + let server = Server::start(topology).await; + let client = server.client("", DEFAULT_TIMEOUT).await; + client.rpush("list", "item").await.unwrap(); + let before = client_id(&client, "list").await; + + let rejected = client.incr_by("list", 1).await; + + assert!(matches!(rejected, Err(Error::Server(_))), "{rejected:?}"); + assert_eq!(client.rpush("list", "next").await.unwrap(), 2); + assert_eq!(client_id(&client, "list").await, before); + } + + #[rstest] + #[tokio::test] + async fn a_reply_of_another_type_is_a_decode_error_and_a_missing_key_is_none( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + ) { + let server = Server::start(topology).await; + let client = server.client("", DEFAULT_TIMEOUT).await; + client + .set("float", "15.0", Duration::from_secs(60)) + .await + .unwrap(); + + let decoded = client.get::("float").await; + + assert!(matches!(decoded, Err(Error::Decode(_))), "{decoded:?}"); + assert_eq!(client.get::("missing").await.unwrap(), None); + } + + #[rstest] + #[tokio::test] + async fn a_timed_out_command_does_not_leak_its_reply_into_the_next_one( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + ) { + let server = Server::start(topology).await; + let client = server.client("", Duration::from_millis(500)).await; + client + .set("first", "1", Duration::from_secs(60)) + .await + .unwrap(); + client + .set("second", "2", Duration::from_secs(60)) + .await + .unwrap(); + + server.freeze().await; + let timed_out = client.get::("first").await; + server.thaw().await; + + assert!( + matches!(timed_out, Err(Error::Transport(_))), + "{timed_out:?}" + ); + assert_eq!( + client.get::("second").await.unwrap().as_deref(), + Some("2") + ); + } + + #[rstest] + #[tokio::test] + async fn the_same_client_recovers_after_redis_restarts( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + ) { + let server = Server::start(topology).await; + let client = server.client("", Duration::from_millis(500)).await; + + server.restart().await; + + let deadline = Instant::now() + Duration::from_secs(30); + let recovered = loop { + match client.incr_by("after_restart", 1).await { + Ok(value) => break Some(value), + Err(_) if Instant::now() < deadline => { + tokio::time::sleep(Duration::from_millis(200)).await + } + Err(_) => break None, + } + }; + assert_eq!(recovered, Some(1)); + } + + #[rstest] + #[tokio::test] + async fn a_client_created_while_redis_is_down_works_once_redis_is_up( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + ) { + let server = Server::start(topology).await; + server.stop().await; + + let client = server.client("", Duration::from_millis(500)).await; + let while_down = client.ping().await; + server.resume().await; + + assert!( + matches!(while_down, Err(Error::Transport(_))), + "{while_down:?}" + ); + let deadline = Instant::now() + Duration::from_secs(30); + let recovered = loop { + match client.ping().await { + Ok(()) => break true, + Err(_) if Instant::now() < deadline => { + tokio::time::sleep(Duration::from_millis(200)).await + } + Err(_) => break false, + } + }; + assert!(recovered); + } + + #[rstest] + #[tokio::test] + async fn empty_batches_succeed_without_asking_redis( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + ) { + let server = Server::start(topology).await; + let client = server.client("", DEFAULT_TIMEOUT).await; + + assert_eq!( + client.mget::(&[]).await.unwrap(), + Vec::>::new() + ); + assert_eq!(client.del(&[]).await.unwrap(), 0); + } + + #[rstest] + #[tokio::test] + async fn mget_and_del_span_slots_in_input_order( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + ) { + let server = Server::start(topology).await; + let client = server.client("", DEFAULT_TIMEOUT).await; + client.set("a", "1", Duration::from_secs(60)).await.unwrap(); + client.set("c", "3", Duration::from_secs(60)).await.unwrap(); + + let values = client.mget::(&["c", "b", "a"]).await.unwrap(); + let deleted = client.del(&["a", "b", "c"]).await.unwrap(); + + assert_eq!( + values, + vec![Some("3".to_owned()), None, Some("1".to_owned())] + ); + assert_eq!(deleted, 2); + assert_eq!( + client.mget::(&["a", "c"]).await.unwrap(), + vec![None, None] + ); + } + + #[rstest] + #[tokio::test] + async fn set_nx_only_writes_a_missing_key( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + ) { + let server = Server::start(topology).await; + let client = server.client("", DEFAULT_TIMEOUT).await; + + let first = client + .set_nx("lock", "owner-1", Duration::from_secs(60)) + .await + .unwrap(); + let second = client + .set_nx("lock", "owner-2", Duration::from_secs(60)) + .await + .unwrap(); + + assert!(first); + assert!(!second); + assert_eq!( + client.get::("lock").await.unwrap().as_deref(), + Some("owner-1") + ); + } + + #[rstest] + #[tokio::test] + async fn ttl_reports_missing_persistent_and_remaining_time( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + ) { + let server = Server::start(topology).await; + let client = server.client("", DEFAULT_TIMEOUT).await; + client.incr_by("persistent", 1).await.unwrap(); + client + .set("expiring", "1", Duration::from_secs(60)) + .await + .unwrap(); + client.incr_by_float("expired_later", 0.5).await.unwrap(); + assert!( + client + .expire("expired_later", Duration::from_secs(30)) + .await + .unwrap() + ); + + assert_eq!(client.ttl("missing").await.unwrap(), Ttl::Missing); + assert_eq!(client.ttl("persistent").await.unwrap(), Ttl::Persistent); + assert!(matches!( + client.ttl("expiring").await.unwrap(), + Ttl::Expires(left) if left > Duration::from_secs(50) && left <= Duration::from_secs(60) + )); + assert!(matches!( + client.ttl("expired_later").await.unwrap(), + Ttl::Expires(left) if left > Duration::from_secs(20) && left <= Duration::from_secs(30) + )); + assert!( + !client + .expire("missing", Duration::from_secs(30)) + .await + .unwrap() + ); + } + + #[rstest] + #[tokio::test] + async fn incr_by_float_accumulates_fractions( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + ) { + let server = Server::start(topology).await; + let client = server.client("", DEFAULT_TIMEOUT).await; + + client.incr_by_float("spend", 0.25).await.unwrap(); + + assert_eq!(client.incr_by_float("spend", 0.5).await.unwrap(), 0.75); + } + + #[rstest] + #[tokio::test] + async fn concurrent_commands_on_one_cloned_client_all_apply( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + ) { + let server = Server::start(topology).await; + let client = server.client("", DEFAULT_TIMEOUT).await; + + let tasks: Vec<_> = (0..100) + .map(|_| { + let client = client.clone(); + tokio::spawn(async move { + for _ in 0..10 { + client.incr_by("shared", 1).await.unwrap(); + } + }) + }) + .collect(); + for task in tasks { + task.await.unwrap(); + } + + assert_eq!(client.get::("shared").await.unwrap(), Some(1000)); + } + + fn breaker(failure_threshold: u32) -> Breaker { + Breaker::new( + Config { + failure_threshold, + recovery_timeout: Duration::from_millis(200), + timeout_min_duration: Duration::ZERO, + }, + Box::new(|_| {}), + ) + } + + #[rstest] + #[tokio::test] + async fn an_open_breaker_rejects_without_waiting_on_redis_and_a_probe_closes_it( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + ) { + let server = Server::start(topology).await; + let breaker = breaker(2); + let client = server + .guarded_client("", Duration::from_millis(300), Some(breaker.clone())) + .await; + + server.freeze().await; + let first = client.get::("k").await; + let second = client.get::("k").await; + let started = Instant::now(); + let third = client.get::("k").await; + let rejected_in = started.elapsed(); + server.thaw().await; + + assert!(matches!(first, Err(Error::Transport(_))), "{first:?}"); + assert!(matches!(second, Err(Error::Transport(_))), "{second:?}"); + assert!(matches!(third, Err(Error::BreakerOpen)), "{third:?}"); + assert!(rejected_in < Duration::from_millis(50), "{rejected_in:?}"); + + let deadline = Instant::now() + Duration::from_secs(30); + let recovered = loop { + match client.get::("k").await { + Ok(value) => break Some(value), + Err(_) if Instant::now() < deadline => { + tokio::time::sleep(Duration::from_millis(50)).await + } + Err(_) => break None, + } + }; + assert_eq!(recovered, Some(None)); + assert!( + matches!(breaker.state(), State::Closed { .. }), + "{:?}", + breaker.state() + ); + } + + #[rstest] + #[tokio::test] + async fn rejected_commands_never_open_the_breaker( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + ) { + let server = Server::start(topology).await; + let breaker = breaker(2); + let client = server + .guarded_client("", DEFAULT_TIMEOUT, Some(breaker.clone())) + .await; + client.rpush("list", "item").await.unwrap(); + + let rejected = [ + client.incr_by("list", 1).await, + client.incr_by("list", 1).await, + client.incr_by("list", 1).await, + ]; + + assert!( + matches!(rejected[2], Err(Error::Server(_))), + "{:?}", + rejected[2] + ); + assert!( + matches!(breaker.state(), State::Closed { hard: 0, .. }), + "{:?}", + breaker.state() + ); + } + + #[rstest] + #[tokio::test] + async fn an_open_breaker_fails_every_pipeline_op_without_sending( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + ) { + let server = Server::start(topology).await; + let client = server + .guarded_client("", Duration::from_millis(300), Some(breaker(1))) + .await; + let pipeline = client.pipeline().incr_by("a", 1).incr_by("c", 1); + + server.freeze().await; + let failed = client.run(&pipeline).await; + let rejected = client.run(&pipeline).await; + server.thaw().await; + + assert!( + matches!(failed[0], Err(Error::Transport(_))), + "{:?}", + failed[0] + ); + assert!( + matches!(rejected[0], Err(Error::BreakerOpen)), + "{:?}", + rejected[0] + ); + assert!( + matches!(rejected[1], Err(Error::BreakerOpen)), + "{:?}", + rejected[1] + ); + } + + #[rstest] + #[tokio::test] + async fn a_host_endpoint_connects_to_its_database() { + let server = RedisServer::start().await.unwrap(); + let port = server.url().rsplit(':').next().unwrap().parse().unwrap(); + let connect = |endpoint| { + RedisClient::new(RedisConfig { + endpoint, + namespace: Namespace::default(), + timeout: DEFAULT_TIMEOUT, + breaker: None, + }) + }; + let in_db_3 = connect(Endpoint::Host(HostEndpoint { + host: "127.0.0.1".to_owned(), + port, + username: None, + password: None, + tls: false, + db: 3, + })) + .unwrap(); + let in_db_0 = connect(Endpoint::Url(server.url().to_owned())).unwrap(); + let via_url = connect(Endpoint::Url(format!("{}/3", server.url()))).unwrap(); + + in_db_3 + .set("k", "3", Duration::from_secs(60)) + .await + .unwrap(); + + assert_eq!(in_db_0.get::("k").await.unwrap(), None); + assert_eq!( + via_url.get::("k").await.unwrap().as_deref(), + Some("3") + ); + } +} diff --git a/litellm-rust/crates/redis-client/src/commands.rs b/litellm-rust/crates/redis-client/src/commands.rs new file mode 100644 index 00000000000..fb8a4a28c94 --- /dev/null +++ b/litellm-rust/crates/redis-client/src/commands.rs @@ -0,0 +1,62 @@ +use std::time::Duration; + +use redis::{Cmd, ToRedisArgs, ToSingleRedisArg}; + +pub(crate) fn ping() -> Cmd { + redis::cmd("PING") +} + +pub(crate) fn get(key: String) -> Cmd { + redis::cmd("GET").arg(key).to_owned() +} + +pub(crate) fn mget(keys: Vec) -> Cmd { + redis::cmd("MGET").arg(keys).to_owned() +} + +pub(crate) fn set(key: String, value: impl ToSingleRedisArg, ttl: Duration) -> Cmd { + redis::cmd("SET") + .arg(key) + .arg(value) + .arg("PX") + .arg(millis(ttl)) + .to_owned() +} + +pub(crate) fn set_nx(key: String, value: impl ToSingleRedisArg, ttl: Duration) -> Cmd { + redis::cmd("SET") + .arg(key) + .arg(value) + .arg("NX") + .arg("PX") + .arg(millis(ttl)) + .to_owned() +} + +pub(crate) fn del(keys: Vec) -> Cmd { + redis::cmd("DEL").arg(keys).to_owned() +} + +pub(crate) fn expire(key: String, ttl: Duration) -> Cmd { + redis::cmd("PEXPIRE").arg(key).arg(millis(ttl)).to_owned() +} + +pub(crate) fn ttl(key: String) -> Cmd { + redis::cmd("PTTL").arg(key).to_owned() +} + +pub(crate) fn incr_by(key: String, delta: i64) -> Cmd { + redis::cmd("INCRBY").arg(key).arg(delta).to_owned() +} + +pub(crate) fn incr_by_float(key: String, delta: f64) -> Cmd { + redis::cmd("INCRBYFLOAT").arg(key).arg(delta).to_owned() +} + +pub(crate) fn rpush(key: String, values: impl ToRedisArgs) -> Cmd { + redis::cmd("RPUSH").arg(key).arg(values).to_owned() +} + +fn millis(duration: Duration) -> u64 { + u64::try_from(duration.as_millis()).unwrap_or(u64::MAX) +} diff --git a/litellm-rust/crates/redis-client/src/config.rs b/litellm-rust/crates/redis-client/src/config.rs new file mode 100644 index 00000000000..64934c68891 --- /dev/null +++ b/litellm-rust/crates/redis-client/src/config.rs @@ -0,0 +1,152 @@ +use std::{fmt, time::Duration}; + +use redis::{ConnectionAddr, ConnectionInfo, IntoConnectionInfo, RedisConnectionInfo, RedisResult}; + +use litellm_circuit_breaker::Breaker; + +use crate::Namespace; + +pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(5); + +pub struct RedisConfig { + pub endpoint: Endpoint, + pub namespace: Namespace, + pub timeout: Duration, + pub breaker: Option, +} + +#[derive(Clone, PartialEq)] +pub enum Endpoint { + Url(String), + Host(HostEndpoint), + Cluster(ClusterEndpoint), +} + +#[derive(Clone, PartialEq)] +pub struct HostEndpoint { + pub host: String, + pub port: u16, + pub username: Option, + pub password: Option, + pub tls: bool, + pub db: i64, +} + +#[derive(Clone, PartialEq)] +pub struct ClusterEndpoint { + pub nodes: Vec, + pub username: Option, + pub password: Option, + pub tls: bool, +} + +#[derive(Clone, Debug, PartialEq)] +pub struct ClusterNode { + pub host: String, + pub port: u16, +} + +impl ClusterEndpoint { + pub(crate) fn addrs(&self) -> impl Iterator + '_ { + self.nodes + .iter() + .map(|node| addr(&node.host, node.port, self.tls)) + } +} + +impl HostEndpoint { + pub(crate) fn connection_info(&self) -> RedisResult { + let settings = RedisConnectionInfo::default().set_db(self.db); + let settings = match &self.username { + Some(username) => settings.set_username(username), + None => settings, + }; + let settings = match &self.password { + Some(password) => settings.set_password(password), + None => settings, + }; + Ok(addr(&self.host, self.port, self.tls) + .into_connection_info()? + .set_redis_settings(settings)) + } +} + +fn addr(host: &str, port: u16, tls: bool) -> ConnectionAddr { + match tls { + true => ConnectionAddr::TcpTls { + host: host.to_owned(), + port, + insecure: false, + tls_params: None, + }, + false => ConnectionAddr::Tcp(host.to_owned(), port), + } +} + +impl fmt::Debug for Endpoint { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Url(url) => formatter + .debug_tuple("Url") + .field(&redacted_url(url)) + .finish(), + Self::Host(host) => formatter + .debug_struct("Host") + .field("host", &host.host) + .field("port", &host.port) + .field("username", &host.username) + .field("password", &host.password.as_ref().map(|_| "")) + .field("tls", &host.tls) + .field("db", &host.db) + .finish(), + Self::Cluster(cluster) => formatter + .debug_struct("Cluster") + .field("nodes", &cluster.nodes) + .field("username", &cluster.username) + .field("password", &cluster.password.as_ref().map(|_| "")) + .field("tls", &cluster.tls) + .finish(), + } + } +} + +fn redacted_url(url: &str) -> String { + match url + .split_once("://") + .and_then(|(scheme, rest)| Some((scheme, rest.rsplit_once('@')?.1))) + { + Some((scheme, host)) => format!("{scheme}://@{host}"), + None => url.to_owned(), + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::{Endpoint, HostEndpoint}; + + #[rstest] + #[case::url_with_password(Endpoint::Url("redis://user:secret@cache:6379/0".to_owned()), "redis://@cache:6379/0")] + #[case::url_with_at_in_password(Endpoint::Url("rediss://:p@ss@cache:6380".to_owned()), "rediss://@cache:6380")] + #[case::host_password( + Endpoint::Host(HostEndpoint { + host: "cache".to_owned(), + port: 6379, + username: None, + password: Some("secret".to_owned()), + tls: false, + db: 0, + }), + "" + )] + fn debug_output_never_shows_the_password(#[case] endpoint: Endpoint, #[case] shown: &str) { + let debug = format!("{endpoint:?}"); + + assert!( + !debug.contains("secret") && !debug.contains("p@ss"), + "{debug}" + ); + assert!(debug.contains(shown), "{debug}"); + } +} diff --git a/litellm-rust/crates/redis-client/src/connection.rs b/litellm-rust/crates/redis-client/src/connection.rs new file mode 100644 index 00000000000..aaf26f105da --- /dev/null +++ b/litellm-rust/crates/redis-client/src/connection.rs @@ -0,0 +1,80 @@ +use std::time::Duration; + +use redis::{ + Cmd, ConnectionInfo, Pipeline, RedisFuture, RedisResult, Value, + aio::{ConnectionLike, ConnectionManager, ConnectionManagerConfig}, + cluster::{ClusterClientBuilder, ClusterConfig}, + cluster_async::ClusterConnection, +}; + +use crate::ClusterEndpoint; + +const RETRIES: u32 = 3; + +#[derive(Clone)] +pub(crate) enum Connection { + Standalone(ConnectionManager), + Cluster(ClusterConnection), +} + +impl Connection { + pub(crate) fn standalone(info: ConnectionInfo, timeout: Duration) -> RedisResult { + let config = ConnectionManagerConfig::new() + .set_connection_timeout(Some(timeout)) + .set_response_timeout(Some(timeout)) + .set_number_of_retries(RETRIES as usize); + let manager = ConnectionManager::new_lazy_with_config(redis::Client::open(info)?, config)?; + Ok(Self::Standalone(manager)) + } + + pub(crate) fn cluster(endpoint: &ClusterEndpoint, timeout: Duration) -> RedisResult { + let builder = ClusterClientBuilder::new(endpoint.addrs()) + .connection_timeout(timeout) + .response_timeout(timeout) + .retries(RETRIES); + let builder = match &endpoint.username { + Some(username) => builder.username(username), + None => builder, + }; + let builder = match &endpoint.password { + Some(password) => builder.password(password), + None => builder, + }; + let config = ClusterConfig::new() + .set_connection_timeout(timeout) + .set_response_timeout(timeout); + Ok(Self::Cluster( + builder + .build()? + .get_pending_async_connection_with_config(config), + )) + } +} + +impl ConnectionLike for Connection { + fn req_packed_command<'a>(&'a mut self, cmd: &'a Cmd) -> RedisFuture<'a, Value> { + match self { + Self::Standalone(connection) => connection.req_packed_command(cmd), + Self::Cluster(connection) => connection.req_packed_command(cmd), + } + } + + fn req_packed_commands<'a>( + &'a mut self, + pipeline: &'a Pipeline, + offset: usize, + count: usize, + ) -> RedisFuture<'a, Vec> { + match self { + Self::Standalone(connection) => connection.req_packed_commands(pipeline, offset, count), + Self::Cluster(connection) => connection.req_packed_commands(pipeline, offset, count), + } + } + + fn get_db(&self) -> i64 { + match self { + Self::Standalone(connection) => connection.get_db(), + Self::Cluster(connection) => connection.get_db(), + } + } +} diff --git a/litellm-rust/crates/redis-client/src/error.rs b/litellm-rust/crates/redis-client/src/error.rs new file mode 100644 index 00000000000..ed3b2c74004 --- /dev/null +++ b/litellm-rust/crates/redis-client/src/error.rs @@ -0,0 +1,69 @@ +use redis::{ErrorKind, RedisError, ServerErrorKind}; + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("Redis did not answer, so the command may or may not have applied")] + Transport(#[source] RedisError), + #[error("Redis rejected the command")] + Server(#[source] RedisError), + #[error("Redis replied with a value of a different type")] + Decode(#[source] RedisError), + #[error("the Redis circuit breaker is open, so the command was not sent")] + BreakerOpen, +} + +#[derive(Debug, thiserror::Error)] +#[error("connecting to Redis")] +pub struct ConnectError(#[from] RedisError); + +impl From for Error { + fn from(error: RedisError) -> Self { + match error.kind() { + ErrorKind::Server( + ServerErrorKind::BusyLoading + | ServerErrorKind::ClusterDown + | ServerErrorKind::MasterDown + | ServerErrorKind::TryAgain + | ServerErrorKind::ReadOnly, + ) => Self::Transport(error), + ErrorKind::Server(_) | ErrorKind::Extension => Self::Server(error), + ErrorKind::Parse | ErrorKind::UnexpectedReturnType => Self::Decode(error), + _ => Self::Transport(error), + } + } +} + +#[cfg(test)] +mod tests { + use redis::{RedisError, Value}; + use rstest::rstest; + + use super::Error; + + fn reply(resp: &str) -> RedisError { + match redis::parse_redis_value(resp.as_bytes()).unwrap() { + Value::ServerError(error) => error.into(), + other => panic!("{resp:?} parsed as {other:?}"), + } + } + + #[rstest] + #[case::wrong_type("-WRONGTYPE Operation against a key holding the wrong kind of value\r\n")] + #[case::no_script("-NOSCRIPT No matching script\r\n")] + #[case::cross_slot("-CROSSSLOT Keys in request don't hash to the same slot\r\n")] + #[case::generic("-ERR value is not an integer or out of range\r\n")] + #[case::no_perm("-NOPERM this user has no permissions to access one of the keys\r\n")] + fn rejections_are_server_errors(#[case] resp: &str) { + assert!(matches!(Error::from(reply(resp)), Error::Server(_))); + } + + #[rstest] + #[case::loading("-LOADING Redis is loading the dataset in memory\r\n")] + #[case::cluster_down("-CLUSTERDOWN The cluster is down\r\n")] + #[case::master_down("-MASTERDOWN Link with MASTER is down\r\n")] + #[case::try_again("-TRYAGAIN Multiple keys request during rehashing of slot\r\n")] + #[case::read_only("-READONLY You can't write against a read only replica.\r\n")] + fn unavailable_nodes_are_transport_errors(#[case] resp: &str) { + assert!(matches!(Error::from(reply(resp)), Error::Transport(_))); + } +} diff --git a/litellm-rust/crates/redis-client/src/lib.rs b/litellm-rust/crates/redis-client/src/lib.rs new file mode 100644 index 00000000000..47ff8677e06 --- /dev/null +++ b/litellm-rust/crates/redis-client/src/lib.rs @@ -0,0 +1,20 @@ +mod client; +mod commands; +mod config; +mod connection; +mod error; +mod namespace; +mod pipeline; +mod script; +#[cfg(test)] +mod test_support; + +pub use client::{RedisClient, Ttl}; +pub use config::{ + ClusterEndpoint, ClusterNode, DEFAULT_TIMEOUT, Endpoint, HostEndpoint, RedisConfig, +}; +pub use error::{ConnectError, Error}; +pub use namespace::Namespace; +pub use pipeline::Pipeline; +pub use redis::Value; +pub use script::Script; diff --git a/litellm-rust/crates/redis-client/src/namespace.rs b/litellm-rust/crates/redis-client/src/namespace.rs new file mode 100644 index 00000000000..4db02b5b7bb --- /dev/null +++ b/litellm-rust/crates/redis-client/src/namespace.rs @@ -0,0 +1,57 @@ +#[derive(Clone, Debug, Default, PartialEq)] +pub struct Namespace(Option); + +impl Namespace { + pub fn new(namespace: impl Into) -> Self { + let namespace = namespace.into(); + Self((!namespace.is_empty()).then_some(namespace)) + } + + pub fn as_str(&self) -> Option<&str> { + self.0.as_deref() + } + + pub(crate) fn apply(&self, key: &str) -> String { + match &self.0 { + Some(namespace) if !is_prefixed(key, namespace) => { + format!("{namespace}:{key}") + } + _ => key.to_owned(), + } + } +} + +fn is_prefixed(key: &str, namespace: &str) -> bool { + key.strip_prefix(namespace) + .is_some_and(|rest| rest.starts_with(':')) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::Namespace; + + #[rstest] + #[case::plain(Namespace::new("litellm"), "spend:key:abc", "litellm:spend:key:abc")] + #[case::no_namespace(Namespace::default(), "spend:key:abc", "spend:key:abc")] + #[case::empty_namespace(Namespace::new(""), "spend:key:abc", "spend:key:abc")] + #[case::already_prefixed(Namespace::new("litellm"), "litellm:abc", "litellm:abc")] + #[case::shares_prefix_without_delimiter( + Namespace::new("litellm"), + "litellm_spend_update_buffer", + "litellm:litellm_spend_update_buffer" + )] + #[case::hash_tagged_key( + Namespace::new("litellm"), + "{api_key:sk-1}:requests", + "litellm:{api_key:sk-1}:requests" + )] + fn prefixes_keys_like_python_check_and_fix_namespace( + #[case] namespace: Namespace, + #[case] key: &str, + #[case] expected: &str, + ) { + assert_eq!(namespace.apply(key), expected); + } +} diff --git a/litellm-rust/crates/redis-client/src/pipeline.rs b/litellm-rust/crates/redis-client/src/pipeline.rs new file mode 100644 index 00000000000..e54c65e9605 --- /dev/null +++ b/litellm-rust/crates/redis-client/src/pipeline.rs @@ -0,0 +1,292 @@ +use std::{collections::HashMap, time::Duration}; + +use futures_util::future::join_all; +use redis::{ + Cmd, ErrorKind, RedisResult, ServerErrorKind, ToRedisArgs, ToSingleRedisArg, Value, + cluster_routing::Slot, +}; + +use crate::{Error, Namespace, Script, commands, connection::Connection}; + +pub struct Pipeline { + namespace: Namespace, + ops: Vec, +} + +struct Op { + command: Cmd, + key: String, + script: Option, +} + +struct ScriptCall { + script: Script, + keys: Vec, + args: Vec>, +} + +type Settled = (usize, Result); + +impl Pipeline { + pub(crate) fn new(namespace: Namespace) -> Self { + Self { + namespace, + ops: Vec::new(), + } + } + + pub fn get(self, key: &str) -> Self { + let key = self.namespace.apply(key); + self.command(commands::get(key.clone()), key) + } + + pub fn set(self, key: &str, value: impl ToSingleRedisArg, ttl: Duration) -> Self { + let key = self.namespace.apply(key); + self.command(commands::set(key.clone(), value, ttl), key) + } + + pub fn del(self, key: &str) -> Self { + let key = self.namespace.apply(key); + self.command(commands::del(vec![key.clone()]), key) + } + + pub fn expire(self, key: &str, ttl: Duration) -> Self { + let key = self.namespace.apply(key); + self.command(commands::expire(key.clone(), ttl), key) + } + + pub fn incr_by(self, key: &str, delta: i64) -> Self { + let key = self.namespace.apply(key); + self.command(commands::incr_by(key.clone(), delta), key) + } + + pub fn incr_by_float(self, key: &str, delta: f64) -> Self { + let key = self.namespace.apply(key); + self.command(commands::incr_by_float(key.clone(), delta), key) + } + + pub fn eval(self, script: &Script, keys: &[&str], args: impl ToRedisArgs) -> Self { + let keys: Vec = keys.iter().map(|key| self.namespace.apply(key)).collect(); + let args = args.to_redis_args(); + let op = Op { + command: script.evalsha(&keys, &args), + key: keys.first().cloned().unwrap_or_default(), + script: Some(ScriptCall { + script: script.clone(), + keys, + args, + }), + }; + self.push(op) + } + + fn command(self, command: Cmd, key: String) -> Self { + self.push(Op { + command, + key, + script: None, + }) + } + + fn push(mut self, op: Op) -> Self { + self.ops.push(op); + self + } + + pub(crate) fn len(&self) -> usize { + self.ops.len() + } + + pub(crate) async fn run(&self, connection: &Connection) -> Vec> { + let groups = self.groups(connection); + let settled = join_all( + groups + .iter() + .map(|group| self.run_group(connection.clone(), group)), + ) + .await; + let mut settled: Vec = settled.into_iter().flatten().collect(); + settled.sort_by_key(|(index, _)| *index); + settled.into_iter().map(|(_, result)| result).collect() + } + + fn groups(&self, connection: &Connection) -> Vec> { + match connection { + _ if self.ops.is_empty() => Vec::new(), + Connection::Standalone(_) => vec![(0..self.ops.len()).collect()], + Connection::Cluster(_) => { + let mut by_slot: HashMap> = HashMap::new(); + for (index, op) in self.ops.iter().enumerate() { + by_slot + .entry(Slot::for_key(&op.key)) + .or_default() + .push(index); + } + by_slot.into_values().collect() + } + } + } + + async fn run_group(&self, mut connection: Connection, indices: &[usize]) -> Vec { + let mut pipeline = redis::pipe(); + pipeline.ignore_errors(); + for &index in indices { + pipeline.add_command(self.ops[index].command.clone()); + } + let replies: RedisResult>> = + pipeline.query_async(&mut connection).await; + let replies = match replies { + Ok(replies) => replies, + Err(error) => { + return indices + .iter() + .map(|&index| (index, Err(Error::from(error.clone())))) + .collect(); + } + }; + let mut settled = Vec::with_capacity(indices.len()); + for (&index, reply) in indices.iter().zip(replies) { + settled.push((index, self.settle(&mut connection, index, reply).await)); + } + settled + } + + async fn settle( + &self, + connection: &mut Connection, + index: usize, + reply: RedisResult, + ) -> Result { + match (reply, &self.ops[index].script) { + (Err(error), Some(call)) + if error.kind() == ErrorKind::Server(ServerErrorKind::NoScript) => + { + Ok(call + .script + .invocation(&call.keys, &call.args) + .invoke_async(connection) + .await?) + } + (reply, _) => Ok(reply?), + } + } +} + +#[cfg(test)] +mod tests { + use std::time::{Duration, Instant}; + + use redis::Value; + use rstest::rstest; + + use crate::{ + DEFAULT_TIMEOUT, Error, Script, + test_support::{Server, Topology, flush_scripts}, + }; + + #[rstest] + #[tokio::test] + async fn each_op_settles_on_its_own_in_input_order( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + ) { + let server = Server::start(topology).await; + let client = server.client("", DEFAULT_TIMEOUT).await; + client + .set("text", "not a number", Duration::from_secs(60)) + .await + .unwrap(); + let pipeline = client + .pipeline() + .incr_by("a", 1) + .incr_by("text", 1) + .incr_by("c", 5) + .set("b", "2", Duration::from_secs(60)) + .get("b") + .incr_by("a", 1) + .get("missing"); + + let results = client.run(&pipeline).await; + + assert_eq!(results.len(), 7); + assert_eq!(results[0].as_ref().unwrap(), &Value::Int(1)); + assert!( + matches!(results[1], Err(Error::Server(_))), + "{:?}", + results[1] + ); + assert_eq!(results[2].as_ref().unwrap(), &Value::Int(5)); + assert_eq!(results[3].as_ref().unwrap(), &Value::Okay); + assert_eq!( + results[4].as_ref().unwrap(), + &Value::BulkString(b"2".to_vec()) + ); + assert_eq!(results[5].as_ref().unwrap(), &Value::Int(2)); + assert_eq!(results[6].as_ref().unwrap(), &Value::Nil); + } + + #[rstest] + #[tokio::test] + async fn a_flushed_script_in_a_pipeline_reruns_alone( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + ) { + let server = Server::start(topology).await; + let client = server.client("ns", DEFAULT_TIMEOUT).await; + let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])"); + client.eval::(&incr, &["counter"], 1).await.unwrap(); + flush_scripts(&client).await; + let pipeline = client + .pipeline() + .eval(&incr, &["counter"], 2) + .incr_by("other", 7); + + let results = client.run(&pipeline).await; + + assert_eq!(results[0].as_ref().unwrap(), &Value::Int(3)); + assert_eq!(results[1].as_ref().unwrap(), &Value::Int(7)); + assert_eq!(client.get::("counter").await.unwrap(), Some(3)); + } + + #[rstest] + #[tokio::test] + async fn an_unreachable_server_fails_every_op_as_transport( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + ) { + let server = Server::start(topology).await; + let client = server.client("", Duration::from_millis(500)).await; + let pipeline = client.pipeline().incr_by("a", 1).incr_by("c", 1); + + server.freeze().await; + let results = client.run(&pipeline).await; + server.thaw().await; + + assert_eq!(results.len(), 2); + assert!( + matches!(results[0], Err(Error::Transport(_))), + "{:?}", + results[0] + ); + assert!( + matches!(results[1], Err(Error::Transport(_))), + "{:?}", + results[1] + ); + } + + #[rstest] + #[tokio::test] + async fn an_empty_pipeline_returns_without_asking_redis( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + ) { + let server = Server::start(topology).await; + let client = server.client("", Duration::from_secs(2)).await; + + server.freeze().await; + let started = Instant::now(); + let results = client.run(&client.pipeline()).await; + let elapsed = started.elapsed(); + server.thaw().await; + + assert!(results.is_empty()); + assert!(elapsed < Duration::from_secs(1), "{elapsed:?}"); + } +} diff --git a/litellm-rust/crates/redis-client/src/script.rs b/litellm-rust/crates/redis-client/src/script.rs new file mode 100644 index 00000000000..c5d73133ecc --- /dev/null +++ b/litellm-rust/crates/redis-client/src/script.rs @@ -0,0 +1,88 @@ +use redis::{Cmd, ScriptInvocation}; + +#[derive(Clone)] +pub struct Script(redis::Script); + +impl Script { + pub fn new(source: &str) -> Self { + Self(redis::Script::new(source)) + } + + pub(crate) fn evalsha(&self, keys: &[String], args: &[Vec]) -> Cmd { + redis::cmd("EVALSHA") + .arg(self.0.get_hash()) + .arg(keys.len()) + .arg(keys) + .arg(args) + .to_owned() + } + + pub(crate) fn invocation(&self, keys: &[String], args: &[Vec]) -> ScriptInvocation<'_> { + let mut invocation = self.0.prepare_invoke(); + invocation.key(keys).arg(args); + invocation + } +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use rstest::rstest; + + use super::Script; + use crate::{ + DEFAULT_TIMEOUT, Error, + test_support::{Server, Topology, flush_scripts}, + }; + + #[rstest] + #[tokio::test] + async fn script_keys_are_namespaced( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + ) { + let server = Server::start(topology).await; + let client = server.client("a", DEFAULT_TIMEOUT).await; + let bare = server.client("", DEFAULT_TIMEOUT).await; + let read = Script::new("return redis.call('GET', KEYS[1])"); + bare.set("a:x", "1", Duration::from_secs(60)).await.unwrap(); + + let value: Option = client + .eval(&read, &["x"], Vec::::new()) + .await + .unwrap(); + + assert_eq!(value.as_deref(), Some("1")); + } + + #[rstest] + #[tokio::test] + async fn a_failing_script_is_a_server_error_and_keeps_the_connection( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + ) { + let server = Server::start(topology).await; + let client = server.client("", DEFAULT_TIMEOUT).await; + let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])"); + client.rpush("list", "item").await.unwrap(); + + let failed = client.eval::(&incr, &["list"], 1).await; + + assert!(matches!(failed, Err(Error::Server(_))), "{failed:?}"); + assert_eq!(client.eval::(&incr, &["counter"], 3).await.unwrap(), 3); + } + + #[rstest] + #[tokio::test] + async fn a_flushed_script_is_loaded_again( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + ) { + let server = Server::start(topology).await; + let client = server.client("", DEFAULT_TIMEOUT).await; + let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])"); + client.eval::(&incr, &["counter"], 1).await.unwrap(); + + flush_scripts(&client).await; + + assert_eq!(client.eval::(&incr, &["counter"], 1).await.unwrap(), 2); + } +} diff --git a/litellm-rust/crates/redis-client/src/test_support.rs b/litellm-rust/crates/redis-client/src/test_support.rs new file mode 100644 index 00000000000..64a174b1543 --- /dev/null +++ b/litellm-rust/crates/redis-client/src/test_support.rs @@ -0,0 +1,125 @@ +use std::time::Duration; + +use litellm_circuit_breaker::Breaker; +use litellm_redis_testing::{RedisCluster, RedisServer}; +use redis::cluster_routing::{Route, RoutingInfo, SingleNodeRoutingInfo, SlotAddr}; + +use crate::{ + ClusterEndpoint, ClusterNode, Endpoint, Namespace, RedisClient, RedisConfig, + connection::Connection, +}; + +#[derive(Clone, Copy, Debug)] +pub(crate) enum Topology { + Standalone, + Cluster, +} + +pub(crate) enum Server { + Standalone(RedisServer), + Cluster(RedisCluster), +} + +impl Server { + pub(crate) async fn start(topology: Topology) -> Self { + match topology { + Topology::Standalone => Self::Standalone(RedisServer::start().await.unwrap()), + Topology::Cluster => Self::Cluster(RedisCluster::start().await.unwrap()), + } + } + + pub(crate) fn endpoint(&self) -> Endpoint { + match self { + Self::Standalone(server) => Endpoint::Url(server.url().to_owned()), + Self::Cluster(cluster) => Endpoint::Cluster(ClusterEndpoint { + nodes: cluster + .nodes() + .map(|(host, port)| ClusterNode { + host: host.to_owned(), + port, + }) + .collect(), + username: None, + password: None, + tls: false, + }), + } + } + + pub(crate) async fn client(&self, namespace: &str, timeout: Duration) -> RedisClient { + self.guarded_client(namespace, timeout, None).await + } + + pub(crate) async fn guarded_client( + &self, + namespace: &str, + timeout: Duration, + breaker: Option, + ) -> RedisClient { + RedisClient::new(RedisConfig { + endpoint: self.endpoint(), + namespace: Namespace::new(namespace), + timeout, + breaker, + }) + .unwrap() + } + + pub(crate) async fn restart(&self) { + match self { + Self::Standalone(server) => server.restart().await.unwrap(), + Self::Cluster(cluster) => cluster.restart().await.unwrap(), + } + } + + pub(crate) async fn stop(&self) { + match self { + Self::Standalone(server) => server.stop().await.unwrap(), + Self::Cluster(cluster) => cluster.stop().await.unwrap(), + } + } + + pub(crate) async fn resume(&self) { + match self { + Self::Standalone(server) => server.resume().await.unwrap(), + Self::Cluster(cluster) => cluster.resume().await.unwrap(), + } + } + + pub(crate) async fn freeze(&self) { + match self { + Self::Standalone(server) => server.freeze().await.unwrap(), + Self::Cluster(cluster) => cluster.freeze().await.unwrap(), + } + } + + pub(crate) async fn thaw(&self) { + match self { + Self::Standalone(server) => server.thaw().await.unwrap(), + Self::Cluster(cluster) => cluster.thaw().await.unwrap(), + } + } +} + +pub(crate) async fn client_id(client: &RedisClient, key: &str) -> i64 { + let command = redis::cmd("CLIENT").arg("ID").to_owned(); + match client.connection.clone() { + Connection::Standalone(mut connection) => { + command.query_async(&mut connection).await.unwrap() + } + Connection::Cluster(mut connection) => { + let node = Route::with_key(key, SlotAddr::Master); + let routing = RoutingInfo::SingleNode(SingleNodeRoutingInfo::SpecificNode(node)); + redis::from_redis_value(connection.route_command(command, routing).await.unwrap()) + .unwrap() + } + } +} + +pub(crate) async fn flush_scripts(client: &RedisClient) { + redis::cmd("SCRIPT") + .arg("FLUSH") + .query_async::<()>(&mut client.connection.clone()) + .await + .unwrap(); +} diff --git a/litellm-rust/crates/redis-testing/AGENTS.md b/litellm-rust/crates/redis-testing/AGENTS.md new file mode 100644 index 00000000000..fb7a4cd737a --- /dev/null +++ b/litellm-rust/crates/redis-testing/AGENTS.md @@ -0,0 +1,6 @@ +- Dev-only crate that starts real Redis for tests. Never a dependency outside `[dev-dependencies]` +- Same shape as `db-testing`: one container per test, a handle that removes it on drop +- The host port is fixed at start so `restart` keeps the same URL and a client can reconnect to it +- `RedisCluster` runs 3 masters in one container. Each node listens on the same port inside and outside and announces 127.0.0.1, so MOVED redirects work from the test process +- Outages come from the container (`restart`, `freeze` / `thaw`), never `DEBUG`, which the official image disables +- No env gate that makes a test pass without Redis. A test that can't start its container fails diff --git a/litellm-rust/crates/redis-testing/Cargo.toml b/litellm-rust/crates/redis-testing/Cargo.toml new file mode 100644 index 00000000000..c9bad227acc --- /dev/null +++ b/litellm-rust/crates/redis-testing/Cargo.toml @@ -0,0 +1,11 @@ +[package] +name = "litellm-redis-testing" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +testcontainers-modules = { workspace = true, features = ["redis"] } +thiserror.workspace = true +tokio = { workspace = true, features = ["io-util"] } diff --git a/litellm-rust/crates/redis-testing/src/cluster.rs b/litellm-rust/crates/redis-testing/src/cluster.rs new file mode 100644 index 00000000000..e61f7832b1c --- /dev/null +++ b/litellm-rust/crates/redis-testing/src/cluster.rs @@ -0,0 +1,93 @@ +use testcontainers_modules::testcontainers::{ + ContainerAsync, GenericImage, ImageExt, + core::{IntoContainerPort, WaitFor}, + runners::AsyncRunner, +}; + +use crate::{ + Error, + ports::{HOST, start_on_free_ports, wait_for_pong}, + standalone::REDIS_TAG, +}; + +const READY: &str = "redis-testing cluster ready"; + +const START_CLUSTER: &str = r#" +set -e +bus=17000 +for port in "$@"; do + mkdir -p "/data/$port" + redis-server --port "$port" --cluster-enabled yes --cluster-port "$bus" \ + --cluster-config-file "/data/$port/nodes.conf" --cluster-announce-ip 127.0.0.1 \ + --dir "/data/$port" --save '' --appendonly no --daemonize yes + bus=$((bus + 1)) +done +for port in "$@"; do + until redis-cli -p "$port" ping >/dev/null 2>&1; do sleep 0.1; done +done +if [ ! -f /data/created ]; then + redis-cli --cluster create $(for port in "$@"; do printf '127.0.0.1:%s ' "$port"; done) --cluster-yes >/dev/null + touch /data/created +fi +for port in "$@"; do + until redis-cli -p "$port" cluster info | grep -q cluster_state:ok; do sleep 0.1; done +done +echo "redis-testing cluster ready" +exec sleep infinity +"#; + +pub struct RedisCluster { + container: ContainerAsync, + ports: [u16; 3], +} + +impl RedisCluster { + pub async fn start() -> Result { + start_on_free_ports(Self::start_on).await + } + + async fn start_on(ports: [u16; 3]) -> Result { + let image = GenericImage::new("redis", REDIS_TAG) + .with_wait_for(WaitFor::message_on_stdout(READY)) + .with_entrypoint("sh"); + let request = ports.iter().fold( + image.with_cmd( + ["-c", START_CLUSTER, "sh"] + .into_iter() + .map(str::to_owned) + .chain(ports.iter().map(u16::to_string)), + ), + |request, &port| request.with_mapped_port(port, port.tcp()), + ); + let container = request.start().await?; + for port in ports { + wait_for_pong(port).await?; + } + Ok(Self { container, ports }) + } + + pub fn nodes(&self) -> impl Iterator + '_ { + self.ports.iter().map(|&port| (HOST, port)) + } + + pub async fn restart(&self) -> Result<(), Error> { + self.stop().await?; + self.resume().await + } + + pub async fn stop(&self) -> Result<(), Error> { + Ok(self.container.stop().await?) + } + + pub async fn resume(&self) -> Result<(), Error> { + Ok(self.container.start().await?) + } + + pub async fn freeze(&self) -> Result<(), Error> { + Ok(self.container.pause().await?) + } + + pub async fn thaw(&self) -> Result<(), Error> { + Ok(self.container.unpause().await?) + } +} diff --git a/litellm-rust/crates/redis-testing/src/error.rs b/litellm-rust/crates/redis-testing/src/error.rs new file mode 100644 index 00000000000..f21cc959812 --- /dev/null +++ b/litellm-rust/crates/redis-testing/src/error.rs @@ -0,0 +1,11 @@ +use std::io; + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("reserving a host port for the Redis container")] + Port(#[source] io::Error), + #[error("Redis on host port {port} did not answer PING")] + Unreachable { port: u16 }, + #[error("controlling the Redis container")] + Container(#[from] testcontainers_modules::testcontainers::TestcontainersError), +} diff --git a/litellm-rust/crates/redis-testing/src/lib.rs b/litellm-rust/crates/redis-testing/src/lib.rs new file mode 100644 index 00000000000..6c2afd7f275 --- /dev/null +++ b/litellm-rust/crates/redis-testing/src/lib.rs @@ -0,0 +1,8 @@ +mod cluster; +mod error; +mod ports; +mod standalone; + +pub use cluster::RedisCluster; +pub use error::Error; +pub use standalone::RedisServer; diff --git a/litellm-rust/crates/redis-testing/src/ports.rs b/litellm-rust/crates/redis-testing/src/ports.rs new file mode 100644 index 00000000000..e7b94988c10 --- /dev/null +++ b/litellm-rust/crates/redis-testing/src/ports.rs @@ -0,0 +1,66 @@ +use std::{ + future::Future, + io, + net::TcpListener, + time::{Duration, Instant}, +}; + +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::TcpStream, +}; + +use crate::Error; + +pub(crate) const HOST: &str = "127.0.0.1"; +const READY_TIMEOUT: Duration = Duration::from_secs(10); +const START_ATTEMPTS: usize = 3; + +fn free_ports() -> Result<[u16; N], Error> { + let listeners = (0..N) + .map(|_| TcpListener::bind((HOST, 0))) + .collect::>>() + .map_err(Error::Port)?; + let ports = listeners + .iter() + .map(|listener| listener.local_addr().map(|address| address.port())) + .collect::>>() + .map_err(Error::Port)?; + Ok(ports + .try_into() + .unwrap_or_else(|ports: Vec| unreachable!("{} ports", ports.len()))) +} + +pub(crate) async fn wait_for_pong(port: u16) -> Result<(), Error> { + let deadline = Instant::now() + READY_TIMEOUT; + while Instant::now() < deadline { + if ping(port).await.unwrap_or(false) { + return Ok(()); + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + Err(Error::Unreachable { port }) +} + +async fn ping(port: u16) -> io::Result { + let mut stream = TcpStream::connect((HOST, port)).await?; + stream.write_all(b"PING\r\n").await?; + let mut reply = [0; 7]; + stream.read_exact(&mut reply).await?; + Ok(&reply == b"+PONG\r\n") +} + +pub(crate) async fn start_on_free_ports( + start: impl Fn([u16; N]) -> F, +) -> Result +where + F: Future>, +{ + let mut attempt = 1; + loop { + match start(free_ports()?).await { + Err(_) if attempt < START_ATTEMPTS => attempt += 1, + started => return started, + } + } +} diff --git a/litellm-rust/crates/redis-testing/src/standalone.rs b/litellm-rust/crates/redis-testing/src/standalone.rs new file mode 100644 index 00000000000..8060ccbf190 --- /dev/null +++ b/litellm-rust/crates/redis-testing/src/standalone.rs @@ -0,0 +1,62 @@ +use testcontainers_modules::{ + redis::{REDIS_PORT, Redis}, + testcontainers::core::IntoContainerPort, + testcontainers::{ContainerAsync, ImageExt, runners::AsyncRunner}, +}; + +use crate::{ + Error, + ports::{HOST, start_on_free_ports, wait_for_pong}, +}; + +pub(crate) const REDIS_TAG: &str = + "7.4@sha256:44960b3100c48159a66d28c18b0588c193492e25f7551d60e51dd464fc3eca5b"; + +pub struct RedisServer { + container: ContainerAsync, + url: String, +} + +impl RedisServer { + pub async fn start() -> Result { + start_on_free_ports(Self::start_on).await + } + + async fn start_on([host_port]: [u16; 1]) -> Result { + let container = Redis::default() + .with_tag(REDIS_TAG) + .with_mapped_port(host_port, REDIS_PORT.tcp()) + .start() + .await?; + wait_for_pong(host_port).await?; + Ok(Self { + container, + url: format!("redis://{HOST}:{host_port}"), + }) + } + + pub fn url(&self) -> &str { + &self.url + } + + pub async fn restart(&self) -> Result<(), Error> { + self.stop().await?; + self.resume().await + } + + pub async fn stop(&self) -> Result<(), Error> { + Ok(self.container.stop().await?) + } + + pub async fn resume(&self) -> Result<(), Error> { + Ok(self.container.start().await?) + } + + pub async fn freeze(&self) -> Result<(), Error> { + Ok(self.container.pause().await?) + } + + pub async fn thaw(&self) -> Result<(), Error> { + Ok(self.container.unpause().await?) + } +} diff --git a/litellm-rust/crates/spend-clickhouse/Cargo.toml b/litellm-rust/crates/spend-clickhouse/Cargo.toml index 661a75ac02e..893c5daded3 100644 --- a/litellm-rust/crates/spend-clickhouse/Cargo.toml +++ b/litellm-rust/crates/spend-clickhouse/Cargo.toml @@ -19,7 +19,7 @@ time = { workspace = true, features = ["formatting"] } [dev-dependencies] litellm-http = { workspace = true, features = ["test-support"] } rstest.workspace = true -testcontainers-modules = { version = "0.15.0", features = ["clickhouse"] } +testcontainers-modules = { workspace = true, features = ["clickhouse"] } tokio.workspace = true uuid.workspace = true wiremock.workspace = true diff --git a/litellm-rust/crates/storage-clickhouse/Cargo.toml b/litellm-rust/crates/storage-clickhouse/Cargo.toml index 4e72e823dc2..39d04f2de5d 100644 --- a/litellm-rust/crates/storage-clickhouse/Cargo.toml +++ b/litellm-rust/crates/storage-clickhouse/Cargo.toml @@ -19,6 +19,6 @@ url.workspace = true [dev-dependencies] litellm-http = { workspace = true, features = ["test-support"] } rstest.workspace = true -testcontainers-modules = { version = "0.15.0", features = ["clickhouse"] } +testcontainers-modules = { workspace = true, features = ["clickhouse"] } tokio.workspace = true wiremock.workspace = true diff --git a/litellm/__init__.py b/litellm/__init__.py index c0141f43611..117ba9604e9 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1785,7 +1785,7 @@ if TYPE_CHECKING: from .llms.anthropic.pass_through.messages.transformation import ( AnthropicMessagesConfig as AnthropicMessagesConfig, ) - from .llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( + from .llms.bedrock.messages.invoke_transformations.anthropic_claude_transformation import ( AmazonAnthropicClaudeMessagesConfig as AmazonAnthropicClaudeMessagesConfig, ) from .llms.bedrock.messages.mantle_transformation import ( @@ -1808,7 +1808,7 @@ if TYPE_CHECKING: from .llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import ( VertexAIAnthropicConfig as VertexAIAnthropicConfig, ) - from .llms.vertex_ai.vertex_ai_partner_models.llama3.transformation import ( + from .llms.vertex_ai.vertex_ai_partner_models.llama.transformation import ( VertexAILlama3Config as VertexAILlama3Config, ) from .llms.vertex_ai.vertex_ai_partner_models.ai21.transformation import ( @@ -1835,10 +1835,10 @@ if TYPE_CHECKING: from .llms.bedrock.chat.invoke_transformations.amazon_qwen3_transformation import ( AmazonQwen3Config as AmazonQwen3Config, ) - from .llms.bedrock.chat.invoke_transformations.anthropic_claude2_transformation import ( + from .llms.bedrock.chat.invoke_transformations.anthropic_legacy_text_completions_transformation import ( AmazonAnthropicConfig as AmazonAnthropicConfig, ) - from .llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import ( + from .llms.bedrock.chat.invoke_transformations.anthropic_claude_transformation import ( AmazonAnthropicClaudeConfig as AmazonAnthropicClaudeConfig, ) from .llms.bedrock.chat.invoke_transformations.amazon_cohere_transformation import ( @@ -2000,7 +2000,7 @@ if TYPE_CHECKING: GradientAIConfig as GradientAIConfig, ) from .llms.openai.chat.gpt_transformation import OpenAIGPTConfig as OpenAIGPTConfig - from .llms.openai.chat.gpt_5_transformation import ( + from .llms.openai.chat.gpt_reasoning_transformation import ( OpenAIGPT5Config as OpenAIGPT5Config, ) from .llms.openai.transcriptions.whisper_transformation import ( @@ -2141,7 +2141,7 @@ if TYPE_CHECKING: from .llms.azure.chat.gpt_transformation import ( AzureOpenAIConfig as AzureOpenAIConfig, ) - from .llms.azure.chat.gpt_5_transformation import ( + from .llms.azure.chat.gpt_reasoning_transformation import ( AzureOpenAIGPT5Config as AzureOpenAIGPT5Config, ) from .llms.azure.completion.transformation import ( diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index e0cb7c680bc..a56004d1274 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -805,7 +805,7 @@ _LLM_CONFIGS_IMPORT_MAP: Final = { "BedrockClaudePlatformMessagesConfig", ), "AmazonAnthropicClaudeMessagesConfig": ( - ".llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation", + ".llms.bedrock.messages.invoke_transformations.anthropic_claude_transformation", "AmazonAnthropicClaudeMessagesConfig", ), "AmazonMantleMessagesConfig": ( @@ -835,7 +835,7 @@ _LLM_CONFIGS_IMPORT_MAP: Final = { "VertexAIAnthropicConfig", ), "VertexAILlama3Config": ( - ".llms.vertex_ai.vertex_ai_partner_models.llama3.transformation", + ".llms.vertex_ai.vertex_ai_partner_models.llama.transformation", "VertexAILlama3Config", ), "VertexAIAi21Config": ( @@ -880,11 +880,11 @@ _LLM_CONFIGS_IMPORT_MAP: Final = { "GoogleAIStudioGeminiConfig", ), # Alias "AmazonAnthropicConfig": ( - ".llms.bedrock.chat.invoke_transformations.anthropic_claude2_transformation", + ".llms.bedrock.chat.invoke_transformations.anthropic_legacy_text_completions_transformation", "AmazonAnthropicConfig", ), "AmazonAnthropicClaudeConfig": ( - ".llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation", + ".llms.bedrock.chat.invoke_transformations.anthropic_claude_transformation", "AmazonAnthropicClaudeConfig", ), "AmazonCohereConfig": ( @@ -1101,7 +1101,7 @@ _LLM_CONFIGS_IMPORT_MAP: Final = { "OpenAIOSeriesConfig", ), # Alias "OpenAIGPTConfig": (".llms.openai.chat.gpt_transformation", "OpenAIGPTConfig"), - "OpenAIGPT5Config": (".llms.openai.chat.gpt_5_transformation", "OpenAIGPT5Config"), + "OpenAIGPT5Config": (".llms.openai.chat.gpt_reasoning_transformation", "OpenAIGPT5Config"), "OpenAIWhisperAudioTranscriptionConfig": ( ".llms.openai.transcriptions.whisper_transformation", "OpenAIWhisperAudioTranscriptionConfig", @@ -1175,7 +1175,7 @@ _LLM_CONFIGS_IMPORT_MAP: Final = { "CometAPIConfig": (".llms.cometapi.chat.transformation", "CometAPIConfig"), "AzureOpenAIConfig": (".llms.azure.chat.gpt_transformation", "AzureOpenAIConfig"), "AzureOpenAIGPT5Config": ( - ".llms.azure.chat.gpt_5_transformation", + ".llms.azure.chat.gpt_reasoning_transformation", "AzureOpenAIGPT5Config", ), "AzureOpenAITextConfig": ( diff --git a/litellm/constants.py b/litellm/constants.py index d1daf2ee2d1..2901b7d8948 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -627,7 +627,7 @@ BEDROCK_MIN_THINKING_BUDGET_TOKENS: Final = get_env_int("BEDROCK_MIN_THINKING_BU # Anthropic's Messages API rejects thinking.budget_tokens < 1024. ANTHROPIC_MIN_THINKING_BUDGET_TOKENS: Final = 1024 REPLICATE_POLLING_DELAY_SECONDS: Final = get_env_float("REPLICATE_POLLING_DELAY_SECONDS", 0.5) -DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS: Final = get_env_int("DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS", 4096) +DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS: Final = get_env_int("DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS", 128000) DEFAULT_OCI_CHAT_MAX_TOKENS: Final = 4096 TOGETHER_AI_4_B: Final = get_env_int("TOGETHER_AI_4_B", 4) TOGETHER_AI_8_B: Final = get_env_int("TOGETHER_AI_8_B", 8) diff --git a/litellm/decisions/call.py b/litellm/decisions/call.py index abbeccc10b3..3b73f65c917 100644 --- a/litellm/decisions/call.py +++ b/litellm/decisions/call.py @@ -456,10 +456,14 @@ def prepare_call( extra_headers: Mapping[str, str] | None, kwargs: Mapping[str, object], ) -> DecisionsCall: + base_url_alias: Final = kwargs.get("base_url") + call_api_base: Final = ( + base_url_alias if api_base is None and isinstance(base_url_alias, str) and base_url_alias else api_base + ) upstream_model, provider, dynamic_api_key, dynamic_api_base = litellm.get_llm_provider( model=model, custom_llm_provider=custom_llm_provider, - api_base=api_base, + api_base=call_api_base, api_key=api_key, ) provider_config: Final = _provider_config(upstream_model, provider) @@ -488,7 +492,7 @@ def prepare_call( try: resolved_api_base, static_api_key = provider_config.resolve_connection( - dynamic_api_base or api_base, dynamic_api_key or api_key + dynamic_api_base or call_api_base, dynamic_api_key or api_key ) except ValueError as error: raise litellm.BadRequestError(message=str(error), model=model, llm_provider=provider) from error diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 19d161b8aa4..bc8344fa686 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -14,6 +14,7 @@ from functools import partial from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, Protocol, TypeAlias, TypeVar, cast +import openai from pydantic import BaseModel, Field, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict, assert_never @@ -22,11 +23,13 @@ from litellm._internal_context import with_service_target from litellm._logging import print_verbose, verbose_logger from litellm.constants import ( AUTH_OBJECTS_TARGET, + LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL, PROMETHEUS_OVERFLOW_SERIES_LABEL_VALUE, PROXY_LLM_PROVIDER_FALLBACK, PROXY_REJECTED_BEFORE_ROUTING_KEY, ) from litellm.exceptions import ( + CallerCredentialAuthenticationError, MalformedRequestFieldError, validate_rate_limit_category, validate_rate_limit_type, @@ -51,6 +54,7 @@ from litellm.litellm_core_utils.get_llm_provider_logic import is_registered_cust from litellm.litellm_core_utils.service_tier_utils import ( get_service_tier_from_standard_logging_payload, ) +from litellm.llms.github_copilot.per_user_auth import github_copilot_user_session_from from litellm.llms.openai_like.json_loader import JSONProviderRegistry from litellm.models.end_user import LiteLLM_EndUserTable from litellm.models.team import LiteLLM_TeamTable @@ -269,7 +273,10 @@ _NO_LOGGED_FIELDS: Final[Mapping[str, object]] = MappingProxyType({}) def _logged_fields(value: object) -> Mapping[str, object]: - return _LOGGED_FIELDS.validate_python(value) if isinstance(value, Mapping) else _NO_LOGGED_FIELDS + try: + return _LOGGED_FIELDS.validate_python(value) + except ValidationError: + return _NO_LOGGED_FIELDS def _is_llm_provider(name: str) -> bool: @@ -359,6 +366,87 @@ def _get_proxy_llm_router() -> Router | None: return llm_router +_CALLER_CREDENTIAL_FIELDS: Final = frozenset( + { + "aws_access_key_id", + "aws_secret_access_key", + "aws_session_token", + "aws_web_identity_token", + "azure_ad_token", + "azure_password", + "azure_scope", + "azure_username", + "client_id", + "client_secret", + "gigachat_access_token", + "litellm_credential_name", + "oci_fingerprint", + "oci_key", + "oci_key_file", + "oci_tenancy", + "oci_user", + "tenant_id", + "token", + "user_config", + "vertex_ai_credentials", + "vertex_credentials", + "watsonx_credentials", + "wx_credentials", + "zen_api_key", + } +) + + +_CALLER_ENDPOINT_FIELDS: Final = frozenset( + {"aws_bedrock_runtime_endpoint", "aws_sts_endpoint", "deployment_url", "s3_endpoint_url", "sagemaker_base_url"} +) + + +_CALLER_TRANSPORT_FIELDS: Final = frozenset( + {"api_key", "extra_headers", "headers", "organization", "provider_specific_header"} +) + + +def _sent_a_provider_credential_or_endpoint(caller_body: Mapping[str, object]) -> bool: + """True when the caller's request body, or its ``extra_body``, carries a provider credential field or a + provider endpoint field.""" + sent: Final = {*caller_body, *_logged_fields(caller_body.get("extra_body"))} + return not (_CALLER_CREDENTIAL_FIELDS.isdisjoint(sent) and _CALLER_ENDPOINT_FIELDS.isdisjoint(sent)) + + +def _used_a_caller_credential(logged_call: Mapping[str, object]) -> bool: + """True when the provider was called with a credential or at an endpoint that belongs to the caller rather than + the deployment: an api_key, header, OpenAI organization or other provider credential or endpoint field the caller + sent with the request, a stored credential the caller picked by name, the caller's own Anthropic OAuth token, the caller's own GitHub Copilot connection, or the + caller's identity token that a token-exchange deployment trades in on every call.""" + litellm_params: Final = _logged_fields(logged_call.get("litellm_params")) + proxy_server_request: Final = _logged_fields(litellm_params.get("proxy_server_request")) + caller_credential_fields: Final = proxy_server_request.get("caller_credential_fields") + logged_metadata: Final = _logged_fields(_logged_fields(logged_call.get("standard_logging_object")).get("metadata")) + return ( + ( + isinstance(caller_credential_fields, tuple) + and any(field in caller_credential_fields for field in _CALLER_TRANSPORT_FIELDS) + ) + or _sent_a_provider_credential_or_endpoint(_logged_fields(proxy_server_request.get("body"))) + or bool(litellm_params.get("token_exchange_endpoint")) + or logged_metadata.get("used_client_oauth_token") is True + or github_copilot_user_session_from(litellm_params) is not None + ) + + +def _is_deployment_failure(exception: openai.APIError, logged_call: Mapping[str, object]) -> bool: + """True for an error raised while calling the deployment with its own credential, on a request the proxy + did not reject itself. An error LiteLLM raises before it calls the provider, such as a missing deployment key, + is not one.""" + return not ( + isinstance(exception, CallerCredentialAuthenticationError) + or bool(logged_call.get(LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL)) + or logged_call.get("api_call_start_time") is None + or _used_a_caller_credential(logged_call) + ) + + def _bounded_requested_model_label(requested_model: object, router_originated: bool = False) -> str | None: """ Bound ``requested_model`` label cardinality: names the router recognizes @@ -2817,6 +2905,8 @@ class PrometheusLogger(CustomLogger): information from various callback function signatures and determines if the request represents an invalid API key attempt that should be filtered from metrics. + An error raised while calling the deployment with its own credential is never skipped. + Args: kwargs: Dictionary potentially containing exception and other data user_api_key_dict: User API key authentication object (currently unused) @@ -2836,6 +2926,11 @@ class PrometheusLogger(CustomLogger): if exception is None and kwargs: exception = kwargs.get("exception") + if isinstance(exception, openai.APIError) and _is_deployment_failure( + exception, _LOGGED_FIELDS.validate_python(kwargs or {}) + ): + return False + if self._is_invalid_api_key_request(status_code, exception=exception): verbose_logger.debug( "Skipping Prometheus metrics for invalid API key request: status_code=%s, exception=%s", diff --git a/litellm/llms/azure/chat/gpt_5_transformation.py b/litellm/llms/azure/chat/gpt_reasoning_transformation.py similarity index 98% rename from litellm/llms/azure/chat/gpt_5_transformation.py rename to litellm/llms/azure/chat/gpt_reasoning_transformation.py index 4bec335064e..5bd03f511c5 100644 --- a/litellm/llms/azure/chat/gpt_5_transformation.py +++ b/litellm/llms/azure/chat/gpt_reasoning_transformation.py @@ -4,7 +4,7 @@ from collections.abc import Collection from typing import ClassVar, Final import litellm -from litellm.llms.openai.chat.gpt_5_transformation import ( +from litellm.llms.openai.chat.gpt_reasoning_transformation import ( OpenAIGPT5Config, is_gpt_reasoning_series_name, ) diff --git a/litellm/llms/azure/chat/gpt_transformation.py b/litellm/llms/azure/chat/gpt_transformation.py index 03fe27437a5..26ae30635d4 100644 --- a/litellm/llms/azure/chat/gpt_transformation.py +++ b/litellm/llms/azure/chat/gpt_transformation.py @@ -16,7 +16,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_azure_openai_messages, ) from litellm.llms.base_llm.chat.transformation import BaseLLMException -from litellm.llms.openai.chat.gpt_5_transformation import GPT_REASONING_SERIES_MARKERS +from litellm.llms.openai.chat.gpt_reasoning_transformation import GPT_REASONING_SERIES_MARKERS from litellm.types.llms.azure import ( API_VERSION_MONTH_SUPPORTED_RESPONSE_FORMAT, API_VERSION_YEAR_SUPPORTED_RESPONSE_FORMAT, diff --git a/litellm/llms/azure/responses/transformation.py b/litellm/llms/azure/responses/transformation.py index 6d904503df8..647ad97e02c 100644 --- a/litellm/llms/azure/responses/transformation.py +++ b/litellm/llms/azure/responses/transformation.py @@ -12,7 +12,7 @@ from pydantic import TypeAdapter from litellm._logging import verbose_logger from litellm.constants import AZURE_DEFAULT_RESPONSES_API_VERSION from litellm.litellm_core_utils.url_utils import encode_url_path_segment -from litellm.llms.azure.chat.gpt_5_transformation import AzureOpenAIGPT5Config +from litellm.llms.azure.chat.gpt_reasoning_transformation import AzureOpenAIGPT5Config from litellm.llms.azure.common_utils import BaseAzureLLM from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.types.llms.openai import * diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index 3b541310d28..6aed7eb3c54 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -21,7 +21,7 @@ from litellm.llms.azure_ai.common_utils import ( is_foundry_model_inference_base, ) from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj -from litellm.llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config +from litellm.llms.openai.chat.gpt_reasoning_transformation import OpenAIGPT5Config from litellm.llms.openai.common_utils import drop_params_from_unprocessable_entity_error from litellm.llms.openai.openai import OpenAIChatCompletionResponseIterator, OpenAIConfig from litellm.llms.xai.chat.transformation import XAIChatConfig diff --git a/litellm/llms/azure_ai/common_utils.py b/litellm/llms/azure_ai/common_utils.py index cffe9049de6..65d73363b30 100644 --- a/litellm/llms/azure_ai/common_utils.py +++ b/litellm/llms/azure_ai/common_utils.py @@ -6,7 +6,7 @@ from urllib.parse import urlparse import litellm from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter -from litellm.llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config +from litellm.llms.openai.chat.gpt_reasoning_transformation import OpenAIGPT5Config from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.types.router import GenericLiteLLMParams diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py similarity index 100% rename from litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py rename to litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude2_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_legacy_text_completions_transformation.py similarity index 100% rename from litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude2_transformation.py rename to litellm/llms/bedrock/chat/invoke_transformations/anthropic_legacy_text_completions_transformation.py diff --git a/litellm/llms/bedrock/chat/mantle/transformation.py b/litellm/llms/bedrock/chat/mantle/transformation.py index a1aa2034e4c..b5a352bab48 100644 --- a/litellm/llms/bedrock/chat/mantle/transformation.py +++ b/litellm/llms/bedrock/chat/mantle/transformation.py @@ -12,7 +12,7 @@ from typing import TYPE_CHECKING, Any, Final from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media from litellm.llms.anthropic.chat.transformation import JsonModeRequest -from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import ( +from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude_transformation import ( AmazonAnthropicClaudeConfig, ) from litellm.llms.bedrock.common_utils import build_mantle_messages_url diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index cb10c2384cd..8c4702bcf99 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -1052,7 +1052,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): # --- Anthropic: use existing AmazonAnthropicClaudeConfig --- if provider == LlmProviders.ANTHROPIC: - from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import ( + from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude_transformation import ( AmazonAnthropicClaudeConfig, ) diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude_transformation.py similarity index 100% rename from litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py rename to litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude_transformation.py diff --git a/litellm/llms/bedrock/messages/mantle_transformation.py b/litellm/llms/bedrock/messages/mantle_transformation.py index d4a2d2eabe9..7b9dde42c4a 100644 --- a/litellm/llms/bedrock/messages/mantle_transformation.py +++ b/litellm/llms/bedrock/messages/mantle_transformation.py @@ -19,7 +19,7 @@ from litellm.llms.anthropic.pass_through.messages.transformation import ( AnthropicMessagesConfig, ) from litellm.llms.bedrock.common_utils import build_mantle_messages_url -from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( +from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude_transformation import ( AmazonAnthropicClaudeMessagesConfig, ) from litellm.types.llms.anthropic_messages.anthropic_response import ( diff --git a/litellm/llms/bedrock_mantle/chat/transformation.py b/litellm/llms/bedrock_mantle/chat/transformation.py index 61c7d51a889..9e168f14401 100644 --- a/litellm/llms/bedrock_mantle/chat/transformation.py +++ b/litellm/llms/bedrock_mantle/chat/transformation.py @@ -24,7 +24,7 @@ from litellm.llms.bedrock_mantle.common_utils import ( BEDROCK_MANTLE_DEFAULT_REGION, BedrockMantleAuthMixin, ) -from litellm.llms.openai.chat.gpt_5_transformation import get_effort_level, is_gpt_reasoning_series_name +from litellm.llms.openai.chat.gpt_reasoning_transformation import get_effort_level, is_gpt_reasoning_series_name from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.types.router import GenericLiteLLMParams diff --git a/litellm/llms/openai/chat/gpt_5_transformation.py b/litellm/llms/openai/chat/gpt_reasoning_transformation.py similarity index 100% rename from litellm/llms/openai/chat/gpt_5_transformation.py rename to litellm/llms/openai/chat/gpt_reasoning_transformation.py diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 73f40f5e0e4..9d4df354a30 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -48,7 +48,7 @@ from litellm.utils import ( from ...types.llms.openai import * from ..base import BaseLLM -from .chat.gpt_5_transformation import OpenAIGPT5Config +from .chat.gpt_reasoning_transformation import OpenAIGPT5Config from .chat.gpt_transformation import OpenAIGPTConfig, OpenAIUnknownModelConfig from .chat.o_series_transformation import OpenAIOSeriesConfig from .common_utils import ( diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index b90f6aa01da..07822bb17b3 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -25,7 +25,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig -from litellm.llms.openai.chat.gpt_5_transformation import ( +from litellm.llms.openai.chat.gpt_reasoning_transformation import ( is_gpt_reasoning_series_name, unsupported_reasoning_effort_message, ) @@ -291,7 +291,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): level the row cannot gate because the same model takes it here and refuses it on chat. Azure overrides this so a bare deployment name reads the azure/ entry. """ - from litellm.llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config + from litellm.llms.openai.chat.gpt_reasoning_transformation import OpenAIGPT5Config return OpenAIGPT5Config.row_disables_reasoning_effort_level(model, level) @@ -305,7 +305,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): copy of the rule here is what let this surface keep forwarding temperature after the chat surface stopped. """ - from litellm.llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config + from litellm.llms.openai.chat.gpt_reasoning_transformation import OpenAIGPT5Config return OpenAIGPT5Config.effort_resolves_to_none(model, effort) diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index 36337fddbdc..115d914d3ac 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -249,5 +249,13 @@ "/v1/responses", "/v1/messages" ] + }, + "refuseless": { + "base_url": "https://api.refuseless.com/v1", + "api_key_env": "REFUSELESS_API_KEY", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ] } } diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/__init__.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/__init__.py index a9c1e5819f2..991d7e8f1bc 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/__init__.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/__init__.py @@ -12,7 +12,7 @@ def get_vertex_ai_partner_model_config(model: str, vertex_publisher_or_api_spec: return VertexAIAi21Config() elif vertex_publisher_or_api_spec == "openapi" or vertex_publisher_or_api_spec == "mistralai": - from .llama3.transformation import VertexAILlama3Config + from .llama.transformation import VertexAILlama3Config return VertexAILlama3Config() else: diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/llama/transformation.py similarity index 100% rename from litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/transformation.py rename to litellm/llms/vertex_ai/vertex_ai_partner_models/llama/transformation.py diff --git a/litellm/main.py b/litellm/main.py index b3d62daa4d8..23eac8ce0aa 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -130,7 +130,7 @@ from litellm.llms.bedrock.common_utils import ( from litellm.llms.bedrock_mantle.chat.claude_transformation import bedrock_mantle_chat_config from litellm.llms.cohere.common_utils import CohereModelInfo from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler, http2_enabled -from litellm.llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config +from litellm.llms.openai.chat.gpt_reasoning_transformation import OpenAIGPT5Config from litellm.llms.openai.common_utils import get_openai_credentials from litellm.llms.openai_like.json_loader import JSONProviderRegistry from litellm.llms.vertex_ai.common_utils import ( @@ -237,7 +237,7 @@ from .llms.azure_ai.anthropic.transformation import azure_anthropic_messages_url from .llms.azure_ai.embed import AzureAIEmbedding from .llms.bedrock.chat import BedrockConverseLLM from .llms.bedrock.chat.converse_transformation import AmazonConverseConfig -from .llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import AmazonAnthropicClaudeConfig +from .llms.bedrock.chat.invoke_transformations.anthropic_claude_transformation import AmazonAnthropicClaudeConfig from .llms.bedrock.chat.native_output_fallback import NativeOutputAttemptResult, call_with_native_output_fallback from .llms.bedrock.embed.embedding import BedrockEmbedding from .llms.bedrock.image_edit.handler import BedrockImageEdit diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index c913c879ccf..b820b058eb3 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -69,6 +69,106 @@ "mode": "image_generation", "output_cost_per_image": 0.036 }, + "abliteration/abliterated-model": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 1e-06, + "litellm_provider": "abliteration", + "max_input_tokens": 262144, + "max_output_tokens": 262134, + "mode": "chat", + "output_cost_per_token": 3e-06, + "source": "https://abliteration.ai/pricing", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "reasoning_effort_levels": [ + "none", + "minimal", + "low", + "medium", + "high", + "xhigh", + "max" + ] + }, + "abliteration/abliterated-model-large": { + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "litellm_provider": "abliteration", + "max_input_tokens": 1000000, + "max_output_tokens": 999990, + "mode": "chat", + "output_cost_per_token": 5e-06, + "source": "https://abliteration.ai/pricing", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false, + "reasoning_effort_levels": [ + "none", + "minimal", + "low", + "medium", + "high", + "xhigh", + "max" + ] + }, + "abliteration/abliterated-model-large-v2": { + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "litellm_provider": "abliteration", + "max_input_tokens": 1000000, + "max_output_tokens": 999990, + "mode": "chat", + "output_cost_per_token": 5e-06, + "source": "https://abliteration.ai/pricing", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_none_reasoning_effort": false, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false, + "reasoning_effort_levels": [ + "minimal", + "low", + "medium", + "high", + "xhigh", + "max" + ] + }, + "abliteration/abliterated-model-large-flash": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 1e-06, + "litellm_provider": "abliteration", + "max_input_tokens": 1048576, + "max_output_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2e-06, + "source": "https://abliteration.ai/pricing", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "reasoning_effort_levels": [ + "none", + "minimal", + "low", + "medium", + "high", + "xhigh", + "max" + ] + }, "ai21.j2-mid-v1": { "input_cost_per_token": 1.25e-05, "litellm_provider": "bedrock", @@ -75800,5 +75900,110 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_vision": true + }, + "refuseless/glm-5.3-abliterated": { + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "refuseless", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.25e-06, + "source": "https://refuseless.com/pricing", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false, + "reasoning_effort_levels": [ + "low", + "high", + "max" + ] + }, + "refuseless/glm-5.3-flash-abliterated": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 2.5e-07, + "litellm_provider": "refuseless", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 8e-07, + "source": "https://refuseless.com/pricing", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false, + "reasoning_effort_levels": [ + "low", + "high", + "max" + ] + }, + "refuseless/deepseek-v4.1-flash-abliterated": { + "cache_read_input_token_cost": 4.5e-08, + "input_cost_per_token": 2.2e-07, + "litellm_provider": "refuseless", + "max_input_tokens": 1048576, + "max_output_tokens": 130000, + "mode": "chat", + "output_cost_per_token": 7.5e-07, + "source": "https://refuseless.com/pricing", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false, + "reasoning_effort_levels": [ + "low", + "high", + "max" + ] + }, + "refuseless/deepseek-v4.1-flash-fast": { + "cache_read_input_token_cost": 3e-09, + "input_cost_per_token": 5e-08, + "litellm_provider": "refuseless", + "max_input_tokens": 1048576, + "max_output_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 1.5e-07, + "source": "https://refuseless.com/pricing", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false, + "reasoning_effort_levels": [ + "low", + "high", + "max" + ] + }, + "refuseless/minimax-m3.1-flash": { + "cache_read_input_token_cost": 2e-09, + "input_cost_per_token": 5e-08, + "litellm_provider": "refuseless", + "max_input_tokens": 1048576, + "max_output_tokens": 512000, + "mode": "chat", + "output_cost_per_token": 2.5e-07, + "source": "https://refuseless.com/pricing", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false, + "reasoning_effort_levels": [ + "low", + "high", + "max" + ] } } diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index f95b553c2b4..49b667af819 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -32,6 +32,23 @@ } }, "providers": { + "refuseless": { + "display_name": "Refuseless (`refuseless`)", + "url": "https://docs.litellm.ai/docs/providers/refuseless", + "endpoints": { + "chat_completions": true, + "messages": false, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": false + } + }, "a2a": { "display_name": "A2A (Agent-to-Agent) (`a2a`)", "url": "https://docs.litellm.ai/docs/providers/a2a", diff --git a/litellm/proxy/agent_endpoints/model_list_helpers.py b/litellm/proxy/agent_endpoints/model_list_helpers.py index 4a88644bae5..447b4f52aaf 100644 --- a/litellm/proxy/agent_endpoints/model_list_helpers.py +++ b/litellm/proxy/agent_endpoints/model_list_helpers.py @@ -51,7 +51,7 @@ async def append_agents_to_model_group( async def append_agents_to_model_info( models: list[dict], user_api_key_dict: UserAPIKeyAuth, -) -> list[dict]: +) -> list[dict[str, object]]: """ Append A2A agents to model info list for UI display. diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index e7a06df6bab..baa150efc86 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -99,6 +99,8 @@ _CREDENTIAL_HEADER_NAMES: Final = SpecialHeaders.litellm_credential_header_names {"cookie", "proxy-authorization"} ) _TRANSPORT_ONLY_CREDENTIAL_KEYS: Final = frozenset({"provider_specific_header", "headers", "api_key"}) +_REQUEST_FIELDS: Final = TypeAdapter(Mapping[str, object]) +_REQUEST_ENTRIES: Final = TypeAdapter(tuple[object, ...]) # Matches any header of the form x--session-id (case-insensitive). # Excludes the two explicit litellm headers which are handled with higher priority. @@ -289,6 +291,9 @@ _DESTINATION_ONLY_CALLBACK_VARS: Final[frozenset[str]] = frozenset({CAPTURE_MESS _CALLBACK_CREDENTIAL_KEYS: Final = frozenset(StandardCallbackDynamicParams.__annotations__) | frozenset( {TRUSTED_CALLBACK_VARS_FIELD} ) +_HEADERS_THAT_CANNOT_CHANGE_PROVIDER_AUTH: Final = frozenset( + {*(f"x-litellm-{key}" for key in StandardLoggingUserAPIKeyMetadata.__annotations__), *ANTHROPIC_API_HEADERS} +) service_logger_obj: Final = ServiceLogging() # used for tracking latency on OTEL # Bounded dedup for stale-alias warnings (FIFO eviction when over cap). _MAX_STALE_ALIAS_WARNING_KEYS: Final = 10_000 @@ -1451,6 +1456,44 @@ def clean_headers( return clean_headers +def _request_fields(value: object) -> Mapping[str, object]: + try: + return _REQUEST_FIELDS.validate_python(value) + except PydanticValidationError: + return MappingProxyType({}) + + +def _request_entries(value: object) -> tuple[object, ...]: + if value is None: + return () + try: + return _REQUEST_ENTRIES.validate_python(value) + except PydanticValidationError: + return (value,) + + +def _carries_a_caller_header(headers: object) -> bool: + return any(name.lower() not in _HEADERS_THAT_CANNOT_CHANGE_PROVIDER_AUTH for name in _request_fields(headers)) + + +def _caller_credential_fields(data: Mapping[str, object]) -> tuple[str, ...]: + """Names of the request fields through which the caller can change how the provider authenticates the call: + ``api_key``, ``organization``, and ``headers``, ``extra_headers`` or a ``provider_specific_header`` entry's + ``extra_headers`` that carry any header besides the key and user details the proxy adds under + ``add_user_information_to_llm_headers`` and the ``anthropic-version`` and ``anthropic-beta`` headers.""" + carried: Final = { + "api_key": "api_key" in data, + "extra_headers": _carries_a_caller_header(data.get("extra_headers")), + "headers": _carries_a_caller_header(data.get("headers")), + "organization": "organization" in data, + "provider_specific_header": any( + _carries_a_caller_header(_request_fields(entry).get("extra_headers")) + for entry in _request_entries(data.get("provider_specific_header")) + ), + } + return tuple(name for name, sent in sorted(carried.items()) if sent) + + def _is_credential_header(header: str) -> bool: """Whether `header` carries a caller credential rather than request context.""" return header.lower().removeprefix("x-pass-") in _CREDENTIAL_HEADER_NAMES @@ -2605,6 +2648,8 @@ async def add_litellm_data_to_request( data[_metadata_variable_name]["used_client_oauth_token"] = add_provider_specific_headers_to_request( data=data, headers=_headers ) + caller_fields: Final = _caller_credential_fields(_REQUEST_FIELDS.validate_python(data)) + data["proxy_server_request"]["caller_credential_fields"] = caller_fields # rebind-ok: data is an out-param ## Cache Controls cache_control_header: Final = _headers.get("Cache-Control", None) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 11f616a5e18..e625ea6f9f2 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -145,6 +145,7 @@ from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.bedrock import AwsSessionTag from litellm.types.proxy.management_endpoints.model_management_endpoints import ( AutoRouterClassifierDefaultPromptResponse, + PatchModelRequest, UpdateUsefulLinksRequest, ) from litellm.types.router import ( @@ -160,6 +161,7 @@ from litellm.types.utils import ( OIDC_SECRET_REFERENCE_PREFIX, echoed_cost_map_fields, echoed_cost_map_pricing_fields, + is_server_derived_pricing_key, without_server_derived_pricing, ) from litellm.utils import get_utc_datetime @@ -995,6 +997,7 @@ def _ptu_zeroed_pricing( def _ptu_pricing_delta( *, stored_model_info: Mapping[str, object], + stored_litellm_params: Mapping[str, object], model_info: Mapping[str, object], litellm_params: Mapping[str, object], patch: updateDeployment, @@ -1014,6 +1017,24 @@ def _ptu_pricing_delta( supplied: Final = patch.litellm_params.model_dump(exclude_none=True) if patch.litellm_params else _EMPTY_MODEL_INFO zeroed: Final = _ptu_zeroed_pricing(model_info=model_info, litellm_params=litellm_params, supplied=supplied) if zeroed: + retains_ptu: Final = all(stored_model_info.get(field) is not None for field in _PTU_PRICED_PAIR) + protected_pricing: Final = frozenset( + field + for field in _CUSTOM_PRICING_FIELDS + if retains_ptu + and (stored_litellm_params.get(field) is not None or stored_model_info.get(field) is not None) + ).union(zeroed) + removed_ptu_pricing: Final = tuple( + field for field in _param_keys_to_remove(patch) if field in protected_pricing + ) + if removed_ptu_pricing: + raise HTTPException( + status_code=400, + detail=( + f"PTU billing requires {', '.join(removed_ptu_pricing)}, so these pricing parameters cannot be removed. " + "Keep them when enabling PTU, or clear ptu_count and cost_per_ptu_per_hour together to disable it." + ), + ) return zeroed, frozenset() was_ptu: Final = any(stored_model_info.get(field) is not None for field in _PTU_PRICED_PAIR) if not was_ptu or not _explicitly_cleared_ptu_fields(patch.model_info) & _PTU_PRICED_PAIR: @@ -1069,6 +1090,10 @@ def _loaded_catalog_entry( return loaded_catalog().get(catalog_key, MappingProxyType({})) +def _param_keys_to_remove(patch: updateDeployment) -> tuple[str, ...]: + return patch.litellm_params_to_remove if isinstance(patch, PatchModelRequest) else () + + def update_db_model( db_model: Deployment, updated_patch: updateDeployment, @@ -1077,7 +1102,8 @@ def update_db_model( if updated_patch.model_info is not None: _raise_if_ptu_cost_attribution_disabled(updated_patch.model_info.model_dump(exclude_none=True)) merged_model_name: Final = updated_patch.model_name or db_model.model_name - merged_litellm_params: Final = db_model.litellm_params.model_dump(exclude_none=True) + stored_litellm_params: Final = db_model.litellm_params.model_dump(exclude_none=True, exclude_unset=True) + merged_litellm_params: Final = dict(stored_litellm_params) stored_model_info: Final = _dump_without_none(db_model.model_info) echoed_pricing: Final = echoed_cost_map_pricing_fields(stored_model_info) merged_model_info: Final[dict[str, object]] = { @@ -1093,7 +1119,7 @@ def update_db_model( if k == "complexity_router_config" else encrypt_value_helper(v) ) - for k, v in updated_patch.litellm_params.model_dump(exclude_none=True).items() + for k, v in updated_patch.litellm_params.model_dump(exclude_none=True, exclude_unset=True).items() } merged_litellm_params.update(encrypted_params) @@ -1144,11 +1170,17 @@ def update_db_model( for field in _explicitly_cleared_ptu_fields(updated_patch.model_info): merged_model_info.pop(field, None) + for field in _param_keys_to_remove(updated_patch): + merged_litellm_params.pop(field, None) + if is_server_derived_pricing_key(field): + merged_model_info.pop(field, None) + _validate_ptu_model_info(merged_model_info) ptu_pricing, ptu_released = _ptu_pricing_delta( stored_model_info=db_model.model_info.model_dump(exclude_none=True) if db_model.model_info else _EMPTY_MODEL_INFO, + stored_litellm_params=stored_litellm_params, model_info=merged_model_info, litellm_params=merged_litellm_params, patch=updated_patch, @@ -1255,7 +1287,7 @@ async def get_kubernetes_pods( ) async def patch_model( model_id: str, # Get model_id from path parameter - patch_data: updateDeployment, # Create a specific schema for PATCH operations + patch_data: PatchModelRequest, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -1319,18 +1351,34 @@ async def patch_model( param=None, ) + removals: Final = _param_keys_to_remove(patch_data) + incoming_params: Final = ( + updateLiteLLMParams.model_validate( + { + **(patch_data.litellm_params.model_dump(exclude_unset=True) if patch_data.litellm_params else {}), + **dict.fromkeys(removals), + } + ) + if removals + else patch_data.litellm_params + ) + remaining_params: Final = db_model.litellm_params.model_copy(update=dict.fromkeys(removals)) write_authorization: Final = await ModelManagementAuthChecks.can_user_make_model_call( model_params=db_model, user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, premium_user=premium_user, - incoming_params=submitted_litellm_params(patch_data.litellm_params), + incoming_params=submitted_litellm_params(incoming_params), member_operation="update", incoming_model_params=patch_data, ) member_write: Final = write_authorization if isinstance(write_authorization, MemberAutoRouterWrite) else None - member_marker: Final = _member_auto_router_marker_for_update( - incoming_params=patch_data.litellm_params, existing=db_model, member_write=member_write + member_marker: Final = ( + False + if db_model.model_info.member_auto_router and STRATEGY_ROUTER_PARAM_FIELDS.intersection(removals) + else _member_auto_router_marker_for_update( + incoming_params=incoming_params, existing=db_model, member_write=member_write + ) ) marker_info: Final = ( ModelInfo(id=db_model.model_info.id) @@ -1367,27 +1415,26 @@ async def patch_model( ) ModelManagementAuthChecks.can_user_attach_credential( - litellm_params=patch_data.litellm_params, + litellm_params=incoming_params, user_api_key_dict=user_api_key_dict, existing_litellm_params=db_model.litellm_params, null_detaches=True, ) - await _raise_on_invalid_credential_name(patch_data.litellm_params, db_model.litellm_params, prisma_client) + await _raise_on_invalid_credential_name(incoming_params, db_model.litellm_params, prisma_client) ModelManagementAuthChecks.can_user_set_aws_session_tags( - litellm_params=patch_data.litellm_params, + litellm_params=incoming_params, user_api_key_dict=user_api_key_dict, existing_litellm_params=db_model.litellm_params, + null_detaches="aws_session_tags" in removals, ) _raise_on_strategy_router_write_violation( - incoming_params=patch_data.litellm_params, - existing_params=db_model.litellm_params, + incoming_params=incoming_params, + existing_params=remaining_params, ) - effective_params: Final = _effective_complexity_router_params( - patch_data.litellm_params, db_model.litellm_params - ) + effective_params: Final = _effective_complexity_router_params(incoming_params, remaining_params) requested_model_name: Final = patch_data.model_name stored_model_name: str | None = None @@ -1818,20 +1865,20 @@ async def _update_team_model_in_db( db_model=db_model, ) - # Ensure model_info exists and set team_public_model_name - if patch_data.model_info is None: - from litellm.types.router import ModelInfo - - patch_data.model_info = ModelInfo() - patch_data.model_info.team_public_model_name = public_model_name - # Check if team assignment is new or changed db_team_id: Final = db_model.model_info.team_id if db_model.model_info else None is_new_team_assignment: Final = db_team_id != patch_team_id # Team rows keep their internal UUID-based model_name; the public name lives in model_info - patch_data.model_name = f"model_name_{patch_team_id}_{uuid.uuid4()}" if is_new_team_assignment else None - row: Final = await write_row(update_db_model(db_model=db_model, updated_patch=patch_data)) + team_patch: Final = patch_data.model_copy( + update={ + "model_name": f"model_name_{patch_team_id}_{uuid.uuid4()}" if is_new_team_assignment else None, + "model_info": (patch_data.model_info or ModelInfo()).model_copy( + update={"team_public_model_name": public_model_name} + ), + } + ) + row: Final = await write_row(update_db_model(db_model=db_model, updated_patch=team_patch)) if is_new_team_assignment: await _setup_new_team_model_assignment( @@ -2228,14 +2275,16 @@ class ModelManagementAuthChecks: litellm_params: GenericLiteLLMParams | None, user_api_key_dict: UserAPIKeyAuth, existing_litellm_params: GenericLiteLLMParams | None = None, + *, + null_detaches: bool = False, ) -> Literal[True]: - if litellm_params is None or litellm_params.aws_session_tags is None: + if litellm_params is None or (litellm_params.aws_session_tags is None and not null_detaches): return True if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: return True existing_tags: Final = existing_litellm_params.aws_session_tags if existing_litellm_params is not None else None if existing_tags is not None and _canonical_session_tags(existing_tags) == _canonical_session_tags( - litellm_params.aws_session_tags + litellm_params.aws_session_tags or () ): return True raise ProxyException( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 5c7a966b702..35f6acb0fbd 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -16760,6 +16760,35 @@ async def _find_model_by_id( return all_models, search_total_count +async def _models_with_stored_param_keys( + models: Sequence[Mapping[str, object]], model_id: str, prisma_client: PrismaClient +) -> Sequence[Mapping[str, object]]: + visible_model: Final = next( + (model for model in models if _with_get(model.get("model_info", {})).get("id") == model_id), None + ) + if visible_model is None: + return models + stored: Final = await ModelRepository(WriterPinnedClient(prisma_client.db)).table.find_unique( + where={"model_id": model_id} + ) + if stored is None: + return models + raw_params: Final = _UNTYPED_VALUE.validate_python(stored.litellm_params) + stored_params: Final = ( + _LITELLM_PARAMS_COPY.validate_json(raw_params) + if isinstance(raw_params, str) + else _LITELLM_PARAMS_COPY.validate_python(raw_params) + ) + visible_params: Final = _LITELLM_PARAMS_COPY.validate_python(visible_model.get("litellm_params", {})) + editor_model: Final = MappingProxyType( + { + **visible_model, + "litellm_params": {key: value for key, value in visible_params.items() if key in stored_params}, + } + ) + return tuple(editor_model if model is visible_model else model for model in models) + + @router.get( "/v2/model/info", tags=["model management"], @@ -16777,6 +16806,10 @@ async def model_info_v2( size: int = Query(50, description="Page size", ge=1), search: str | None = fastapi.Query(None, description="Search model names (case-insensitive partial match)"), modelId: str | None = fastapi.Query(None, description="Search for a specific model by its unique ID"), + use_stored_params: bool = fastapi.Query( + False, + description="With modelId, omit runtime-only litellm_params defaults from a DB-stored model's editor view.", + ), teamId: str | None = fastapi.Query( None, description="Filter models by team ID. Returns models with direct_access=True or teamId in access_via_team_ids", @@ -16828,6 +16861,7 @@ async def model_info_v2( page / size: Pagination controls (defaults: page=1, size=50). search: Case-insensitive partial match on model name or team public name. modelId: Return a single deployment by LiteLLM model id. + use_stored_params: With modelId, return only persisted litellm_params keys, retaining normal redaction. teamId: Filter to models with direct access or team membership for this team id. sortBy / sortOrder: Sort by model_name, created_at, updated_at, costs, or status. blocked: Filter by routing status (false = active, true = paused). @@ -16864,6 +16898,9 @@ async def model_info_v2( """ global llm_model_list, general_settings, user_config_file_path, proxy_config, llm_router + if use_stored_params is True and not isinstance(modelId, str): + raise HTTPException(status_code=400, detail="use_stored_params requires modelId") + # Return empty data array when no models are configured (graceful handling for fresh installs) if llm_router is None or not llm_router.model_list: return { @@ -16999,6 +17036,14 @@ async def model_info_v2( # Translate `model_name` to the public name for team-scoped rows. all_models = [_translate_model_name_for_response(m) for m in all_models] + if use_stored_params is True and isinstance(modelId, str): + return _paginate_models_response( + all_models=await _models_with_stored_param_keys(all_models, modelId, prisma_client), + page=page, + size=size, + total_count=search_total_count, + search=search, + ) return _paginate_models_response( all_models=all_models, page=page, @@ -17512,7 +17557,7 @@ def _filter_v1_model_info_deployments( return [model for model in all_models if _deployment_matches_allowed_model_names(model, allowed_model_names)] -def _translate_model_name_for_response(model: dict) -> dict: +def _translate_model_name_for_response(model: dict) -> dict[str, object]: """For team-scoped DB rows, replace `model_name` with the public name in `model_info.team_public_model_name` before returning. The DB column and the in-memory router index keep the internal mangled name diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index 19ee465c7dc..d96f2976508 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -17,6 +17,34 @@ ], "default_model_placeholder": "scaledown/summarize" }, + { + "provider": "ABLITERATION", + "provider_display_name": "Abliteration", + "litellm_provider": "abliteration", + "credential_fields": [ + { + "key": "api_key", + "label": "API Key", + "placeholder": null, + "tooltip": null, + "required": true, + "field_type": "password", + "options": null, + "default_value": null + }, + { + "key": "api_base", + "label": "API Base", + "placeholder": "https://api.abliteration.ai/v1", + "tooltip": "Optional. Leave blank to use the provider's default endpoint.", + "required": false, + "field_type": "text", + "options": null, + "default_value": null + } + ], + "default_model_placeholder": "abliteration/abliterated-model-large-v2" + }, { "provider": "AIML", "provider_display_name": "AI/ML API", @@ -3203,6 +3231,34 @@ ], "default_model_placeholder": "gpt-3.5-turbo" }, + { + "provider": "REFUSELESS", + "provider_display_name": "Refuseless", + "litellm_provider": "refuseless", + "credential_fields": [ + { + "key": "api_key", + "label": "API Key", + "placeholder": null, + "tooltip": null, + "required": true, + "field_type": "password", + "options": null, + "default_value": null + }, + { + "key": "api_base", + "label": "API Base", + "placeholder": "https://api.refuseless.com/v1", + "tooltip": "Optional. Leave blank to use the provider's default endpoint.", + "required": false, + "field_type": "text", + "options": null, + "default_value": null + } + ], + "default_model_placeholder": "refuseless/glm-5.3-abliterated" + }, { "provider": "REPLICATE", "provider_display_name": "Replicate", diff --git a/litellm/responses/bridge_check.py b/litellm/responses/bridge_check.py index 8b00cdf4813..504150b330c 100644 --- a/litellm/responses/bridge_check.py +++ b/litellm/responses/bridge_check.py @@ -10,7 +10,7 @@ from litellm.llms.azure_ai.common_utils import ( azure_ai_supports_native_responses, foundry_chat_rejects_function_tools_while_reasoning, ) -from litellm.llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config +from litellm.llms.openai.chat.gpt_reasoning_transformation import OpenAIGPT5Config from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import OpenAIWebSearchOptions from litellm.types.utils import LlmProviders diff --git a/litellm/router_utils/reasoning_effort_capability.py b/litellm/router_utils/reasoning_effort_capability.py index 4357722970f..16b8a22a40f 100644 --- a/litellm/router_utils/reasoning_effort_capability.py +++ b/litellm/router_utils/reasoning_effort_capability.py @@ -469,7 +469,7 @@ def _supports_none_reasoning_effort(model_info: Mapping[str, object], flag: obje if model_info.get("litellm_provider") != "azure": return flag is not False - from litellm.llms.azure.chat.gpt_5_transformation import AzureOpenAIGPT5Config + from litellm.llms.azure.chat.gpt_reasoning_transformation import AzureOpenAIGPT5Config key: Final = model_info.get("key") if not isinstance(key, str) or not AzureOpenAIGPT5Config.is_model_gpt_5_model(key): diff --git a/litellm/types/proxy/management_endpoints/model_management_endpoints.py b/litellm/types/proxy/management_endpoints/model_management_endpoints.py index c012f872280..f032bb20c2a 100644 --- a/litellm/types/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/types/proxy/management_endpoints/model_management_endpoints.py @@ -1,11 +1,37 @@ from collections.abc import Mapping from datetime import datetime +from typing import Annotated, Final -from pydantic import ConfigDict, Field +from pydantic import ConfigDict, Field, model_validator from litellm.types.llms.base import LiteLLMBaseModel -from ...router import ModelGroupInfo +from ...router import ModelGroupInfo, updateDeployment + + +class PatchModelRequest(updateDeployment): + litellm_params_to_remove: tuple[Annotated[str, Field(min_length=1)], ...] = Field( + default=(), + description=( + "Top-level litellm_params keys to delete. Omitted parameters remain unchanged. " + "The required model key and keys also supplied in litellm_params cannot be removed. " + "Removing an absent key or repeating a key is a no-op." + ), + ) + + model_config = ConfigDict(frozen=True, protected_namespaces=()) + + @model_validator(mode="after") + def validate_param_removals(self) -> "PatchModelRequest": + removals: Final = frozenset(self.litellm_params_to_remove) + if "model" in removals: + raise ValueError("The required litellm_params.model cannot be removed") + conflicts: Final[frozenset[str]] = ( + removals.intersection(self.litellm_params.model_fields_set) if self.litellm_params else frozenset() + ) + if conflicts: + raise ValueError(f"Cannot set and remove the same litellm_params keys: {', '.join(sorted(conflicts))}") + return self class ModelGroupInfoProxy(ModelGroupInfo): diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 4f39741a079..7e8645fd0dd 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -4415,6 +4415,7 @@ GenericBudgetConfigType = dict[str, BudgetConfig] class LlmProviders(str, Enum): + REFUSELESS = "refuseless" OPENAI = "openai" CHATGPT = "chatgpt" OPENAI_LIKE = "openai_like" # embedding only @@ -4442,6 +4443,7 @@ class LlmProviders(str, Enum): VERTEX_AI = "vertex_ai" VERTEX_AI_BETA = "vertex_ai_beta" GEMINI = "gemini" + ABLITERATION = "abliteration" AI21 = "ai21" BASETEN = "baseten" BLACK_FOREST_LABS = "black_forest_labs" diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index c913c879ccf..b820b058eb3 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -69,6 +69,106 @@ "mode": "image_generation", "output_cost_per_image": 0.036 }, + "abliteration/abliterated-model": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 1e-06, + "litellm_provider": "abliteration", + "max_input_tokens": 262144, + "max_output_tokens": 262134, + "mode": "chat", + "output_cost_per_token": 3e-06, + "source": "https://abliteration.ai/pricing", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "reasoning_effort_levels": [ + "none", + "minimal", + "low", + "medium", + "high", + "xhigh", + "max" + ] + }, + "abliteration/abliterated-model-large": { + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "litellm_provider": "abliteration", + "max_input_tokens": 1000000, + "max_output_tokens": 999990, + "mode": "chat", + "output_cost_per_token": 5e-06, + "source": "https://abliteration.ai/pricing", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false, + "reasoning_effort_levels": [ + "none", + "minimal", + "low", + "medium", + "high", + "xhigh", + "max" + ] + }, + "abliteration/abliterated-model-large-v2": { + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "litellm_provider": "abliteration", + "max_input_tokens": 1000000, + "max_output_tokens": 999990, + "mode": "chat", + "output_cost_per_token": 5e-06, + "source": "https://abliteration.ai/pricing", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_none_reasoning_effort": false, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false, + "reasoning_effort_levels": [ + "minimal", + "low", + "medium", + "high", + "xhigh", + "max" + ] + }, + "abliteration/abliterated-model-large-flash": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 1e-06, + "litellm_provider": "abliteration", + "max_input_tokens": 1048576, + "max_output_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2e-06, + "source": "https://abliteration.ai/pricing", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "reasoning_effort_levels": [ + "none", + "minimal", + "low", + "medium", + "high", + "xhigh", + "max" + ] + }, "ai21.j2-mid-v1": { "input_cost_per_token": 1.25e-05, "litellm_provider": "bedrock", @@ -75800,5 +75900,110 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_vision": true + }, + "refuseless/glm-5.3-abliterated": { + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "refuseless", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.25e-06, + "source": "https://refuseless.com/pricing", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false, + "reasoning_effort_levels": [ + "low", + "high", + "max" + ] + }, + "refuseless/glm-5.3-flash-abliterated": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 2.5e-07, + "litellm_provider": "refuseless", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 8e-07, + "source": "https://refuseless.com/pricing", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false, + "reasoning_effort_levels": [ + "low", + "high", + "max" + ] + }, + "refuseless/deepseek-v4.1-flash-abliterated": { + "cache_read_input_token_cost": 4.5e-08, + "input_cost_per_token": 2.2e-07, + "litellm_provider": "refuseless", + "max_input_tokens": 1048576, + "max_output_tokens": 130000, + "mode": "chat", + "output_cost_per_token": 7.5e-07, + "source": "https://refuseless.com/pricing", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false, + "reasoning_effort_levels": [ + "low", + "high", + "max" + ] + }, + "refuseless/deepseek-v4.1-flash-fast": { + "cache_read_input_token_cost": 3e-09, + "input_cost_per_token": 5e-08, + "litellm_provider": "refuseless", + "max_input_tokens": 1048576, + "max_output_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 1.5e-07, + "source": "https://refuseless.com/pricing", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false, + "reasoning_effort_levels": [ + "low", + "high", + "max" + ] + }, + "refuseless/minimax-m3.1-flash": { + "cache_read_input_token_cost": 2e-09, + "input_cost_per_token": 5e-08, + "litellm_provider": "refuseless", + "max_input_tokens": 1048576, + "max_output_tokens": 512000, + "mode": "chat", + "output_cost_per_token": 2.5e-07, + "source": "https://refuseless.com/pricing", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": false, + "reasoning_effort_levels": [ + "low", + "high", + "max" + ] } } diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index d65da3f41fd..e62b74aa87d 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -33,6 +33,23 @@ } }, "providers": { + "refuseless": { + "display_name": "Refuseless (`refuseless`)", + "url": "https://docs.litellm.ai/docs/providers/refuseless", + "endpoints": { + "chat_completions": true, + "messages": false, + "responses": true, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": false + } + }, "a2a": { "display_name": "A2A (Agent-to-Agent) (`a2a`)", "url": "https://docs.litellm.ai/docs/providers/a2a", diff --git a/tests/e2e/coverage_registry/llm_conversational.yaml b/tests/e2e/coverage_registry/llm_conversational.yaml index e3994e598c6..ea095d4510c 100644 --- a/tests/e2e/coverage_registry/llm_conversational.yaml +++ b/tests/e2e/coverage_registry/llm_conversational.yaml @@ -71,9 +71,9 @@ - {id: llm.messages.anthropic.prompt_cache_5m.nonstream.works, module: llm, tier: P0, subject_endpoint: messages, route: anthropic, capability: prompt_cache_5m, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Prompt caching via Messages API"} - {id: llm.messages.anthropic.thinking.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: anthropic, capability: thinking, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Extended thinking via Messages API"} - {id: llm.messages.bedrock_converse.basic.nonstream.works_for_kimi_k3_max_tokens_1, module: llm, tier: P0, subject_endpoint: messages, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works_for_kimi_k3_max_tokens_1], source: "llm_translation/test_bedrock_kimi_k3_messages_e2e.py", rationale: "Bedrock Converse accepts max_tokens=1 on the Anthropic Messages surface after LiteLLM applies the model minimum", fail_before_fix: proven} -- {id: llm.messages.bedrock_invoke.mid_conversation_system.nonstream.cache_hit, module: llm, tier: P0, subject_endpoint: messages, route: bedrock_invoke, capability: mid_conversation_system, streaming: nonstream, assertions: [works, cache_hit], source: "llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py", rationale: "Flagged Claude 4.8+/5 must keep mid-conversation system reminders in messages; hoisting mutates the system prefix and collapses the prompt cache (#32578/#32831/#32882)", fail_before_fix: proven} -- {id: llm.messages.bedrock_invoke.mid_conversation_system.nonstream.works, module: llm, tier: P0, subject_endpoint: messages, route: bedrock_invoke, capability: mid_conversation_system, streaming: nonstream, assertions: [works], source: "llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py", rationale: "Claude <= 4.7 rejects role system inside messages; unflagged models must convert reminders to user turns in place (hoisting collapses the prompt cache) or every Claude Code session 400s (#32831)", fail_before_fix: proven} -- {id: llm.messages.bedrock_invoke.web_search_server_tool.nonstream.works, module: llm, tier: P0, subject_endpoint: messages, route: bedrock_invoke, capability: web_search_server_tool, streaming: nonstream, assertions: [works], source: "llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py", rationale: "Bedrock hosts no web_search server tool, so this only works because interception rewrites it before the upstream call and the agentic loop feeds the results back in native shape; a regression that short-circuits or forwards it instead yields raw text or AWS's 400", fail_before_fix: unproven} +- {id: llm.messages.bedrock_invoke.mid_conversation_system.nonstream.cache_hit, module: llm, tier: P0, subject_endpoint: messages, route: bedrock_invoke, capability: mid_conversation_system, streaming: nonstream, assertions: [works, cache_hit], source: "llms/bedrock/messages/invoke_transformations/anthropic_claude_transformation.py", rationale: "Flagged Claude 4.8+/5 must keep mid-conversation system reminders in messages; hoisting mutates the system prefix and collapses the prompt cache (#32578/#32831/#32882)", fail_before_fix: proven} +- {id: llm.messages.bedrock_invoke.mid_conversation_system.nonstream.works, module: llm, tier: P0, subject_endpoint: messages, route: bedrock_invoke, capability: mid_conversation_system, streaming: nonstream, assertions: [works], source: "llms/bedrock/messages/invoke_transformations/anthropic_claude_transformation.py", rationale: "Claude <= 4.7 rejects role system inside messages; unflagged models must convert reminders to user turns in place (hoisting collapses the prompt cache) or every Claude Code session 400s (#32831)", fail_before_fix: proven} +- {id: llm.messages.bedrock_invoke.web_search_server_tool.nonstream.works, module: llm, tier: P0, subject_endpoint: messages, route: bedrock_invoke, capability: web_search_server_tool, streaming: nonstream, assertions: [works], source: "llms/bedrock/messages/invoke_transformations/anthropic_claude_transformation.py", rationale: "Bedrock hosts no web_search server tool, so this only works because interception rewrites it before the upstream call and the agentic loop feeds the results back in native shape; a regression that short-circuits or forwards it instead yields raw text or AWS's 400", fail_before_fix: unproven} - {id: llm.messages.azure_foundry.mid_conversation_system.nonstream.cache_hit, module: llm, tier: P0, subject_endpoint: messages, route: azure_foundry, capability: mid_conversation_system, streaming: nonstream, assertions: [works, cache_hit], source: "llms/azure_ai/anthropic/messages_transformation.py", rationale: "Azure Foundry serves Claude on the native Anthropic contract, so flagged 4.8+/5 must keep mid-conversation system reminders in messages; hoisting mutates the system prefix and collapses the prompt cache (customer RCA gap)", fail_before_fix: proven} - {id: llm.messages.azure_foundry.mid_conversation_system.nonstream.works, module: llm, tier: P0, subject_endpoint: messages, route: azure_foundry, capability: mid_conversation_system, streaming: nonstream, assertions: [works], source: "llms/azure_ai/anthropic/messages_transformation.py", rationale: "Azure Foundry Claude <= 4.7 rejects role system inside messages; unflagged models must convert reminders to user turns in place (hoisting collapses the prompt cache) or every Claude Code session 400s (customer RCA gap)", fail_before_fix: proven} - {id: llm.messages.vertex.mid_conversation_system.nonstream.cache_hit, module: llm, tier: P0, subject_endpoint: messages, route: vertex, capability: mid_conversation_system, streaming: nonstream, assertions: [works, cache_hit], source: "llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py", rationale: "Vertex serves Claude on the native Anthropic contract, so flagged 4.8+/5 must keep mid-conversation system reminders in messages; hoisting mutates the system prefix and collapses the prompt cache (customer RCA gap)", fail_before_fix: proven} diff --git a/tests/integration/_support/copilot_session_prometheus_worker.py b/tests/integration/_support/copilot_session_prometheus_worker.py new file mode 100644 index 00000000000..a6b2c5d4692 --- /dev/null +++ b/tests/integration/_support/copilot_session_prometheus_worker.py @@ -0,0 +1,162 @@ +"""One SDK process for the SDK provider-401 metric cells: a provider 401 on a caller's own Copilot session, one on +the router deployment's own key, one on a key the router caller passes and one on a direct call's own key, then a +report of the deployment series on that process's default registry. Run as a script by ``run_copilot_session_worker`` so each process registers the +Prometheus metrics exactly once.""" + +from __future__ import annotations + +import asyncio +import json +import subprocess +import sys +from itertools import chain +from dataclasses import dataclass +from typing import Final + +import litellm +import openai +from litellm.integrations.prometheus import PrometheusLogger +from litellm.llms.github_copilot.per_user_auth import GithubCopilotUserSession +from prometheus_client import REGISTRY +from pydantic import TypeAdapter + +COPILOT_DEPLOYMENT: Final = "copilot-per-user-deployment" +OPERATOR_DEPLOYMENT: Final = "operator-key-deployment" +CALLER_SESSION_TOKEN: Final = "caller-copilot-session-token" +OPERATOR_KEY: Final = "rotated-operator-key" +ROUTER_CALLER_KEY: Final = "rotated-router-caller-key" +DIRECT_CALL: Final = "direct-sdk-call" +DIRECT_KEY: Final = "rotated-direct-call-key" +FAILURES: Final = "litellm_deployment_failure_responses_total" +TOTALS: Final = "litellm_deployment_total_requests_total" +LLM_FAILED: Final = "litellm_llm_api_failed_requests_metric_total" +STATE: Final = "litellm_deployment_state" +DEPLOYMENT_SERIES: Final = frozenset({FAILURES, TOTALS, LLM_FAILED, STATE}) +WORKER_TIMEOUT_SECONDS: Final = 120 + + +@dataclass(frozen=True, slots=True) +class Outcome: + deployment: str + status: int | None + + +@dataclass(frozen=True, slots=True) +class Series: + metric: str + model_id: str + api_base: str + value: float + + +@dataclass(frozen=True, slots=True) +class WorkerReport: + api_base: str + outcomes: tuple[Outcome, ...] + series: tuple[Series, ...] + + def values(self, metric: str, model_id: str) -> tuple[float, ...]: + return tuple(s.value for s in self.series if s.metric == metric and s.model_id == model_id) + + def on_the_endpoint(self, metric: str, model_id: str) -> tuple[float, ...]: + return tuple( + s.value for s in self.series if (s.metric, s.model_id, s.api_base) == (metric, model_id, self.api_base) + ) + + def outside_the_router(self, metric: str) -> tuple[float, ...]: + return tuple(s.value for s in self.series if s.metric == metric and not s.model_id) + + def on_a_one_off_deployment(self, metric: str) -> tuple[float, ...]: + named_ids: Final = ("", COPILOT_DEPLOYMENT, OPERATOR_DEPLOYMENT) + return tuple(s.value for s in self.series if s.metric == metric and s.model_id not in named_ids) + + +_REPORT: Final = TypeAdapter(WorkerReport) + + +def run_copilot_session_worker(api_base: str) -> WorkerReport: + completed: Final = subprocess.run( + (sys.executable, "-P", __file__, api_base), + capture_output=True, + text=True, + timeout=WORKER_TIMEOUT_SECONDS, + check=False, + ) + assert completed.returncode == 0, completed.stderr[-3000:] + return _REPORT.validate_json(completed.stdout.splitlines()[-1]) + + +def _deployment_series() -> tuple[Series, ...]: + samples: Final = tuple(chain.from_iterable(family.samples for family in REGISTRY.collect())) + return tuple( + Series( + metric=sample.name, + model_id=str(sample.labels.get("model_id", "")), + api_base=str(sample.labels.get("api_base", "")), + value=float(sample.value), + ) + for sample in samples + if sample.name in DEPLOYMENT_SERIES + ) + + +def _router(api_base: str) -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": COPILOT_DEPLOYMENT, + "litellm_params": { + "model": "github_copilot/gpt-4o", + "litellm_credential_name": "copilot-per-user", + "github_copilot_auth_type": "per_user_oauth", + }, + "model_info": {"id": COPILOT_DEPLOYMENT}, + }, + { + "model_name": OPERATOR_DEPLOYMENT, + "litellm_params": {"model": "openai/gpt-5.6", "api_base": api_base, "api_key": OPERATOR_KEY}, + "model_info": {"id": OPERATOR_DEPLOYMENT}, + }, + ], + num_retries=0, + disable_cooldowns=True, + ) + + +async def _attempt(router: litellm.Router, model: str, **extra: object) -> Outcome: + try: + await router.acompletion(model=model, messages=[{"role": "user", "content": f"{model} worker"}], **extra) + except openai.APIStatusError as error: + return Outcome(deployment=model, status=error.status_code) + return Outcome(deployment=model, status=None) + + +async def _direct_attempt(api_base: str) -> Outcome: + try: + await litellm.acompletion( + model="openai/gpt-5.6", + api_base=api_base, + api_key=DIRECT_KEY, + messages=[{"role": "user", "content": f"{DIRECT_CALL} worker"}], + ) + except openai.APIStatusError as error: + return Outcome(deployment=DIRECT_CALL, status=error.status_code) + return Outcome(deployment=DIRECT_CALL, status=None) + + +async def main(api_base: str) -> None: + litellm.callbacks = [PrometheusLogger()] # test-quality-ok: this worker process exits right after the report + session: Final = GithubCopilotUserSession(token=CALLER_SESSION_TOKEN, api_base=api_base) + router: Final = _router(api_base) + outcomes: Final = ( + await _attempt(router, COPILOT_DEPLOYMENT, github_copilot_user_session=session), + await _attempt(router, OPERATOR_DEPLOYMENT), + await _attempt(router, OPERATOR_DEPLOYMENT, api_key=ROUTER_CALLER_KEY), + await _direct_attempt(api_base), + ) + report: Final = WorkerReport(api_base=api_base, outcomes=outcomes, series=_deployment_series()) + sys.stdout.write(json.dumps(_REPORT.dump_python(report)) + "\n") + + +if __name__ == "__main__": + asyncio.run(main(sys.argv[1])) diff --git a/tests/integration/observability/test_prometheus_provider_auth_failure.py b/tests/integration/observability/test_prometheus_provider_auth_failure.py new file mode 100644 index 00000000000..13d264ca6b0 --- /dev/null +++ b/tests/integration/observability/test_prometheus_provider_auth_failure.py @@ -0,0 +1,1074 @@ +from __future__ import annotations + +import asyncio +import json +import threading +import uuid +from collections import Counter +from collections.abc import Iterator, Sequence +from dataclasses import dataclass +from hashlib import sha256 +from pathlib import Path +from typing import Final + +import httpx +import openai +import pytest +import yaml +from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment +from integration._support.database import write_rows +from integration._support.mcp import ENTRY_POINTS, McpCaller, McpPeer, echo_tool, register_mcp, scripted_peer +from integration._support.process import owned_proxy +from integration._support.prometheus_series import Sample, scrape +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +RUN: Final = uuid.uuid4().hex[:8] +REJECTS: Final = "provider-rejects-credential" +HELD_REJECTS: Final = "provider-rejects-credential-when-released" +KEY_NOT_VALID: Final = "provider-key-not-valid" +RATE_LIMITS: Final = "provider-rate-limits" +RECOVERS: Final = "provider-rejects-once" +HEALTHY: Final = "provider-healthy" +AGENT_365_GUARDRAIL: Final = f"agent-365-{RUN}" +CALLER_OAUTH_TOKEN: Final = f"sk-ant-oat01-caller-{RUN}" +REVOKED_GATEWAY_KEY: Final = f"revoked-gateway-key-{RUN}" +UNKNOWN_ORGANIZATION: Final = f"org-unknown-{RUN}" +REVOKED_STORED_KEY: Final = f"sk-revoked-stored-credential-{RUN}" +AZURE_API_VERSION: Final = "2024-02-01" +CALLER_AZURE_CELLS: Final = ( + ("chat", False), + ("embeddings", False), + ("responses", False), + ("messages", False), + ("responses", True), +) +CALLER_CREDENTIAL_CELLS: Final = ( + ("caller-watsonx-token", "token"), + ("caller-gigachat-access-token", "gigachat_access_token"), +) +TOKEN_EXCHANGE_CELLS: Final = ("token-exchange", "token-exchange-credential") +KEYLESS_CELLS: Final = ("keyless", "messages-keyless") +TOKEN_EXCHANGE_SETTINGS: Final = { + "token_exchange_endpoint": "https://login.microsoftonline.com/00000000-0000-4000-8000-0000000c0b01/oauth2/v2.0/token", + "client_id": "synthetic-obo-client-id", + "client_secret": "synthetic-obo-client-secret", +} + +FAILURES: Final = "litellm_deployment_failure_responses_total" +TOTALS: Final = "litellm_deployment_total_requests_total" +LLM_FAILED: Final = "litellm_llm_api_failed_requests_metric_total" +STATE: Final = "litellm_deployment_state" +PROXY_FAILED: Final = "litellm_proxy_failed_requests_metric_total" +SUCCESSES: Final = "litellm_deployment_success_responses_total" +DEPLOYMENT_SERIES: Final = (FAILURES, TOTALS, LLM_FAILED, STATE) + +ENDPOINT_CELLS: Final = ( + "chat", + "chat-stream", + "messages", + "messages-stream", + "responses", + "responses-stream", + "embeddings", +) +REJECTED_CELLS: Final = ( + *ENDPOINT_CELLS, + "sdk-sync", + "sdk-async", + "repeated", + "fallback", + "burst", + "caller-api-key", + "caller-extra-headers", + "caller-headers", + "caller-provider-specific-header", + "plain-extra-headers", + "plain-headers", + "plain-provider-specific-header", + "forwarded-provider-header", + "forwarded-plain-header", + "forwarded-user-information", +) +GATED_CELLS: Final = ("invalid-key", "missing-key", "expired-key", "blocked-key", "no-model-access") +CALLER_SCOPE_CELLS: Final = ("forwarded-gateway-key", "caller-organization", "forwarded-organization") +HEALTHY_CELLS: Final = (*GATED_CELLS, *CALLER_SCOPE_CELLS, "healthy", "key-rpm-limit") +GUARDRAIL_ROWS: Final = ( + "SELECT COALESCE(jsonb_path_query_first(metadata, " + "'$.guardrail_information[*] ? (@.guardrail_name == $name).guardrail_status', " + "jsonb_build_object('name', %s::text)) #>> '{}', 'none') AS status " + 'FROM "LiteLLM_SpendLogs" WHERE api_key = %s AND call_type = %s' +) + + +def name(cell: str) -> str: + return f"prometheus-{cell}-{RUN}" + + +def identity(cell: str) -> str: + return f"{name(cell)}-deployment" + + +def _error(status: int, message: str, kind: str, code: str) -> Reply: + return Reply(status=status, body=json.dumps({"error": {"message": message, "type": kind, "code": code}}).encode()) + + +def _completion() -> Reply: + return Reply( + body=json.dumps( + { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}, + } + ).encode() + ) + + +class Provider: + """Answers by upstream model name. `HELD_REJECTS` waits for `release` before rejecting the credential.""" + + def __init__(self, release: threading.Event) -> None: + self._lock: Final = threading.Lock() + self._seen: Final[Counter[str]] = Counter() + self._release: Final = release + + def __call__(self, request: Request) -> Reply: + sent: Final = json.loads(request.body) + model: Final = str(sent.get("model") or sent.get("model_id") or request.target.split("/")[2]) + if request.headers.get("x-portkey-api-key") == REVOKED_GATEWAY_KEY: + return _error(401, "Invalid gateway API key", "invalid_request_error", "invalid_api_key") + if request.headers.get("authorization") == f"Bearer {REVOKED_STORED_KEY}": + return _error(401, "Incorrect API key provided: sk-revok*****", "invalid_request_error", "invalid_api_key") + if request.headers.get("openai-organization") == UNKNOWN_ORGANIZATION: + return _error( + 401, + "OpenAI-Organization header should match organization for API key", + "invalid_request_error", + "mismatched_organization", + ) + with self._lock: + self._seen[model] += 1 + attempt: Final = self._seen[model] + if model == HELD_REJECTS: + assert self._release.wait(timeout=30), "held provider reply was never released" + if model in (REJECTS, HELD_REJECTS) or (model == RECOVERS and attempt == 1): + return _error(401, "Incorrect API key provided: sk-rotat*****", "invalid_request_error", "invalid_api_key") + if model == KEY_NOT_VALID: + return _error(400, "API key not valid. Please pass a valid API key.", "invalid_request_error", "400") + if model == RATE_LIMITS: + return _error(429, "Rate limit reached for requests", "rate_limit_error", "rate_limit_exceeded") + return _completion() + + +def deployment(cell: str, upstream_model: str, api_base: str) -> dict[str, JsonValue]: + return { + "model_name": name(cell), + "litellm_params": { + "model": f"openai/{upstream_model}", + "api_key": f"sk-deployment-credential-{cell}", + "api_base": api_base, + }, + "model_info": {"id": identity(cell)}, + } + + +def anthropic_deployment(cell: str, api_base: str) -> dict[str, JsonValue]: + return { + "model_name": name(cell), + "litellm_params": { + "model": f"anthropic/{REJECTS}", + "api_key": f"sk-ant-api03-deployment-credential-{cell}", + "api_base": api_base, + }, + "model_info": {"id": identity(cell)}, + } + + +def bedrock_deployment(cell: str, opt_in: dict[str, JsonValue]) -> dict[str, JsonValue]: + return { + "model_name": name(cell), + "litellm_params": { + "model": f"bedrock/converse/{REJECTS}", + "aws_access_key_id": "AKIADEPLOYMENTCREDENTIAL", + "aws_secret_access_key": f"deployment-credential-{cell}", + "aws_region_name": "us-east-1", + **opt_in, + }, + "model_info": {"id": identity(cell)}, + } + + +def keyless_anthropic_deployment(cell: str, api_base: str) -> dict[str, JsonValue]: + return { + "model_name": name(cell), + "litellm_params": {"model": f"anthropic/{REJECTS}", "api_base": api_base}, + "model_info": {"id": identity(cell)}, + } + + +def azure_deployment(cell: str, api_base: str, credential: dict[str, JsonValue]) -> dict[str, JsonValue]: + return { + "model_name": name(cell), + "litellm_params": { + "model": f"azure/{REJECTS}", + "api_base": api_base, + "api_version": AZURE_API_VERSION, + **credential, + }, + "model_info": {"id": identity(cell)}, + } + + +def caller_credential_deployment(cell: str, field: str, api_base: str) -> dict[str, JsonValue]: + provider: Final = "watsonx" if field == "token" else "gigachat" + project: Final[dict[str, JsonValue]] = {"project_id": f"project-{RUN}"} if provider == "watsonx" else {} + return { + "model_name": name(cell), + "litellm_params": { + "model": f"{provider}/{REJECTS}", + "api_base": api_base, + field: f"deployment-credential-{cell}", + **project, + }, + "model_info": {"id": identity(cell)}, + } + + +def token_exchange_deployment(cell: str) -> dict[str, JsonValue]: + settings: Final[dict[str, JsonValue]] = ( + {"litellm_credential_name": f"m365-obo-{RUN}"} if cell.endswith("-credential") else {**TOKEN_EXCHANGE_SETTINGS} + ) + return { + "model_name": name(cell), + "litellm_params": {"model": "microsoft_365_copilot/chat", **settings}, + "model_info": {"id": identity(cell)}, + } + + +def caller_azure_cell(opt_in: str, route: str, inside_extra_body: bool) -> str: + return f"{opt_in}-caller-azure-token-{route}{'-extra-body' if inside_extra_body else ''}" + + +def config(directory: Path, upstream: Wire) -> Path: + api_base: Final = f"{upstream.url}/v1" + configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + configuration["model_list"] = [ + *configuration["model_list"], + *(deployment(cell, REJECTS, api_base) for cell in REJECTED_CELLS), + *(deployment(cell, HEALTHY, api_base) for cell in HEALTHY_CELLS), + deployment("held", HELD_REJECTS, api_base), + deployment("key-not-valid", KEY_NOT_VALID, api_base), + deployment("rate-limited", RATE_LIMITS, api_base), + deployment("recovers", RECOVERS, api_base), + anthropic_deployment("anthropic-own-key", upstream.url), + anthropic_deployment("anthropic-sdk-own-key", upstream.url), + anthropic_deployment("caller-oauth", upstream.url), + *(keyless_anthropic_deployment(cell, upstream.url) for cell in KEYLESS_CELLS), + bedrock_deployment( + "deployment-caller-bedrock-endpoint", + {"configurable_clientside_auth_params": ["aws_bedrock_runtime_endpoint"]}, + ), + azure_deployment("azure-own-token", upstream.url, {"azure_ad_token": "deployment-entra-token-azure-own-token"}), + *( + azure_deployment( + caller_azure_cell("deployment", route, inside_extra_body), + upstream.url, + {"configurable_clientside_auth_params": ["azure_ad_token"]}, + ) + for route, inside_extra_body in CALLER_AZURE_CELLS + ), + *(caller_credential_deployment(cell, field, upstream.url) for cell, field in CALLER_CREDENTIAL_CELLS), + *(token_exchange_deployment(cell) for cell in TOKEN_EXCHANGE_CELLS), + { + "model_name": name("caller-api-base"), + "litellm_params": { + "model": f"openai/{REJECTS}", + "api_key": "sk-deployment-credential-caller-api-base", + "api_base": api_base, + "configurable_clientside_auth_params": ["api_base"], + }, + "model_info": {"id": identity("caller-api-base")}, + }, + { + "model_name": name("caller-credential"), + "litellm_params": { + "model": "github_copilot/gpt-4o", + "litellm_credential_name": f"copilot-per-user-{RUN}", + "github_copilot_auth_type": "per_user_oauth", + }, + "model_info": {"id": identity("caller-credential")}, + }, + { + "model_name": name("caller-credential-name"), + "litellm_params": { + "model": f"openai/{HEALTHY}", + "api_base": api_base, + "litellm_credential_name": f"deployment-key-{RUN}", + }, + "model_info": {"id": identity("caller-credential-name")}, + }, + ] + configuration["credential_list"] = [ + *configuration.get("credential_list", []), + {"credential_name": f"m365-obo-{RUN}", "credential_values": {**TOKEN_EXCHANGE_SETTINGS}, "credential_info": {}}, + { + "credential_name": f"deployment-key-{RUN}", + "credential_values": {"api_key": "sk-deployment-credential-caller-credential-name"}, + "credential_info": {}, + }, + { + "credential_name": f"revoked-key-{RUN}", + "credential_values": {"api_key": REVOKED_STORED_KEY}, + "credential_info": {}, + }, + ] + configuration["guardrails"] = [ + { + "guardrail_name": AGENT_365_GUARDRAIL, + "litellm_params": { + "guardrail": "agent_365", + "mode": "pre_mcp_call", + "default_on": True, + "tenant_id": "00000000-0000-4000-8000-0000000a3650", + "client_id": "synthetic-client-id", + "client_secret": "synthetic-client-secret", + }, + } + ] + configuration["litellm_settings"]["callbacks"] = ["prometheus"] + configuration["litellm_settings"]["model_group_settings"] = { + "forward_client_headers_to_llm_api": [ + name("forwarded-provider-header"), + name("forwarded-plain-header"), + name("forwarded-gateway-key"), + name("forwarded-user-information"), + ] + } + configuration["litellm_settings"]["add_user_information_to_llm_headers"] = True + configuration["general_settings"]["forward_llm_provider_auth_headers"] = True + configuration["general_settings"]["forward_openai_org_id"] = True + configuration.setdefault("router_settings", {}).update({"num_retries": 0, "disable_cooldowns": True}) + path: Final = directory / "prometheus-provider-auth-failure.yaml" + path.write_text(yaml.safe_dump(configuration)) + return path + + +@dataclass(frozen=True, slots=True) +class Rig: + gateway: Gateway + upstream: Wire + + @property + def api_base(self) -> str: + return f"{self.upstream.url}/v1" + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + directory: Final = tmp_path_factory.mktemp("prometheus-provider-auth-failure") + released: Final = threading.Event() + released.set() + with ( + wire_server(Provider(released)) as upstream, + gateway_from_environment() as shared, + owned_proxy( + shared, + directory, + {"PROXY_CONFIG_RELOAD_INTERVAL_SECONDS": "2"}, + config=config(directory, upstream), + remove_environment=("ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN"), + ) as owned, + ): + yield Rig(owned, upstream) + + +def opted_in_config(directory: Path, upstream: Wire) -> Path: + configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + configuration["model_list"] = [ + *configuration["model_list"], + *( + azure_deployment(caller_azure_cell("proxy-wide", route, inside_extra_body), upstream.url, {}) + for route, inside_extra_body in CALLER_AZURE_CELLS + ), + bedrock_deployment("proxy-wide-caller-bedrock-endpoint", {}), + ] + configuration["general_settings"]["allow_client_side_credentials"] = True + configuration["litellm_settings"]["callbacks"] = ["prometheus"] + configuration.setdefault("router_settings", {}).update({"num_retries": 0, "disable_cooldowns": True}) + path: Final = directory / "prometheus-caller-credentials-opted-in.yaml" + path.write_text(yaml.safe_dump(configuration)) + return path + + +@pytest.fixture(scope="module") +def opted_in(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + directory: Final = tmp_path_factory.mktemp("prometheus-caller-credentials-opted-in") + released: Final = threading.Event() + released.set() + with ( + wire_server(Provider(released)) as upstream, + gateway_from_environment() as shared, + owned_proxy(shared, directory, {}, config=opted_in_config(directory, upstream)) as owned, + ): + yield Rig(owned, upstream) + + +def series(samples: Sequence[Sample], metric: str, **labels: str) -> tuple[Sample, ...]: + return tuple( + sample + for sample in samples + if sample.name == metric and all(sample.labels.get(label) == value for label, value in labels.items()) + ) + + +def values(samples: Sequence[Sample], metric: str, **labels: str) -> list[float]: + return [sample.value for sample in series(samples, metric, **labels)] + + +def until(gateway: Gateway, metric: str, **labels: str) -> tuple[Sample, ...]: + return eventually(lambda: scrape(gateway), lambda samples: bool(series(samples, metric, **labels)), 30) + + +def messages(cell: str) -> list[JsonValue]: + return [{"role": "user", "content": f"{cell} {uuid.uuid4().hex}"}] + + +def send(gateway: Gateway, cell: str, model: str | None = None, key: str | None = None) -> httpx.Response: + asked: Final = model or name(cell) + stream: Final = cell.endswith("-stream") + if cell.startswith("messages"): + anthropic: Final = {"model": asked, "max_tokens": 16, "messages": messages(cell), "stream": stream} + return gateway.request("POST", "/v1/messages", anthropic, key=key) + if cell.startswith("responses"): + return gateway.request("POST", "/v1/responses", {"model": asked, "input": cell, "stream": stream}, key=key) + if cell == "embeddings": + return gateway.request("POST", "/v1/embeddings", {"model": asked, "input": cell}, key=key) + return gateway.request( + "POST", "/v1/chat/completions", {"model": asked, "messages": messages(cell), "stream": stream}, key=key + ) + + +def assert_reached_with_the_deployment_credential(rig: Rig, cell: str, attempts: int = 1) -> None: + reached: Final = rig.upstream.drain() + assert [request.headers["authorization"] for request in reached] == [ + f"Bearer sk-deployment-credential-{cell}" + ] * attempts, [request.target for request in reached] + + +def recorded(samples: Sequence[Sample], **labels: str) -> list[tuple[str, dict[str, str], float]]: + return [(sample.name, dict(sample.labels), sample.value) for sample in series_of(samples, **labels)] + + +def series_of(samples: Sequence[Sample], **labels: str) -> tuple[Sample, ...]: + return tuple( + sample + for sample in samples + if sample.name.startswith("litellm_") + and all(sample.labels.get(label) == value for label, value in labels.items()) + ) + + +def assert_kept_out_of_the_deployment_metrics(rig: Rig, cell: str) -> None: + samples: Final = until(rig.gateway, PROXY_FAILED, requested_model=name(cell)) + assert {metric: values(samples, metric, model_id=identity(cell)) for metric in DEPLOYMENT_SERIES} == { + metric: [] for metric in DEPLOYMENT_SERIES + }, recorded(samples, model_id=identity(cell)) + + +def assert_recorded_as_a_deployment_failure(rig: Rig, cell: str, status: str, count: float = 1.0) -> None: + samples: Final = until(rig.gateway, FAILURES, model_id=identity(cell)) + observed: Final = recorded(samples, model_id=identity(cell)) + assert values(samples, FAILURES, model_id=identity(cell), exception_status=status) == [count], observed + assert values(samples, TOTALS, model_id=identity(cell)) == [count], observed + assert values(samples, LLM_FAILED, model_id=identity(cell)) == [count], observed + assert values(samples, STATE, model_id=identity(cell)) == [1.0], observed + + +@pytest.mark.parametrize("cell", ENDPOINT_CELLS) +def test_a_provider_401_on_every_endpoint_counts_as_a_deployment_failure_and_a_partial_outage( + rig: Rig, cell: str +) -> None: + rig.upstream.drain() + response: Final = send(rig.gateway, cell) + assert response.status_code == 401, response.text + assert_reached_with_the_deployment_credential(rig, cell) + assert_recorded_as_a_deployment_failure(rig, cell, "401") + + +def test_a_provider_401_seen_by_the_sync_openai_client_is_recorded(rig: Rig) -> None: + rig.upstream.drain() + with ( + openai.OpenAI( + base_url=f"{rig.gateway.client.base_url}/v1", + api_key=rig.gateway.key, + max_retries=0, + http_client=httpx.Client(trust_env=False), + ) as client, + pytest.raises(openai.AuthenticationError, match="Incorrect API key provided"), + ): + client.chat.completions.create(model=name("sdk-sync"), messages=[{"role": "user", "content": "sync"}]) + assert_reached_with_the_deployment_credential(rig, "sdk-sync") + assert_recorded_as_a_deployment_failure(rig, "sdk-sync", "401") + + +def test_a_provider_401_seen_by_the_async_openai_client_is_recorded(rig: Rig) -> None: + async def ask() -> None: + async with openai.AsyncOpenAI( + base_url=f"{rig.gateway.client.base_url}/v1", + api_key=rig.gateway.key, + max_retries=0, + http_client=httpx.AsyncClient(trust_env=False), + ) as client: + await client.chat.completions.create(model=name("sdk-async"), messages=[{"role": "user", "content": "a"}]) + + rig.upstream.drain() + with pytest.raises(openai.AuthenticationError, match="Incorrect API key provided"): + asyncio.run(ask()) + assert_reached_with_the_deployment_credential(rig, "sdk-async") + assert_recorded_as_a_deployment_failure(rig, "sdk-async", "401") + + +def test_each_repeated_provider_401_is_counted_once(rig: Rig) -> None: + rig.upstream.drain() + assert [send(rig.gateway, "repeated").status_code for _ in range(3)] == [401, 401, 401] + assert_reached_with_the_deployment_credential(rig, "repeated", attempts=3) + assert_recorded_as_a_deployment_failure(rig, "repeated", "401", count=3.0) + + +def test_a_provider_401_rescued_by_a_fallback_still_counts_against_the_rejected_deployment(rig: Rig) -> None: + rig.upstream.drain() + response: Final = rig.gateway.request( + "POST", + "/v1/chat/completions", + {"model": name("fallback"), "messages": messages("fallback"), "fallbacks": [name("healthy")]}, + ) + assert response.status_code == 200, response.text + assert [request.headers["authorization"] for request in rig.upstream.drain()] == [ + "Bearer sk-deployment-credential-fallback", + "Bearer sk-deployment-credential-healthy", + ] + assert_recorded_as_a_deployment_failure(rig, "fallback", "401") + rescued: Final = eventually( + lambda: scrape(rig.gateway), lambda samples: values(samples, SUCCESSES, model_id=identity("healthy")) == [1.0] + ) + assert values(rescued, FAILURES, model_id=identity("healthy")) == [] + + +def test_a_provider_auth_message_on_a_non_401_status_is_recorded(rig: Rig) -> None: + rig.upstream.drain() + response: Final = send(rig.gateway, "key-not-valid") + assert response.status_code == 400, response.text + assert_reached_with_the_deployment_credential(rig, "key-not-valid") + assert_recorded_as_a_deployment_failure(rig, "key-not-valid", "400") + + +def test_a_provider_429_is_recorded_as_before(rig: Rig) -> None: + rig.upstream.drain() + response: Final = send(rig.gateway, "rate-limited") + assert response.status_code == 429, response.text + assert_reached_with_the_deployment_credential(rig, "rate-limited") + assert_recorded_as_a_deployment_failure(rig, "rate-limited", "429") + + +def test_a_deployment_marked_down_by_a_provider_401_is_healthy_again_after_a_success(rig: Rig) -> None: + rig.upstream.drain() + assert send(rig.gateway, "recovers").status_code == 401 + assert_recorded_as_a_deployment_failure(rig, "recovers", "401") + recovered: Final = send(rig.gateway, "recovers") + assert recovered.status_code == 200, recovered.text + assert len(rig.upstream.drain()) == 2 + eventually( + lambda: scrape(rig.gateway), + lambda samples: values(samples, STATE, model_id=identity("recovers"), api_base=rig.api_base) == [0.0], + ) + + +def test_a_per_caller_credential_rejection_stays_out_of_the_deployment_metrics(rig: Rig) -> None: + rig.upstream.drain() + response: Final = send(rig.gateway, "caller-credential") + assert response.status_code == 401, response.text + assert "Connect GitHub Copilot" in response.text, response.text + assert rig.upstream.drain() == () + samples: Final = until(rig.gateway, PROXY_FAILED, requested_model=name("caller-credential")) + assert values(samples, FAILURES, model_id=identity("caller-credential")) == [] + assert values(samples, TOTALS, model_id=identity("caller-credential")) == [] + assert values(samples, LLM_FAILED, model_id=identity("caller-credential")) == [] + assert values(samples, STATE, model_id=identity("caller-credential")) == [] + + +def deployment_series(samples: Sequence[Sample]) -> frozenset[tuple[str, tuple[tuple[str, str], ...]]]: + return frozenset((sample.name, sample.identity()) for sample in samples if sample.name in DEPLOYMENT_SERIES) + + +def test_a_provider_401_on_an_api_key_the_caller_sent_stays_out_of_the_deployment_metrics(rig: Rig) -> None: + caller_key: Final = f"sk-caller-supplied-{uuid.uuid4().hex}" + before: Final = deployment_series(scrape(rig.gateway)) + rig.upstream.drain() + response: Final = rig.gateway.request( + "POST", + "/v1/chat/completions", + {"model": name("caller-api-key"), "messages": messages("caller-api-key"), "api_key": caller_key}, + ) + assert response.status_code == 401, response.text + assert [request.headers["authorization"] for request in rig.upstream.drain()] == [f"Bearer {caller_key}"] + samples: Final = until(rig.gateway, PROXY_FAILED, requested_model=name("caller-api-key")) + assert deployment_series(samples) - before == frozenset(), recorded(samples, requested_model=name("caller-api-key")) + + +def test_a_provider_401_on_a_stored_credential_the_caller_names_leaves_a_healthy_deployment_healthy(rig: Rig) -> None: + cell: Final = "caller-credential-name" + rig.upstream.drain() + assert send(rig.gateway, cell).status_code == 200 + assert_reached_with_the_deployment_credential(rig, cell) + response: Final = rig.gateway.request( + "POST", + "/v1/chat/completions", + {"model": name(cell), "messages": messages(cell), "litellm_credential_name": f"revoked-key-{RUN}"}, + ) + assert response.status_code == 401, response.text + assert [request.headers["authorization"] for request in rig.upstream.drain()] == [f"Bearer {REVOKED_STORED_KEY}"] + samples: Final = until(rig.gateway, PROXY_FAILED, requested_model=name(cell)) + observed: Final = recorded(samples, model_id=identity(cell)) + assert values(samples, FAILURES, model_id=identity(cell)) == [], observed + assert values(samples, LLM_FAILED, model_id=identity(cell)) == [], observed + assert values(samples, TOTALS, model_id=identity(cell)) == [1.0], observed + assert values(samples, STATE, model_id=identity(cell)) == [0.0], observed + + +def anthropic_messages(gateway: Gateway, cell: str, key: str, headers: dict[str, str]) -> httpx.Response: + body: Final[dict[str, JsonValue]] = {"model": name(cell), "max_tokens": 16, "messages": messages(cell)} + return gateway.request("POST", "/v1/messages", body, key=key, headers=headers) + + +def caller_header(field: str, header: str, value: str) -> JsonValue: + if field == "provider_specific_header": + return {"custom_llm_provider": "openai", "extra_headers": {header: value}} + return {header: value} + + +@pytest.mark.parametrize( + "cell,field", + [ + ("caller-extra-headers", "extra_headers"), + ("caller-headers", "headers"), + ("caller-provider-specific-header", "provider_specific_header"), + ], +) +def test_a_provider_401_on_an_authorization_header_the_caller_sends_stays_out_of_the_deployment_metrics( + rig: Rig, cell: str, field: str +) -> None: + caller_value: Final = f"Bearer sk-{cell}-{uuid.uuid4().hex}" + rig.upstream.drain() + response: Final = rig.gateway.request( + "POST", + "/v1/chat/completions", + {"model": name(cell), "messages": messages(cell), field: caller_header(field, "Authorization", caller_value)}, + ) + assert response.status_code == 401, response.text + assert [request.headers["authorization"] for request in rig.upstream.drain()] == [caller_value] + samples: Final = until(rig.gateway, PROXY_FAILED, requested_model=name(cell)) + observed: Final = recorded(samples, model_id=identity(cell)) + assert values(samples, FAILURES, model_id=identity(cell)) == [], observed + assert values(samples, TOTALS, model_id=identity(cell)) == [], observed + assert values(samples, LLM_FAILED, model_id=identity(cell)) == [], observed + assert values(samples, STATE, model_id=identity(cell)) == [], observed + + +@pytest.mark.parametrize( + "cell,field", + [ + ("plain-extra-headers", "extra_headers"), + ("plain-headers", "headers"), + ("plain-provider-specific-header", "provider_specific_header"), + ], +) +def test_a_provider_401_on_the_deployments_own_key_stays_out_when_the_caller_sends_a_plain_header( + rig: Rig, cell: str, field: str +) -> None: + trace: Final = f"{cell}-{uuid.uuid4().hex}" + rig.upstream.drain() + response: Final = rig.gateway.request( + "POST", + "/v1/chat/completions", + {"model": name(cell), "messages": messages(cell), field: caller_header(field, "X-Caller-Trace", trace)}, + ) + assert response.status_code == 401, response.text + reached: Final = rig.upstream.drain() + assert [(request.headers["authorization"], request.headers.get("x-caller-trace")) for request in reached] == [ + (f"Bearer sk-deployment-credential-{cell}", trace) + ] + assert_kept_out_of_the_deployment_metrics(rig, cell) + + +def test_a_provider_401_on_a_provider_credential_header_the_proxy_forwards_stays_out_of_the_deployment_metrics( + rig: Rig, +) -> None: + caller_key: Final = f"AIza-forwarded-{uuid.uuid4().hex}" + rig.upstream.drain() + response: Final = rig.gateway.request( + "POST", + "/v1/chat/completions", + {"model": name("forwarded-provider-header"), "messages": messages("forwarded-provider-header")}, + headers={"x-goog-api-key": caller_key}, + ) + assert response.status_code == 401, response.text + assert [request.headers.get("x-goog-api-key") for request in rig.upstream.drain()] == [caller_key] + samples: Final = until(rig.gateway, PROXY_FAILED, requested_model=name("forwarded-provider-header")) + observed: Final = recorded(samples, model_id=identity("forwarded-provider-header")) + assert values(samples, FAILURES, model_id=identity("forwarded-provider-header")) == [], observed + assert values(samples, TOTALS, model_id=identity("forwarded-provider-header")) == [], observed + assert values(samples, LLM_FAILED, model_id=identity("forwarded-provider-header")) == [], observed + assert values(samples, STATE, model_id=identity("forwarded-provider-header")) == [], observed + + +def test_a_provider_401_on_the_deployments_own_key_stays_out_when_the_proxy_forwards_a_plain_header(rig: Rig) -> None: + trace: Final = f"forwarded-plain-header-{uuid.uuid4().hex}" + rig.upstream.drain() + response: Final = rig.gateway.request( + "POST", + "/v1/chat/completions", + {"model": name("forwarded-plain-header"), "messages": messages("forwarded-plain-header")}, + headers={"x-caller-trace": trace}, + ) + assert response.status_code == 401, response.text + assert [ + (request.headers["authorization"], request.headers.get("x-caller-trace")) for request in rig.upstream.drain() + ] == [("Bearer sk-deployment-credential-forwarded-plain-header", trace)] + assert_kept_out_of_the_deployment_metrics(rig, "forwarded-plain-header") + + +def test_a_provider_401_on_the_deployments_own_key_counts_when_the_proxy_forwards_only_its_own_user_information( + rig: Rig, +) -> None: + rig.upstream.drain() + response: Final = send(rig.gateway, "forwarded-user-information") + assert response.status_code == 401, response.text + reached: Final = rig.upstream.drain() + assert [ + (request.headers["authorization"], "x-litellm-user_api_key_hash" in request.headers) for request in reached + ] == [("Bearer sk-deployment-credential-forwarded-user-information", True)] + assert_recorded_as_a_deployment_failure(rig, "forwarded-user-information", "401") + + +def test_a_provider_401_on_a_gateway_key_the_proxy_forwards_stays_out_of_a_healthy_deployments_metrics( + rig: Rig, +) -> None: + rig.upstream.drain() + response: Final = rig.gateway.request( + "POST", + "/v1/chat/completions", + {"model": name("forwarded-gateway-key"), "messages": messages("forwarded-gateway-key")}, + headers={"x-portkey-api-key": REVOKED_GATEWAY_KEY}, + ) + assert response.status_code == 401, response.text + assert [ + (request.headers["authorization"], request.headers.get("x-portkey-api-key")) for request in rig.upstream.drain() + ] == [("Bearer sk-deployment-credential-forwarded-gateway-key", REVOKED_GATEWAY_KEY)] + assert_kept_out_of_the_deployment_metrics(rig, "forwarded-gateway-key") + + +@pytest.mark.parametrize( + "cell,body,headers", + [ + ("caller-organization", {"organization": UNKNOWN_ORGANIZATION}, {}), + ("forwarded-organization", {}, {"OpenAI-Organization": UNKNOWN_ORGANIZATION}), + ], +) +def test_a_provider_401_on_an_organization_the_caller_sends_stays_out_of_a_healthy_deployments_metrics( + rig: Rig, cell: str, body: dict[str, JsonValue], headers: dict[str, str] +) -> None: + rig.upstream.drain() + response: Final = rig.gateway.request( + "POST", "/v1/chat/completions", {"model": name(cell), "messages": messages(cell), **body}, headers=headers + ) + assert response.status_code == 401, response.text + assert [ + (request.headers["authorization"], request.headers.get("openai-organization")) + for request in rig.upstream.drain() + ] == [(f"Bearer sk-deployment-credential-{cell}", UNKNOWN_ORGANIZATION)] + assert_kept_out_of_the_deployment_metrics(rig, cell) + + +@pytest.mark.parametrize( + ("cell", "headers"), + [ + ("anthropic-own-key", {}), + ("anthropic-sdk-own-key", {"anthropic-version": "2023-06-01", "anthropic-beta": "tools-2024-04-04"}), + ], + ids=["no-client-headers", "anthropic-sdk-headers"], +) +def test_a_provider_401_on_an_anthropic_deployments_own_key_is_recorded( + rig: Rig, cell: str, headers: dict[str, str] +) -> None: + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(models=[name(cell)]) + rig.upstream.drain() + response: Final = anthropic_messages(rig.gateway, cell, key, headers) + assert response.status_code == 401, response.text + assert [request.headers.get("x-api-key") for request in rig.upstream.drain()] == [ + f"sk-ant-api03-deployment-credential-{cell}" + ] + assert_recorded_as_a_deployment_failure(rig, cell, "401") + + +def test_a_provider_401_on_an_anthropic_oauth_token_the_caller_sent_stays_out_of_the_deployment_metrics( + rig: Rig, +) -> None: + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(models=[name("caller-oauth")]) + rig.upstream.drain() + response: Final = anthropic_messages( + rig.gateway, + "caller-oauth", + CALLER_OAUTH_TOKEN, + {"x-litellm-api-key": key}, + ) + assert response.status_code == 401, response.text + assert [request.headers.get("authorization") for request in rig.upstream.drain()] == [ + f"Bearer {CALLER_OAUTH_TOKEN}" + ] + samples: Final = until(rig.gateway, PROXY_FAILED, requested_model=name("caller-oauth")) + assert {metric: values(samples, metric, model_id=identity("caller-oauth")) for metric in DEPLOYMENT_SERIES} == { + metric: [] for metric in DEPLOYMENT_SERIES + }, recorded(samples, requested_model=name("caller-oauth")) + + +def test_a_provider_401_on_an_azure_deployments_own_ad_token_is_recorded(rig: Rig) -> None: + rig.upstream.drain() + response: Final = send(rig.gateway, "azure-own-token") + assert response.status_code == 401, response.text + assert [request.headers.get("authorization") for request in rig.upstream.drain()] == [ + "Bearer deployment-entra-token-azure-own-token" + ] + assert_recorded_as_a_deployment_failure(rig, "azure-own-token", "401") + + +def caller_azure_request(cell: str, route: str, credential: dict[str, JsonValue]) -> tuple[str, dict[str, JsonValue]]: + if route == "embeddings": + return "/v1/embeddings", {"model": name(cell), "input": cell, **credential} + if route == "responses": + return "/v1/responses", {"model": name(cell), "input": cell, **credential} + if route == "messages": + return "/v1/messages", {"model": name(cell), "max_tokens": 16, "messages": messages(cell), **credential} + return "/v1/chat/completions", {"model": name(cell), "messages": messages(cell), **credential} + + +@pytest.mark.parametrize( + ("route", "inside_extra_body"), + CALLER_AZURE_CELLS, + ids=[f"{route}{'-extra-body' if inside_extra_body else ''}" for route, inside_extra_body in CALLER_AZURE_CELLS], +) +@pytest.mark.parametrize("opt_in", ("deployment", "proxy-wide")) +def test_a_provider_401_on_an_azure_ad_token_the_caller_sent_stays_out_of_the_deployment_metrics( + request: pytest.FixtureRequest, opt_in: str, route: str, inside_extra_body: bool +) -> None: + opted: Final[Rig] = request.getfixturevalue("rig" if opt_in == "deployment" else "opted_in") + cell: Final = caller_azure_cell(opt_in, route, inside_extra_body) + token: Final = f"caller-entra-token-{uuid.uuid4().hex}" + credential: Final[dict[str, JsonValue]] = ( + {"extra_body": {"azure_ad_token": token}} if inside_extra_body else {"azure_ad_token": token} + ) + opted.upstream.drain() + response: Final = opted.gateway.request("POST", *caller_azure_request(cell, route, credential)) + assert response.status_code == 401, response.text + assert [reached.headers.get("authorization") for reached in opted.upstream.drain()] == [ + None if inside_extra_body else f"Bearer {token}" + ] + samples: Final = until(opted.gateway, PROXY_FAILED, requested_model=name(cell)) + assert {metric: values(samples, metric, model_id=identity(cell)) for metric in DEPLOYMENT_SERIES} == { + metric: [] for metric in DEPLOYMENT_SERIES + }, recorded(samples, requested_model=name(cell)) + + +@pytest.mark.parametrize(("cell", "field"), CALLER_CREDENTIAL_CELLS, ids=[cell for cell, _ in CALLER_CREDENTIAL_CELLS]) +def test_a_provider_401_on_a_provider_credential_the_caller_sent_stays_out_of_the_deployment_metrics( + rig: Rig, cell: str, field: str +) -> None: + caller_credential: Final = f"caller-credential-{uuid.uuid4().hex}" + rig.upstream.drain() + response: Final = rig.gateway.request( + "POST", "/v1/chat/completions", {"model": name(cell), "messages": messages(cell), field: caller_credential} + ) + assert response.status_code == 401, response.text + assert [reached.headers.get("authorization") for reached in rig.upstream.drain()] == [f"Bearer {caller_credential}"] + samples: Final = until(rig.gateway, PROXY_FAILED, requested_model=name(cell)) + assert {metric: values(samples, metric, model_id=identity(cell)) for metric in DEPLOYMENT_SERIES} == { + metric: [] for metric in DEPLOYMENT_SERIES + }, recorded(samples, requested_model=name(cell)) + + +@pytest.mark.parametrize("cell", TOKEN_EXCHANGE_CELLS) +def test_a_token_exchange_deployment_rejecting_a_caller_without_an_identity_token_stays_out_of_the_deployment_metrics( + rig: Rig, cell: str +) -> None: + rig.upstream.drain() + response: Final = send(rig.gateway, cell) + assert response.status_code == 401, response.text + assert "IdP-issued access token" in response.text, response.text + assert rig.upstream.drain() == () + samples: Final = until(rig.gateway, PROXY_FAILED, requested_model=name(cell)) + assert {metric: values(samples, metric, model_id=identity(cell)) for metric in DEPLOYMENT_SERIES} == { + metric: [] for metric in DEPLOYMENT_SERIES + }, recorded(samples, requested_model=name(cell)) + + +@pytest.mark.parametrize("cell", KEYLESS_CELLS) +def test_a_missing_deployment_key_reported_before_the_provider_is_called_stays_out_of_the_deployment_metrics( + rig: Rig, cell: str +) -> None: + rig.upstream.drain() + response: Final = send(rig.gateway, cell) + assert response.status_code == 401, response.text + assert "Missing Anthropic API Key" in response.text, response.text + assert rig.upstream.drain() == () + assert_kept_out_of_the_deployment_metrics(rig, cell) + + +@pytest.mark.parametrize("opt_in", ("deployment", "proxy-wide")) +def test_a_provider_401_at_a_bedrock_endpoint_the_caller_sends_stays_out_of_the_deployment_metrics( + request: pytest.FixtureRequest, opt_in: str +) -> None: + opted: Final[Rig] = request.getfixturevalue("rig" if opt_in == "deployment" else "opted_in") + cell: Final = f"{opt_in}-caller-bedrock-endpoint" + opted.upstream.drain() + response: Final = opted.gateway.request( + "POST", + "/v1/chat/completions", + {"model": name(cell), "messages": messages(cell), "aws_bedrock_runtime_endpoint": opted.upstream.url}, + ) + assert response.status_code == 401, response.text + assert [reached.target for reached in opted.upstream.drain()] == [f"/model/{REJECTS}/converse"] + assert_kept_out_of_the_deployment_metrics(opted, cell) + + +def test_a_provider_401_at_an_api_base_the_caller_sends_counts_against_a_one_off_deployment(rig: Rig) -> None: + caller_api_base: Final = f"{rig.upstream.url}/caller-endpoint-{RUN}/v1" + rig.upstream.drain() + response: Final = rig.gateway.request( + "POST", + "/v1/chat/completions", + {"model": name("caller-api-base"), "messages": messages("caller-api-base"), "api_base": caller_api_base}, + ) + assert response.status_code == 401, response.text + assert [(reached.target, reached.headers["authorization"]) for reached in rig.upstream.drain()] == [ + (f"/caller-endpoint-{RUN}/v1/chat/completions", "Bearer sk-deployment-credential-caller-api-base") + ] + samples: Final = until(rig.gateway, FAILURES, api_base=caller_api_base) + one_off: Final = {sample.labels["model_id"] for sample in series(samples, FAILURES, api_base=caller_api_base)} + assert len(one_off) == 1 and identity("caller-api-base") not in one_off, recorded(samples, api_base=caller_api_base) + assert values(samples, FAILURES, api_base=caller_api_base, exception_status="401") == [1.0] + assert values(samples, TOTALS, api_base=caller_api_base) == [1.0] + assert values(samples, STATE, api_base=caller_api_base) == [1.0] + assert values(samples, LLM_FAILED, model_id=next(iter(one_off))) == [1.0] + assert {metric: values(samples, metric, model_id=identity("caller-api-base")) for metric in DEPLOYMENT_SERIES} == { + metric: [] for metric in DEPLOYMENT_SERIES + } + + +def _gated_key(scenario: Scenario, cell: str) -> str | None: + if cell == "invalid-key": + return f"sk-not-a-virtual-key-{uuid.uuid4().hex}" + if cell == "missing-key": + return None + if cell == "no-model-access": + return scenario.key(models=[name("healthy")]) + key: Final = scenario.key(models=[name(cell)]) + if cell == "expired-key": + write_rows( + "UPDATE \"LiteLLM_VerificationToken\" SET expires = now() - interval '1 hour' WHERE token = %s", + (sha256(key.encode()).hexdigest(),), + database_url=scenario.gateway.database_url, + ) + return key + scenario.gateway.post("/key/block", {"key": key}) + return key + + +def _gated(gateway: Gateway, cell: str, key: str | None) -> httpx.Response: + if key is None: + return gateway.client.post("/v1/chat/completions", json={"model": name(cell), "messages": messages(cell)}) + return send(gateway, cell, key=key) + + +def recorded_before(cell: str) -> list[tuple[str, str, str, float]]: + if cell in ("invalid-key", "expired-key"): + return [] + return sorted([(FAILURES, "", "", 1.0), (TOTALS, "", "", 1.0), (LLM_FAILED, identity(cell), "", 1.0)]) + + +def deployment_rows(samples: Sequence[Sample], cell: str) -> list[tuple[str, str, str, float]]: + return sorted( + (sample.name, sample.labels.get("model_id", ""), sample.labels.get("api_base", ""), sample.value) + for sample in samples + if sample.name in DEPLOYMENT_SERIES + and name(cell) in (sample.labels.get("requested_model"), sample.labels.get("model")) + ) + + +@pytest.mark.parametrize("cell", GATED_CELLS) +def test_a_request_the_proxy_rejects_itself_never_marks_the_deployment_down_and_is_recorded_as_before( + rig: Rig, cell: str +) -> None: + with rig.gateway.scenario() as scenario: + key: Final = _gated_key(scenario, cell) + rig.upstream.drain() + response: Final = _gated(rig.gateway, cell, key) + assert response.status_code == (403 if cell == "no-model-access" else 401), response.text + assert rig.upstream.drain() == () + samples: Final = until(rig.gateway, PROXY_FAILED, requested_model=name(cell)) + assert deployment_rows(samples, cell) == recorded_before(cell), recorded(samples, requested_model=name(cell)) + + +def test_a_key_rate_limit_rejected_at_the_proxy_is_still_counted_as_before(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(models=[name("key-rpm-limit")], rpm_limit=1) + rig.upstream.drain() + admitted: Final = send(rig.gateway, "key-rpm-limit", key=key) + assert admitted.status_code == 200, admitted.text + limited: Final = send(rig.gateway, "key-rpm-limit", key=key) + assert limited.status_code == 429, limited.text + assert len(rig.upstream.drain()) == 1 + samples: Final = until(rig.gateway, FAILURES, requested_model=name("key-rpm-limit"), exception_status="429") + assert values(samples, FAILURES, requested_model=name("key-rpm-limit"), exception_status="429") == [1.0] + eventually( + lambda: scrape(rig.gateway), + lambda observed: ( + values(observed, STATE, model_id=identity("key-rpm-limit"), api_base=rig.api_base) == [0.0] + ), + ) + + +@pytest.fixture(scope="module") +def mcp_peer() -> Iterator[McpPeer]: + with scripted_peer(echo_tool("add")) as peer: + yield peer + + +def test_an_mcp_tool_call_the_proxy_rejects_for_the_caller_stays_out_of_the_llm_failure_metrics( + rig: Rig, mcp_peer: McpPeer +) -> None: + alias: Final = "a365" + RUN + with rig.gateway.scenario() as scenario: + server: Final = register_mcp(scenario, mcp_peer, alias) + key: Final = scenario.key(object_permission={"mcp_servers": [server]}) + hashed: Final = sha256(key.encode()).hexdigest() + outcomes: Final = { + entry: McpCaller(rig.gateway, key, entry, alias).call(f"{alias}-add", {"entry": entry}, server) + for entry in ENTRY_POINTS + } + assert all( + outcome.error is not None and "Agent 365 guardrail rejected" in outcome.raw for outcome in outcomes.values() + ), {entry: outcome.raw for entry, outcome in outcomes.items()} + eventually( + lambda: rig.gateway.rows(GUARDRAIL_ROWS, (AGENT_365_GUARDRAIL, hashed, "call_mcp_tool")), + lambda rows: [row["status"] for row in rows] == ["guardrail_intervened"] * len(ENTRY_POINTS), + seconds=70, + ) + samples: Final = scrape(rig.gateway) + assert values(samples, LLM_FAILED, hashed_api_key=hashed) == [] + assert values(samples, FAILURES, hashed_api_key=hashed) == [] + assert values(samples, TOTALS, hashed_api_key=hashed) == [] diff --git a/tests/integration/observability/test_prometheus_provider_auth_failure_chaos.py b/tests/integration/observability/test_prometheus_provider_auth_failure_chaos.py new file mode 100644 index 00000000000..a1fe1cf4f33 --- /dev/null +++ b/tests/integration/observability/test_prometheus_provider_auth_failure_chaos.py @@ -0,0 +1,147 @@ +import threading +import uuid +from collections.abc import Iterator, Sequence +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import httpx +import psutil +from integration._support.client import Gateway, eventually +from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.prometheus_series import Sample, scrape +from integration._support.wire import Wire, wire_server +from test_prometheus_provider_auth_failure import ( + FAILURES, + LLM_FAILED, + STATE, + TOTALS, + Provider, + config, + identity, + name, + send, + values, +) + +BURST: Final = 12 +ENDPOINTS: Final = ("chat", "chat-stream", "messages", "responses") + + +@dataclass(frozen=True, slots=True) +class ChaosRig: + owned: OwnedProxy + upstream: Wire + release: threading.Event + + +@contextmanager +def _rig(gateway: Gateway, tmp_path: Path) -> Iterator[ChaosRig]: + release: Final = threading.Event() + prom_dir: Final = tmp_path / "prom" + prom_dir.mkdir() + with ( + wire_server(Provider(release)) as upstream, + owned_proxy_process( + gateway, + tmp_path, + {"PROMETHEUS_MULTIPROC_DIR": str(prom_dir)}, + config=config(tmp_path, upstream), + workers=2, + ) as owned, + ): + try: + yield ChaosRig(owned, upstream, release) + finally: + release.set() + + +def _status(gateway: Gateway, cell: str, model: str, key: str | None = None) -> int: + try: + return send(gateway, cell, model=model, key=key).status_code + except httpx.HTTPError: + return -1 + + +def _holds_a_connection_to(pid: int, port: int) -> bool: + return any(connection.raddr[1:] == (port,) for connection in psutil.Process(pid).net_connections(kind="tcp")) + + +def _workers_calling_upstream(rig: ChaosRig) -> frozenset[int]: + port: Final = int(rig.upstream.url.rsplit(":", 1)[1]) + return frozenset(pid for pid in rig.owned.worker_pids() if _holds_a_connection_to(pid, port)) + + +def _counted(samples: Sequence[Sample], cell: str) -> tuple[float, float, float]: + return ( + sum(values(samples, FAILURES, model_id=identity(cell))), + sum(values(samples, TOTALS, model_id=identity(cell))), + sum(values(samples, LLM_FAILED, model_id=identity(cell))), + ) + + +def test_a_burst_of_provider_401s_and_proxy_rejects_across_two_workers_counts_only_the_provider_401s( + gateway: Gateway, tmp_path: Path +) -> None: + with _rig(gateway, tmp_path) as rig: + rig.release.set() + workers: Final = eventually(rig.owned.worker_pids, lambda pids: len(pids) == 2) + invalid: Final = f"sk-not-a-virtual-key-{uuid.uuid4().hex}" + calls: Final[tuple[tuple[str, str, str | None], ...]] = tuple( + (ENDPOINTS[index % len(ENDPOINTS)], name("burst"), None) for index in range(BURST) + ) + tuple(("chat", name("invalid-key"), invalid) for _ in range(BURST)) + + def call(cell: str, model: str, key: str | None) -> int: + return _status(rig.owned.gateway, cell, model, key) + + with ThreadPoolExecutor(max_workers=8) as pool: + statuses: Final = tuple(pool.map(call, *zip(*calls, strict=True))) + assert statuses == (401,) * (2 * BURST), statuses + assert len(rig.upstream.drain()) == BURST + samples: Final = eventually( + lambda: scrape(rig.owned.gateway), + lambda observed: _counted(observed, "burst") == (BURST, BURST, BURST), + seconds=70, + return_last_on_timeout=True, + ) + assert _counted(samples, "burst") == (BURST, BURST, BURST), (workers, _counted(samples, "burst")) + assert set(values(samples, STATE, model_id=identity("burst"))) == {1.0} + assert values(samples, FAILURES, requested_model=name("invalid-key")) == [] + assert values(samples, LLM_FAILED, model=name("invalid-key")) == [] + + +def test_a_worker_killed_while_provider_401s_are_in_flight_loses_no_answered_failure( + gateway: Gateway, tmp_path: Path +) -> None: + with _rig(gateway, tmp_path) as rig: + eventually(rig.owned.worker_pids, lambda pids: len(pids) == 2) + with ThreadPoolExecutor(max_workers=8) as pool: + futures: Final = tuple(pool.submit(_status, rig.owned.gateway, "chat", name("held")) for _ in range(BURST)) + eventually(lambda: rig.upstream.received.qsize() >= 8, bool, seconds=30) + victim: Final = min(eventually(lambda: _workers_calling_upstream(rig), bool)) + with rig.owned.replacing_worker(victim) as kill: + kill() + rig.release.set() + statuses: Final = tuple(future.result() for future in futures) + reached: Final = len(rig.upstream.drain()) + answered: Final = statuses.count(401) + lost: Final = statuses.count(-1) + assert answered + lost == BURST and lost > 0, statuses + samples: Final = eventually( + lambda: scrape(rig.owned.gateway), + lambda observed: answered <= _counted(observed, "held")[0] <= reached, + seconds=70, + return_last_on_timeout=True, + ) + counted: Final = _counted(samples, "held") + assert answered <= counted[0] <= reached and counted[0] == counted[1] == counted[2], (statuses, counted) + follow_up: Final = _status(rig.owned.gateway, "chat", name("held")) + assert follow_up == 401 + recounted: Final = eventually( + lambda: _counted(scrape(rig.owned.gateway), "held"), + lambda observed: observed[0] == counted[0] + 1, + seconds=70, + ) + assert recounted == (counted[0] + 1,) * 3 diff --git a/tests/integration/routing/test_team_wildcard_precedence.py b/tests/integration/routing/test_team_wildcard_precedence.py new file mode 100644 index 00000000000..38fda682129 --- /dev/null +++ b/tests/integration/routing/test_team_wildcard_precedence.py @@ -0,0 +1,61 @@ +"""A team's own wildcard deployment serves that team's keys even when a proxy-wide wildcard of the +same pattern exists, and the proxy-wide one serves every other key (#48490).""" + +import uuid +from typing import Final + +import pytest +from integration._support.client import Gateway, Scenario + +pytestmark: Final = pytest.mark.timeout(120) + + +def _wildcard(scenario: Scenario, pattern: str, *, team_id: str | None) -> str: + identity: Final = f"wildcard-{uuid.uuid4().hex}" + scenario.gateway.post( + "/model/new", + { + "model_name": pattern, + "litellm_params": { + "model": "openai/*", + "api_key": f"provider-key-{identity}", + "api_base": f"{scenario.gateway.upstream_url}/v1", + }, + "model_info": {"id": identity, **({"team_id": team_id} if team_id is not None else {})}, + }, + ) + scenario.cleanups.callback(scenario.delete_model, identity) + return identity + + +def _served_by(gateway: Gateway, key: str, model: str) -> str: + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "wildcard precedence control"}]}, + key=key, + ) + assert response.status_code == 200, response.text + return response.headers["x-litellm-model-id"] + + +def test_a_teams_wildcard_serves_its_keys_before_a_proxy_wide_wildcard_of_the_same_pattern(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + byok_team: Final = scenario.team(models=[]) + other_team: Final = scenario.team(models=[]) + pattern: Final = f"byok-{uuid.uuid4().hex}/*" + team_wildcard: Final = _wildcard(scenario, pattern, team_id=byok_team) + shared_wildcard: Final = _wildcard(scenario, pattern, team_id=None) + model: Final = pattern.replace("*", "gpt-4o-mini") + + served: Final = { + "byok team key": _served_by(gateway, scenario.key(team_id=byok_team), model), + "other team key": _served_by(gateway, scenario.key(team_id=other_team), model), + "key without a team": _served_by(gateway, scenario.key(), model), + } + + assert served == { + "byok team key": team_wildcard, + "other team key": shared_wildcard, + "key without a team": shared_wildcard, + } diff --git a/tests/integration/sdk/test_prometheus_copilot_session_auth_failure_sdk.py b/tests/integration/sdk/test_prometheus_copilot_session_auth_failure_sdk.py new file mode 100644 index 00000000000..1b5e990fe41 --- /dev/null +++ b/tests/integration/sdk/test_prometheus_copilot_session_auth_failure_sdk.py @@ -0,0 +1,99 @@ +"""A provider 401 on a caller's own per-user GitHub Copilot session is the caller's credential failing, so the +shared Copilot deployment keeps no failure, request or outage series for it, while a provider 401 on the router +deployment's own key, on a key a router caller passes or on the key a direct SDK call brings is counted as a failed +request, the router caller's key against its own one-off deployment.""" + +from __future__ import annotations + +import json +from collections.abc import Iterator +from typing import Final + +import pytest +from integration._support.copilot_session_prometheus_worker import ( + CALLER_SESSION_TOKEN, + COPILOT_DEPLOYMENT, + DEPLOYMENT_SERIES, + DIRECT_CALL, + DIRECT_KEY, + FAILURES, + LLM_FAILED, + OPERATOR_DEPLOYMENT, + OPERATOR_KEY, + ROUTER_CALLER_KEY, + STATE, + TOTALS, + Outcome, + WorkerReport, + run_copilot_session_worker, +) +from integration._support.wire import Reply, Request, wire_server + + +@pytest.fixture(scope="module") +def ran() -> Iterator[tuple[WorkerReport, tuple[Request, ...]]]: + def reject(request: Request) -> Reply: + assert request.method == "POST", request.method + return Reply( + status=401, + body=json.dumps({"error": {"message": "bad credentials", "type": "authentication_error"}}).encode(), + ) + + with wire_server(reject) as upstream: + report: Final = run_copilot_session_worker(upstream.url) + yield report, upstream.drain() + + +def test_every_call_reaches_the_provider_with_its_own_credential_and_fails_with_its_401( + ran: tuple[WorkerReport, tuple[Request, ...]], +) -> None: + report, received = ran + assert report.outcomes == ( + Outcome(deployment=COPILOT_DEPLOYMENT, status=401), + Outcome(deployment=OPERATOR_DEPLOYMENT, status=401), + Outcome(deployment=OPERATOR_DEPLOYMENT, status=401), + Outcome(deployment=DIRECT_CALL, status=401), + ) + assert [r.headers.get("authorization") for r in received] == [ + f"Bearer {CALLER_SESSION_TOKEN}", + f"Bearer {OPERATOR_KEY}", + f"Bearer {ROUTER_CALLER_KEY}", + f"Bearer {DIRECT_KEY}", + ] + + +def test_a_provider_401_on_a_callers_copilot_session_stays_out_of_the_shared_deployment_metrics( + ran: tuple[WorkerReport, tuple[Request, ...]], +) -> None: + report, _ = ran + assert {s.metric for s in report.series if s.model_id == COPILOT_DEPLOYMENT} & DEPLOYMENT_SERIES == set() + + +def test_a_provider_401_on_the_deployments_own_key_counts_as_a_deployment_failure( + ran: tuple[WorkerReport, tuple[Request, ...]], +) -> None: + report, _ = ran + assert report.values(FAILURES, OPERATOR_DEPLOYMENT) == (1.0,) + assert report.values(TOTALS, OPERATOR_DEPLOYMENT) == (1.0,) + assert report.values(LLM_FAILED, OPERATOR_DEPLOYMENT) == (1.0,) + assert report.on_the_endpoint(STATE, OPERATOR_DEPLOYMENT) == (1.0,) + + +def test_a_provider_401_on_the_key_a_direct_sdk_call_brings_is_counted_as_a_failed_request( + ran: tuple[WorkerReport, tuple[Request, ...]], +) -> None: + report, _ = ran + assert report.outside_the_router(FAILURES) == (1.0,) + assert report.outside_the_router(TOTALS) == (1.0,) + assert report.outside_the_router(LLM_FAILED) == (1.0,) + + +def test_a_provider_401_on_a_key_a_router_caller_passes_counts_against_its_own_one_off_deployment( + ran: tuple[WorkerReport, tuple[Request, ...]], +) -> None: + report, _ = ran + assert report.on_a_one_off_deployment(FAILURES) == (1.0,) + assert report.on_a_one_off_deployment(TOTALS) == (1.0,) + assert report.on_a_one_off_deployment(LLM_FAILED) == (1.0,) + assert report.values(FAILURES, OPERATOR_DEPLOYMENT) == (1.0,) + assert report.values(TOTALS, OPERATOR_DEPLOYMENT) == (1.0,) diff --git a/tests/unit/decisions/test_main.py b/tests/unit/decisions/test_main.py index e2385a2e215..5b32008b279 100644 --- a/tests/unit/decisions/test_main.py +++ b/tests/unit/decisions/test_main.py @@ -400,6 +400,68 @@ async def test_systemone_uses_evaluation_pricing(respx_mock: respx.MockRouter, d ) +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("base_url", "expected_api_base"), + [ + ("https://onprem.example/v1", "https://onprem.example/v1"), + ("", "https://fallback.example/v1"), + (None, "https://fallback.example/v1"), + (123, "https://fallback.example/v1"), + ], +) +async def test_openai_decisions_honors_base_url_alias( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + base_url: str | int | None, + expected_api_base: str, +) -> None: + monkeypatch.setenv("OPENAI_BASE_URL", "https://fallback.example/v1") + route: Final = respx_mock.post(f"{expected_api_base}/decisions").respond(json=_OPENAI_RESPONSE) + response: Final = await litellm.adecisions( + model="openai/gpt-6-luna", + input="The package arrived broken.", + questions=[{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}], + api_key="backend-secret-key", + base_url=base_url, + ) + + assert route.called + assert route.calls[0].request.headers["authorization"] == "Bearer backend-secret-key" + assert response._hidden_params["model"] == "openai/gpt-6-luna" + + +@pytest.mark.asyncio +async def test_openai_systemone_honors_base_url_alias(respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.post("https://onprem.example/v1/decisions").respond(json=_OPENAI_RESPONSE) + response: Final = await litellm.asystemone( + model="openai/gpt-6-luna", + state="The package arrived broken.", + questions=_QUESTIONS, + api_key="backend-secret-key", + base_url="https://onprem.example/v1", + ) + + assert route.called + assert route.calls[0].request.headers["authorization"] == "Bearer backend-secret-key" + assert response._hidden_params["model"] == "openai/gpt-6-luna" + + +@pytest.mark.asyncio +async def test_openai_decisions_api_base_wins_over_base_url_alias(respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.post("https://api.openai.com/v1/decisions").respond(json=_OPENAI_RESPONSE) + await litellm.adecisions( + model="openai/gpt-6-luna", + input="The package arrived broken.", + questions=[{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}], + api_key="test-key", + api_base="https://api.openai.com/v1", + base_url="https://onprem.example/v1", + ) + + assert route.called + + @pytest.mark.asyncio async def test_systemone_falls_back_to_plain_pricing_without_decisions_row( respx_mock: respx.MockRouter, diff --git a/tests/unit/integrations/test_prometheus_invalid_key_filtering.py b/tests/unit/integrations/test_prometheus_invalid_key_filtering.py index 9e8348e860a..453fd28636d 100644 --- a/tests/unit/integrations/test_prometheus_invalid_key_filtering.py +++ b/tests/unit/integrations/test_prometheus_invalid_key_filtering.py @@ -2,19 +2,63 @@ Unit tests for Prometheus invalid API key request filtering. Tests the 401 detection helpers, that LLM-level metrics skip invalid API key -requests, and that the proxy-level failed request counter still records them. +requests rejected at the proxy gate while recording a provider's 401, and that +the proxy-level failed request counter still records them. """ +from collections.abc import Iterator +from datetime import datetime +from typing import Final from unittest.mock import Mock, patch import pytest from fastapi import HTTPException from prometheus_client import REGISTRY +from prometheus_client.samples import Sample +import litellm +from litellm.constants import LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL from litellm.integrations.prometheus import PrometheusLogger +from litellm.llms.github_copilot.per_user_auth import GithubCopilotUserSession from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth +PROVIDER_CREDENTIAL_FIELDS = [ + "aws_access_key_id", + "aws_secret_access_key", + "aws_session_token", + "aws_web_identity_token", + "azure_ad_token", + "azure_password", + "azure_scope", + "azure_username", + "client_id", + "client_secret", + "gigachat_access_token", + "oci_fingerprint", + "oci_key", + "oci_key_file", + "oci_tenancy", + "oci_user", + "tenant_id", + "token", + "user_config", + "vertex_ai_credentials", + "vertex_credentials", + "watsonx_credentials", + "wx_credentials", + "zen_api_key", +] +PROVIDER_ENDPOINT_FIELDS = [ + "aws_bedrock_runtime_endpoint", + "aws_sts_endpoint", + "deployment_url", + "s3_endpoint_url", + "sagemaker_base_url", +] +CALLED_THE_PROVIDER: Final = {"api_call_start_time": datetime(2026, 10, 19, 12, 0, 0)} + + @pytest.fixture(scope="function") def prometheus_logger(): """Create a PrometheusLogger instance for testing.""" @@ -48,18 +92,13 @@ class TestExtractStatusCode: (ExceptionWithStatusCode, 401, 401), ], ) - def test_extract_from_exception( - self, prometheus_logger, exception_class, code_value, expected - ): + def test_extract_from_exception(self, prometheus_logger, exception_class, code_value, expected): exception = exception_class(code_value) assert prometheus_logger._extract_status_code(exception=exception) == expected def test_extract_from_kwargs(self, prometheus_logger): exception = ExceptionWithCode("401") - assert ( - prometheus_logger._extract_status_code(kwargs={"exception": exception}) - == 401 - ) + assert prometheus_logger._extract_status_code(kwargs={"exception": exception}) == 401 def test_extract_from_enum_values(self, prometheus_logger): enum_values = Mock(status_code="401") @@ -79,50 +118,271 @@ class TestInvalidAPIKeyDetection: ], ) def test_status_code_detection(self, prometheus_logger, status_code, expected): - assert ( - prometheus_logger._is_invalid_api_key_request(status_code=status_code) - == expected - ) + assert prometheus_logger._is_invalid_api_key_request(status_code=status_code) == expected def test_auth_error_message_detection(self, prometheus_logger): exception = AssertionError( "LiteLLM Virtual Key expected. Received=invalid-key-12345, expected to start with 'sk-'." ) - assert ( - prometheus_logger._is_invalid_api_key_request( - status_code=None, exception=exception - ) - is True - ) + assert prometheus_logger._is_invalid_api_key_request(status_code=None, exception=exception) is True def test_non_auth_exception_not_detected(self, prometheus_logger): exception = ValueError("Some other error") - assert ( - prometheus_logger._is_invalid_api_key_request( - status_code=None, exception=exception - ) - is False - ) + assert prometheus_logger._is_invalid_api_key_request(status_code=None, exception=exception) is False class TestSkipMetricsValidation: """Test high-level validation method that orchestrates detection and extraction.""" - def test_skip_for_401_exception(self, prometheus_logger): - """Test full flow: extraction -> detection -> skip decision.""" - exception = ExceptionWithCode("401") - assert ( - prometheus_logger._should_skip_metrics_for_invalid_key(exception=exception) - is True - ) + @pytest.mark.parametrize( + "exception,proxy_gate,expected", + [ + (ExceptionWithCode("401"), True, True), + (AssertionError("expected to start with 'sk-'"), True, True), + (ExceptionWithStatusCode(429), True, False), + (ExceptionWithCode("401"), False, True), + (AssertionError("expected to start with 'sk-'"), False, True), + ( + litellm.AuthenticationError( + message="Incorrect API key provided: sk-rotat*****", + llm_provider="openai", + model="gpt-4o-mini", + ), + False, + False, + ), + ( + litellm.AuthenticationError( + message="API key not valid. Please pass a valid API key.", + llm_provider="gemini", + model="gemini-2.5-flash", + ), + False, + False, + ), + ( + litellm.AuthenticationError( + message="Authentication Error, Invalid proxy server token passed.", + llm_provider="openai", + model="gpt-4o-mini", + ), + True, + True, + ), + ( + litellm.CallerCredentialAuthenticationError( + message="GitHub Copilot connection is missing for this user", + llm_provider="github_copilot", + model="gpt-4o", + ), + False, + True, + ), + ( + HTTPException(status_code=401, detail="caller did not present a token"), + False, + True, + ), + ( + litellm.BadRequestError( + message="API key not valid. Please pass a valid API key.", + llm_provider="gemini", + model="gemini-2.5-flash", + ), + False, + False, + ), + ], + ids=[ + "proxy-gate-401", + "proxy-gate-auth-message", + "proxy-gate-429", + "unmarked-proxy-401", + "unmarked-auth-message", + "provider-401", + "provider-auth-message", + "proxy-gate-litellm-auth-error", + "caller-credential-401", + "unmarked-http-exception-401", + "provider-400-auth-message", + ], + ) + def test_skip_except_for_provider_errors_the_proxy_did_not_raise( + self, prometheus_logger, exception, proxy_gate, expected + ): + kwargs: Final[dict[str, object]] = {"exception": exception, **CALLED_THE_PROVIDER} + if proxy_gate: + kwargs[LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL] = True + assert prometheus_logger._should_skip_metrics_for_invalid_key(kwargs=kwargs) is expected - def test_skip_for_auth_error_message(self, prometheus_logger): - """Test full flow: exception message -> detection -> skip decision.""" - exception = AssertionError("expected to start with 'sk-'") - assert ( - prometheus_logger._should_skip_metrics_for_invalid_key(exception=exception) - is True + @pytest.mark.parametrize( + "logged_call,expected", + [ + ({"litellm_params": {"proxy_server_request": {"caller_credential_fields": ("api_key",)}}}, True), + ( + { + "litellm_params": { + "github_copilot_user_session": GithubCopilotUserSession( + token="caller-session-token", api_base="https://api.githubcopilot.com" + ) + } + }, + True, + ), + ({"standard_logging_object": {"metadata": {"used_client_oauth_token": True}}}, True), + ({"litellm_params": {"proxy_server_request": {"caller_credential_fields": ("headers",)}}}, True), + ( + { + "litellm_params": { + "proxy_server_request": {"caller_credential_fields": ("provider_specific_header",)} + } + }, + True, + ), + ({"litellm_params": {"proxy_server_request": {"caller_credential_fields": ()}}}, False), + ({"litellm_params": {"proxy_server_request": {"caller_credential_fields": "api_key"}}}, False), + ({"standard_logging_object": {"metadata": {"used_client_oauth_token": False}}}, False), + ({"litellm_params": {"proxy_server_request": {"body": {"azure_ad_token": "caller-token"}}}}, True), + ( + { + "litellm_params": { + "proxy_server_request": {"body": {"extra_body": {"azure_ad_token": "caller-token"}}} + } + }, + True, + ), + ( + { + "litellm_params": { + "azure_ad_token": "deployment-token", + "proxy_server_request": {"body": {"model": "azure-deployment"}}, + } + }, + False, + ), + ({"litellm_params": {"proxy_server_request": {"body": {"extra_body": "azure_ad_token"}}}}, False), + ( + { + "litellm_params": { + "proxy_server_request": {"body": {"litellm_credential_name": "caller-credential"}} + } + }, + True, + ), + ( + { + "litellm_params": { + "litellm_credential_name": "deployment-credential", + "proxy_server_request": {"body": {"model": "openai-deployment"}}, + } + }, + False, + ), + ({"litellm_params": {"proxy_server_request": {"caller_credential_fields": ("extra_headers",)}}}, True), + ({"litellm_params": {"proxy_server_request": {"caller_credential_fields": ("organization",)}}}, True), + ( + { + "litellm_params": { + "proxy_server_request": {"credential_fields": ("headers",), "caller_credential_fields": ()} + } + }, + False, + ), + ( + { + "litellm_params": { + "proxy_server_request": {"body": {"extra_headers": {"Authorization": "Bearer sk-caller"}}} + } + }, + False, + ), + ], + ids=[ + "caller-supplied-api-key", + "caller-copilot-session", + "caller-anthropic-oauth-token", + "caller-supplied-headers", + "caller-supplied-provider-specific-header", + "no-caller-credential-fields", + "malformed-credential-fields", + "deployment-anthropic-key", + "caller-supplied-azure-ad-token", + "caller-supplied-azure-ad-token-in-extra-body", + "deployment-azure-ad-token", + "malformed-extra-body", + "caller-named-credential", + "deployment-named-credential", + "caller-supplied-extra-headers", + "caller-supplied-organization", + "headers-that-cannot-change-provider-auth", + "extra-headers-not-stamped-by-the-proxy", + ], + ) + def test_skip_a_provider_401_on_a_credential_the_caller_owns(self, prometheus_logger, logged_call, expected): + exception: Final = litellm.AuthenticationError( + message="Incorrect API key provided: sk-calle*****", + llm_provider="openai", + model="gpt-4o-mini", ) + kwargs: Final = {"exception": exception, **CALLED_THE_PROVIDER, **logged_call} + assert prometheus_logger._should_skip_metrics_for_invalid_key(kwargs=kwargs) is expected + + @pytest.mark.parametrize( + "logged_call,expected", + [(CALLED_THE_PROVIDER, False), ({}, True), ({"api_call_start_time": None}, True)], + ids=["provider-401", "401-before-the-provider-call", "blank-provider-call-start"], + ) + def test_skip_a_401_raised_before_the_provider_was_called(self, prometheus_logger, logged_call, expected): + exception: Final = litellm.AuthenticationError( + message="Missing Anthropic API Key - A call is being made to anthropic but no key is set", + llm_provider="anthropic", + model="claude-haiku-4-5", + ) + kwargs: Final = {"exception": exception, **logged_call} + assert prometheus_logger._should_skip_metrics_for_invalid_key(kwargs=kwargs) is expected + + @pytest.mark.parametrize("field", PROVIDER_CREDENTIAL_FIELDS + PROVIDER_ENDPOINT_FIELDS) + @pytest.mark.parametrize("inside_extra_body", [False, True], ids=["body", "extra-body"]) + def test_skip_a_provider_401_on_a_provider_credential_field_the_caller_sent( + self, prometheus_logger, field, inside_extra_body + ): + sent: Final = {field: "caller-credential"} + body: Final = {"model": "shared-deployment", **({"extra_body": sent} if inside_extra_body else sent)} + kwargs: Final = { + "exception": litellm.AuthenticationError(message="Unauthorized", llm_provider="azure", model="gpt-4o"), + "litellm_params": {"proxy_server_request": {"body": body}}, + **CALLED_THE_PROVIDER, + } + assert prometheus_logger._should_skip_metrics_for_invalid_key(kwargs=kwargs) is True + + @pytest.mark.parametrize("field", PROVIDER_CREDENTIAL_FIELDS + PROVIDER_ENDPOINT_FIELDS) + def test_count_a_provider_401_on_a_provider_credential_field_the_deployment_owns(self, prometheus_logger, field): + kwargs: Final = { + "exception": litellm.AuthenticationError(message="Unauthorized", llm_provider="azure", model="gpt-4o"), + **CALLED_THE_PROVIDER, + "litellm_params": { + field: "deployment-credential", + "proxy_server_request": {"body": {"model": "shared-deployment"}}, + }, + } + assert prometheus_logger._should_skip_metrics_for_invalid_key(kwargs=kwargs) is False + + @pytest.mark.parametrize( + "token_exchange_endpoint,expected", + [("https://login.example.com/oauth2/v2.0/token", True), ("", False), (None, False)], + ids=["token-exchange-deployment", "blank-endpoint", "no-endpoint"], + ) + def test_skip_a_provider_401_on_a_token_exchange_deployment( + self, prometheus_logger, token_exchange_endpoint, expected + ): + kwargs: Final = { + "exception": litellm.AuthenticationError( + message="Unauthorized", llm_provider="microsoft_365_copilot", model="copilot" + ), + "litellm_params": {"token_exchange_endpoint": token_exchange_endpoint, "client_secret": "deployment"}, + **CALLED_THE_PROVIDER, + } + assert prometheus_logger._should_skip_metrics_for_invalid_key(kwargs=kwargs) is expected def test_no_skip_for_valid_request(self, prometheus_logger): assert prometheus_logger._should_skip_metrics_for_invalid_key() is False @@ -147,19 +407,13 @@ class TestAsyncHooks: ), ], ) - async def test_post_call_failure_hook_counts_401_without_key_hash( - self, prometheus_logger, exception - ): + async def test_post_call_failure_hook_counts_401_without_key_hash(self, prometheus_logger, exception): unauthenticated = UserAPIKeyAuth(request_route="/v1/chat/completions") unauthenticated.api_key = "notakeyatall" with ( - patch.object( - prometheus_logger, "litellm_proxy_failed_requests_metric" - ) as mock_failed, - patch.object( - prometheus_logger, "litellm_proxy_total_requests_metric" - ) as mock_total, + patch.object(prometheus_logger, "litellm_proxy_failed_requests_metric") as mock_failed, + patch.object(prometheus_logger, "litellm_proxy_total_requests_metric") as mock_total, ): await prometheus_logger.async_post_call_failure_hook( request_data={"model": "test-model"}, @@ -176,9 +430,7 @@ class TestAsyncHooks: mock_total.labels.return_value.inc.assert_called_once() @pytest.mark.asyncio - async def test_post_call_failure_hook_keeps_resolved_identity_labels_for_401( - self, prometheus_logger - ): + async def test_post_call_failure_hook_keeps_resolved_identity_labels_for_401(self, prometheus_logger): expired_key = UserAPIKeyAuth( api_key="sk-expired", key_alias="expired-alias", @@ -191,9 +443,7 @@ class TestAsyncHooks: code=401, ) - with patch.object( - prometheus_logger, "litellm_proxy_failed_requests_metric" - ) as mock_failed: + with patch.object(prometheus_logger, "litellm_proxy_failed_requests_metric") as mock_failed: await prometheus_logger.async_post_call_failure_hook( request_data={"model": "test-model"}, original_exception=exception, @@ -207,7 +457,7 @@ class TestAsyncHooks: assert failed_labels["team"] == "team-1" @pytest.mark.asyncio - async def test_log_failure_event_skips_401(self, prometheus_logger): + async def test_log_failure_event_skips_401_rejected_at_the_proxy_gate(self, prometheus_logger): exception = ExceptionWithCode("401") kwargs = { "model": "test-model", @@ -220,20 +470,91 @@ class TestAsyncHooks: }, "exception": exception, "litellm_params": {}, + LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL: True, } with ( - patch.object( - prometheus_logger, "litellm_llm_api_failed_requests_metric" - ) as mock_failed, - patch.object( - prometheus_logger, "set_llm_deployment_failure_metrics" - ) as mock_deployment, + patch.object(prometheus_logger, "litellm_llm_api_failed_requests_metric") as mock_failed, + patch.object(prometheus_logger, "set_llm_deployment_failure_metrics") as mock_deployment, ): - await prometheus_logger.async_log_failure_event( kwargs=kwargs, response_obj=None, start_time=None, end_time=None ) mock_failed.labels.assert_not_called() mock_deployment.assert_not_called() + + +def _provider_auth_failure_kwargs() -> dict[str, object]: + exception: Final = litellm.AuthenticationError( + message="Incorrect API key provided: sk-rotat*****", + llm_provider="openai", + model="gpt-4o-mini", + ) + return { + "model": "gpt-4o-mini", + "standard_logging_object": { + "metadata": { + "user_api_key_hash": "test-key", + "user_api_key_alias": "test-alias", + "user_api_key_user_id": "test-user", + "user_api_key_team_id": "test-team", + "user_api_key_team_alias": "test-team-alias", + }, + "model_group": "rotated-credential", + "model_id": "deployment-1", + "api_base": "https://api.openai.com/v1", + "custom_llm_provider": "openai", + }, + "exception": exception, + "litellm_params": {"custom_llm_provider": "openai", "metadata": {}}, + **CALLED_THE_PROVIDER, + } + + +def _metric_values(name: str, model_id: str) -> list[float]: + return [ + sample.value + for sample in _registry_samples() + if sample.name == name and sample.labels.get("model_id") == model_id + ] + + +def _registry_samples() -> Iterator[Sample]: + for metric in REGISTRY.collect(): + yield from metric.samples + + +class TestProviderAuthFailures: + """A 401 raised by the provider is an upstream failure and is recorded.""" + + @pytest.mark.asyncio + async def test_log_failure_event_records_a_provider_401(self, prometheus_logger): + await prometheus_logger.async_log_failure_event( + kwargs=_provider_auth_failure_kwargs(), + response_obj=None, + start_time=None, + end_time=None, + ) + + assert _metric_values("litellm_llm_api_failed_requests_metric_total", "deployment-1") == [1.0] + assert _metric_values("litellm_deployment_failure_responses_total", "deployment-1") == [1.0] + assert _metric_values("litellm_deployment_state", "deployment-1") == [1.0] + + def test_deployment_failure_metrics_record_a_provider_401(self, prometheus_logger): + prometheus_logger.set_llm_deployment_failure_metrics(_provider_auth_failure_kwargs()) + + failures: Final = [ + sample for sample in _registry_samples() if sample.name == "litellm_deployment_failure_responses_total" + ] + assert [ + ( + sample.labels["model_id"], + sample.labels["exception_status"], + sample.labels["api_provider"], + sample.value, + ) + for sample in failures + ] == [("deployment-1", "401", "openai", 1.0)] + assert _metric_values("litellm_deployment_total_requests_total", "deployment-1") == [1.0] + assert _metric_values("litellm_deployment_state", "deployment-1") == [1.0] diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py index 32f18b952ac..b3b04440d16 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -31,7 +31,7 @@ from litellm.llms.anthropic.pass_through.messages.transformation import ( ) from litellm.llms.azure_ai.anthropic.transformation import AzureAnthropicConfig from litellm.llms.bedrock.claude_platform.transformation import BedrockClaudePlatformConfig -from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import ( +from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude_transformation import ( AmazonAnthropicClaudeConfig, ) from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import ( @@ -2549,18 +2549,6 @@ def test_translate_system_message_preserves_cache_control(): # ============ Dynamic max_tokens Tests ============ -def test_get_max_tokens_for_model_claude_3(): - """ - Test that get_max_tokens_for_model returns correct value for Claude 3 models. - Claude 3 models have max_output_tokens of 4096. - """ - config = AnthropicConfig() - - # Claude 3 Sonnet should return 4096 - max_tokens = config.get_max_tokens_for_model("claude-3-sonnet-20240229") - assert max_tokens == 4096 - - def test_get_max_tokens_for_model_claude_35(): """ Test that get_max_tokens_for_model returns correct value for Claude 3.5 models. @@ -2581,24 +2569,24 @@ def test_get_max_tokens_for_model_claude_35(): def test_get_max_tokens_for_model_unknown(): """ - Test that get_max_tokens_for_model returns 4096 fallback for unknown models. + Test that get_max_tokens_for_model returns 128000 fallback for unknown models. """ config = AnthropicConfig() - # Unknown model should return 4096 as fallback + # Unknown model should return 128000 as fallback max_tokens = config.get_max_tokens_for_model("unknown-model-xyz") - assert max_tokens == 4096 + assert max_tokens == 128000 def test_get_max_tokens_for_model_none(): """ - Test that get_max_tokens_for_model returns 4096 fallback when model is None. + Test that get_max_tokens_for_model returns 128000 fallback when model is None. """ config = AnthropicConfig() - # None model should return 4096 as fallback + # None model should return 128000 as fallback max_tokens = config.get_max_tokens_for_model(None) - assert max_tokens == 4096 + assert max_tokens == 128000 def test_get_config_with_model_uses_dynamic_max_tokens(): @@ -2639,10 +2627,10 @@ def test_get_config_with_model_uses_dynamic_max_tokens(): def test_get_config_without_model_uses_fallback(): """ - Test that get_config without model parameter uses 4096 fallback. + Test that get_config without model parameter uses 128000 fallback. """ config = AnthropicConfig.get_config() - assert config["max_tokens"] == 4096 + assert config["max_tokens"] == 128000 def test_get_config_does_not_leak_module_constants(): @@ -6145,7 +6133,7 @@ def test_anthropic_messages_request_strips_billing_header_for_minimax(): def test_translate_system_message_strips_billing_header_for_bedrock_invoke(): - from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import ( + from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude_transformation import ( AmazonAnthropicClaudeConfig, ) @@ -6174,12 +6162,12 @@ def test_translate_system_message_strips_billing_header_for_bedrock_invoke(): True, ), ( - "litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation", + "litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude_transformation", "AmazonAnthropicClaudeConfig", True, ), ( - "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation", + "litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude_transformation", "AmazonAnthropicClaudeMessagesConfig", True, ), diff --git a/tests/unit/llms/anthropic/pass_through/messages/test_reasoning_effort_translation.py b/tests/unit/llms/anthropic/pass_through/messages/test_reasoning_effort_translation.py index c1295305c7a..284c6ea532d 100644 --- a/tests/unit/llms/anthropic/pass_through/messages/test_reasoning_effort_translation.py +++ b/tests/unit/llms/anthropic/pass_through/messages/test_reasoning_effort_translation.py @@ -11,7 +11,7 @@ from litellm.llms.anthropic.common_utils import AnthropicError from litellm.llms.anthropic.pass_through.messages.transformation import ( AnthropicMessagesConfig, ) -from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( +from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude_transformation import ( AmazonAnthropicClaudeMessagesConfig, ) diff --git a/tests/unit/llms/azure/chat/test_azure_base_model_routing.py b/tests/unit/llms/azure/chat/test_azure_base_model_routing.py index aae2587085f..d9607c6db8c 100644 --- a/tests/unit/llms/azure/chat/test_azure_base_model_routing.py +++ b/tests/unit/llms/azure/chat/test_azure_base_model_routing.py @@ -12,7 +12,7 @@ import openai import pytest import litellm -from litellm.llms.azure.chat.gpt_5_transformation import AzureOpenAIGPT5Config +from litellm.llms.azure.chat.gpt_reasoning_transformation import AzureOpenAIGPT5Config from litellm.llms.azure.chat.o_series_transformation import AzureOpenAIO1Config from litellm.utils import ProviderConfigManager, get_optional_params from tests.unit.llms.bedrock.slow_upstream import STREAM_TIMEOUT_SECONDS, slow_upstream_transport diff --git a/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py b/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py index f753fb6b8f9..53df6610913 100644 --- a/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py +++ b/tests/unit/llms/azure/chat/test_azure_chat_gpt_transformation.py @@ -11,7 +11,7 @@ sys.path.insert( import litellm from litellm.litellm_core_utils.prompt_templates.common_utils import TOOL_RESULT_IMAGE_BOUNDARY -from litellm.llms.azure.chat.gpt_5_transformation import AzureOpenAIGPT5Config +from litellm.llms.azure.chat.gpt_reasoning_transformation import AzureOpenAIGPT5Config from litellm.llms.azure.chat.gpt_transformation import AzureOpenAIConfig from litellm.utils import _invalidate_model_cost_lowercase_map, get_optional_params from datetime import datetime diff --git a/tests/unit/llms/azure/chat/test_azure_gpt5_transformation.py b/tests/unit/llms/azure/chat/test_azure_gpt_reasoning_transformation.py similarity index 99% rename from tests/unit/llms/azure/chat/test_azure_gpt5_transformation.py rename to tests/unit/llms/azure/chat/test_azure_gpt_reasoning_transformation.py index 80328022f61..a82dd98e31d 100644 --- a/tests/unit/llms/azure/chat/test_azure_gpt5_transformation.py +++ b/tests/unit/llms/azure/chat/test_azure_gpt_reasoning_transformation.py @@ -5,7 +5,7 @@ import pytest import litellm from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map -from litellm.llms.azure.chat.gpt_5_transformation import AzureOpenAIGPT5Config +from litellm.llms.azure.chat.gpt_reasoning_transformation import AzureOpenAIGPT5Config @pytest.fixture() diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude_transformation.py similarity index 99% rename from tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py rename to tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude_transformation.py index bd98e4eb4de..86c4b29eb53 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude_transformation.py @@ -14,7 +14,7 @@ from pydantic import BaseModel, ConfigDict # Ensure the project root is on the import path so `litellm` can be imported when # tests are executed from any working directory. import litellm -from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import ( +from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude_transformation import ( AmazonAnthropicClaudeConfig, ) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @@ -188,7 +188,7 @@ def test_output_format_conversion_to_inline_schema(): Bedrock Invoke doesn't support the output_format parameter, so LiteLLM converts it by embedding the schema directly into the user message content. """ - from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( + from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude_transformation import ( AmazonAnthropicClaudeMessagesConfig, ) @@ -265,7 +265,7 @@ def test_output_format_conversion_with_string_content(): """ Test that output_format conversion works when message content is a string (not a list). """ - from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( + from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude_transformation import ( AmazonAnthropicClaudeMessagesConfig, ) @@ -313,7 +313,7 @@ def test_output_format_with_no_schema(): """ Test that if output_format has no schema, the conversion is skipped gracefully. """ - from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( + from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude_transformation import ( AmazonAnthropicClaudeMessagesConfig, ) @@ -352,7 +352,7 @@ def test_opus_4_5_model_detection(): Test that the _is_claude_opus_4_5 method correctly identifies Opus 4.5 models with various naming conventions. """ - from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( + from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude_transformation import ( AmazonAnthropicClaudeMessagesConfig, ) @@ -1526,7 +1526,7 @@ class TestBedrockAnthropicPromptCachingRegression: from litellm.llms.bedrock.chat.converse_transformation import ( AmazonConverseConfig, ) - from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import ( + from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude_transformation import ( AmazonAnthropicClaudeConfig, ) @@ -1610,7 +1610,7 @@ class TestBedrockAnthropicPromptCachingRegression: from litellm.llms.bedrock.chat.converse_transformation import ( AmazonConverseConfig, ) - from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import ( + from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude_transformation import ( AmazonAnthropicClaudeConfig, ) @@ -1688,7 +1688,7 @@ class TestBedrockAnthropic1MContextRegression: from litellm.llms.bedrock.chat.converse_transformation import ( AmazonConverseConfig, ) - from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import ( + from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude_transformation import ( AmazonAnthropicClaudeConfig, ) @@ -1753,7 +1753,7 @@ class TestBedrockAnthropic1MContextRegression: from litellm.llms.bedrock.chat.converse_transformation import ( AmazonConverseConfig, ) - from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import ( + from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude_transformation import ( AmazonAnthropicClaudeConfig, ) @@ -1807,7 +1807,7 @@ class TestBedrockAnthropic1MContextRegression: from litellm.llms.bedrock.chat.converse_transformation import ( AmazonConverseConfig, ) - from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import ( + from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude_transformation import ( AmazonAnthropicClaudeConfig, ) @@ -1867,7 +1867,7 @@ class TestBedrockAnthropicCombinedRegressions: from litellm.llms.bedrock.chat.converse_transformation import ( AmazonConverseConfig, ) - from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import ( + from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude_transformation import ( AmazonAnthropicClaudeConfig, ) diff --git a/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude_transformation.py similarity index 99% rename from tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py rename to tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude_transformation.py index 6fe32785281..4d9feea4454 100644 --- a/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude_transformation.py @@ -37,7 +37,7 @@ from litellm.llms.bedrock.common_utils import ( normalize_custom_field_on_tools, normalize_tool_input_schema_types_for_bedrock_invoke, ) -from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( +from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude_transformation import ( AmazonAnthropicClaudeMessagesConfig, AmazonAnthropicClaudeMessagesStreamDecoder, ) diff --git a/tests/unit/llms/bedrock/test_anthropic_beta_support.py b/tests/unit/llms/bedrock/test_anthropic_beta_support.py index a20ec94a99d..73d529ae3ba 100644 --- a/tests/unit/llms/bedrock/test_anthropic_beta_support.py +++ b/tests/unit/llms/bedrock/test_anthropic_beta_support.py @@ -11,11 +11,11 @@ from unittest.mock import MagicMock, patch import pytest from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig -from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import ( +from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude_transformation import ( AmazonAnthropicClaudeConfig, ) from litellm.llms.bedrock.common_utils import get_anthropic_beta_from_headers -from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( +from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude_transformation import ( AmazonAnthropicClaudeMessagesConfig, ) diff --git a/tests/unit/llms/bedrock/test_base_aws_llm.py b/tests/unit/llms/bedrock/test_base_aws_llm.py index 4350fdb699b..c220e39f4d1 100644 --- a/tests/unit/llms/bedrock/test_base_aws_llm.py +++ b/tests/unit/llms/bedrock/test_base_aws_llm.py @@ -3228,7 +3228,7 @@ class TestGetBedrockModelIdArnHandling: def test_arn_url_matches_expected(self): """Full URL built from messages config must match expected encoded form.""" import urllib.parse - from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( + from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude_transformation import ( AmazonAnthropicClaudeMessagesConfig, ) diff --git a/tests/unit/llms/bedrock/test_request_metadata.py b/tests/unit/llms/bedrock/test_request_metadata.py index 6b74ad90cff..8f94ad54c3d 100644 --- a/tests/unit/llms/bedrock/test_request_metadata.py +++ b/tests/unit/llms/bedrock/test_request_metadata.py @@ -15,7 +15,7 @@ from litellm.llms.bedrock.chat.invoke_transformations.amazon_openai_transformati from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import ( AmazonInvokeConfig, ) -from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( +from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude_transformation import ( AmazonAnthropicClaudeMessagesConfig, ) from litellm.llms.bedrock.request_metadata import ( diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index f7850aaf9a1..55fad96ceb1 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -49,7 +49,7 @@ from litellm.llms.custom_httpx.llm_http_handler import ( _rust_responses_websocket_enabled, ) from litellm.llms.azure.videos.transformation import AzureVideoConfig -from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( +from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude_transformation import ( AmazonAnthropicClaudeMessagesConfig, ) from litellm.llms.bedrock.files.transformation import BedrockFilesConfig @@ -3173,7 +3173,7 @@ async def test_anthropic_invalid_thinking_signature_retry_resigns_bedrock_reques """Regression: after Bedrock rejects a replayed thinking block (400 invalid signature), the strip-and-retry re-sign must not inherit attempt 1's SigV4 Authorization/X-Amz-Date; reusing them over the new stripped body makes AWS return 403 SignatureDoesNotMatch.""" - from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( + from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude_transformation import ( AmazonAnthropicClaudeMessagesConfig, ) diff --git a/tests/unit/llms/openai/chat/test_openai_gpt_transformation.py b/tests/unit/llms/openai/chat/test_openai_gpt_transformation.py index 76faf338307..3e96995621e 100644 --- a/tests/unit/llms/openai/chat/test_openai_gpt_transformation.py +++ b/tests/unit/llms/openai/chat/test_openai_gpt_transformation.py @@ -12,7 +12,7 @@ from httpx import Response import litellm from litellm.litellm_core_utils.prompt_templates.common_utils import TOOL_RESULT_IMAGE_BOUNDARY from litellm.litellm_core_utils.prompt_templates.image_handling import in_memory_cache -from litellm.llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config +from litellm.llms.openai.chat.gpt_reasoning_transformation import OpenAIGPT5Config from litellm.llms.openai.chat.gpt_transformation import ( OpenAIChatCompletionStreamingHandler, OpenAIGPTConfig, diff --git a/tests/unit/llms/openai/test_gpt5_transformation.py b/tests/unit/llms/openai/test_gpt_reasoning_transformation.py similarity index 99% rename from tests/unit/llms/openai/test_gpt5_transformation.py rename to tests/unit/llms/openai/test_gpt_reasoning_transformation.py index 7b10411ba20..4cdc674c774 100644 --- a/tests/unit/llms/openai/test_gpt5_transformation.py +++ b/tests/unit/llms/openai/test_gpt_reasoning_transformation.py @@ -7,7 +7,7 @@ import pytest import litellm import litellm.main as litellm_main from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map -from litellm.llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config +from litellm.llms.openai.chat.gpt_reasoning_transformation import OpenAIGPT5Config from litellm.llms.openai.openai import OpenAIConfig from litellm.utils import ( is_explicitly_disabled_factory, diff --git a/tests/unit/llms/openai/test_is_model_gpt_5_model.py b/tests/unit/llms/openai/test_is_model_gpt_5_model.py index f6fef92fc6a..39d7c9da2ce 100644 --- a/tests/unit/llms/openai/test_is_model_gpt_5_model.py +++ b/tests/unit/llms/openai/test_is_model_gpt_5_model.py @@ -35,8 +35,8 @@ from typing import Final import pytest import litellm -from litellm.llms.azure.chat.gpt_5_transformation import AzureOpenAIGPT5Config -from litellm.llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config +from litellm.llms.azure.chat.gpt_reasoning_transformation import AzureOpenAIGPT5Config +from litellm.llms.openai.chat.gpt_reasoning_transformation import OpenAIGPT5Config from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig # --------------------------------------------------------------------------- diff --git a/tests/unit/llms/openai_like/test_abliteration_provider.py b/tests/unit/llms/openai_like/test_abliteration_provider.py index 33ad8abf62c..74bb5f89781 100644 --- a/tests/unit/llms/openai_like/test_abliteration_provider.py +++ b/tests/unit/llms/openai_like/test_abliteration_provider.py @@ -2,12 +2,18 @@ Unit tests for the Abliteration OpenAI-like provider. """ +import json import os import sys +from typing import Final -sys.path.insert( - 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../..")) -) +import pytest +import respx + +import litellm +from litellm.utils import trim_messages + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../.."))) from litellm.llms.openai_like.dynamic_config import create_config_class from litellm.llms.openai_like.json_loader import JSONProviderRegistry @@ -48,3 +54,79 @@ def test_abliteration_complete_url_appends_endpoint(): stream=False, ) assert url == f"{ABLITERATION_BASE_URL}/chat/completions" + + +def test_abliteration_chat_forwards_reasoning_and_accounts_for_cache(local_model_cost_map: None) -> None: + with respx.mock() as upstream: + route: Final = upstream.post("https://api.abliteration.ai/v1/chat/completions").respond( + 200, + json={ + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 1, + "model": "abliterated-model-large-v2", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "hello"}, "finish_reason": "stop"} + ], + "usage": { + "prompt_tokens": 100, + "completion_tokens": 5, + "total_tokens": 105, + "prompt_tokens_details": {"cached_tokens": 80}, + }, + }, + ) + response: Final = litellm.completion( + model="abliteration/abliterated-model-large-v2", + api_key="fake-test-key", + messages=[{"role": "user", "content": "Say hello"}], + reasoning_effort="low", + ) + request: Final = route.calls.last.request + body: Final = json.loads(request.content) + assert request.headers["authorization"] == "Bearer fake-test-key" + assert body["model"] == "abliterated-model-large-v2" + assert body["reasoning_effort"] == "low" + assert response.choices[0].message.content == "hello" + rates: Final = litellm.model_cost["abliteration/abliterated-model-large-v2"] + expected: Final = ( + 20 * rates["input_cost_per_token"] + + 80 * rates["cache_read_input_token_cost"] + + 5 * rates["output_cost_per_token"] + ) + assert litellm.completion_cost(completion_response=response) == pytest.approx(expected) + + +def test_large_v2_model_group_does_not_offer_disabled_reasoning(local_model_cost_map: None) -> None: + """Large v2 cannot disable reasoning per https://docs.abliteration.ai/capabilities/thinking (2026-10-26).""" + router: Final = litellm.Router( + model_list=[ + { + "model_name": "large-v2", + "litellm_params": { + "model": "abliteration/abliterated-model-large-v2", + "api_key": "fake-test-key", + }, + } + ] + ) + info: Final = router.get_model_group_info(model_group="large-v2") + assert info is not None + assert info.supported_reasoning_efforts + assert "none" not in info.supported_reasoning_efforts + assert "low" in info.supported_reasoning_efforts + + +def test_abliteration_models_trim_to_their_input_budget(local_model_cost_map: None) -> None: + model_names: Final = tuple(model for model in litellm.model_cost if model.startswith("abliteration/")) + assert model_names + messages: Final = [{"role": "user", "content": "hello " * 500}] + budget: Final = 100 + for model in model_names: + trimmed: Final = trim_messages( + messages, + model=model, + trim_ratio=budget / litellm.model_cost[model]["max_input_tokens"], + ) + assert litellm.token_counter(model=model, messages=messages) > budget + assert litellm.token_counter(model=model, messages=trimmed) <= budget diff --git a/tests/unit/llms/openai_like/test_refuseless_provider.py b/tests/unit/llms/openai_like/test_refuseless_provider.py new file mode 100644 index 00000000000..4961f100709 --- /dev/null +++ b/tests/unit/llms/openai_like/test_refuseless_provider.py @@ -0,0 +1,111 @@ +import json +from typing import Final + +import pytest +import respx + +import litellm +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.utils import trim_messages + + +def test_refuseless_resolves_environment_and_explicit_credentials(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("REFUSELESS_API_KEY", "fake-environment-key") + assert get_llm_provider(model="refuseless/glm-5.3-abliterated") == ( + "glm-5.3-abliterated", + "refuseless", + "fake-environment-key", + "https://api.refuseless.com/v1", + ) + assert get_llm_provider( + model="refuseless/glm-5.3-abliterated", api_key="fake-explicit-key", api_base="https://gateway.example/v1" + ) == ("glm-5.3-abliterated", "refuseless", "fake-explicit-key", "https://gateway.example/v1") + + +def test_refuseless_chat_forwards_reasoning_and_accounts_for_cache(local_model_cost_map: None) -> None: + with respx.mock() as upstream: + route: Final = upstream.post("https://api.refuseless.com/v1/chat/completions").respond( + 200, + json={ + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 1, + "model": "glm-5.3-abliterated", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "hello"}, "finish_reason": "stop"} + ], + "usage": { + "prompt_tokens": 100, + "completion_tokens": 5, + "total_tokens": 105, + "prompt_tokens_details": {"cached_tokens": 80}, + }, + }, + ) + response: Final = litellm.completion( + model="refuseless/glm-5.3-abliterated", + api_key="fake-test-key", + messages=[{"role": "user", "content": "Say hello"}], + reasoning_effort="low", + ) + request: Final = route.calls.last.request + body: Final = json.loads(request.content) + assert request.headers["authorization"] == "Bearer fake-test-key" + assert body["model"] == "glm-5.3-abliterated" + assert body["reasoning_effort"] == "low" + assert response.choices[0].message.content == "hello" + rates: Final = litellm.model_cost["refuseless/glm-5.3-abliterated"] + expected: Final = ( + 20 * rates["input_cost_per_token"] + + 80 * rates["cache_read_input_token_cost"] + + 5 * rates["output_cost_per_token"] + ) + assert litellm.completion_cost(completion_response=response) == pytest.approx(expected) + + +def test_refuseless_uses_native_responses_endpoint() -> None: + with respx.mock() as upstream: + route: Final = upstream.post("https://api.refuseless.com/v1/responses").respond( + 200, + json={ + "id": "resp-test", + "object": "response", + "created_at": 1, + "model": "glm-5.3-abliterated", + "status": "completed", + "output": [ + { + "id": "msg-test", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "hello", "annotations": []}], + } + ], + "usage": {"input_tokens": 4, "output_tokens": 3, "total_tokens": 7}, + }, + ) + response: Final = litellm.responses( + model="refuseless/glm-5.3-abliterated", + api_key="fake-test-key", + input="Say hello", + store=False, + ) + assert route.calls.last.request.headers["authorization"] == "Bearer fake-test-key" + assert json.loads(route.calls.last.request.content)["model"] == "glm-5.3-abliterated" + assert response.output[0].content[0].text == "hello" + + +def test_refuseless_models_trim_to_their_input_budget(local_model_cost_map: None) -> None: + model_names: Final = tuple(model for model in litellm.model_cost if model.startswith("refuseless/")) + assert model_names + messages: Final = [{"role": "user", "content": "hello " * 500}] + budget: Final = 100 + for model in model_names: + trimmed: Final = trim_messages( + messages, + model=model, + trim_ratio=budget / litellm.model_cost[model]["max_input_tokens"], + ) + assert litellm.token_counter(model=model, messages=messages) > budget + assert litellm.token_counter(model=model, messages=trimmed) <= budget diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/llama3/__init__.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/llama/__init__.py similarity index 100% rename from tests/unit/llms/vertex_ai/vertex_ai_partner_models/llama3/__init__.py rename to tests/unit/llms/vertex_ai/vertex_ai_partner_models/llama/__init__.py diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/llama3/test_vertex_ai_partner_models_llama3_transformation.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/llama/test_vertex_ai_partner_models_llama_transformation.py similarity index 98% rename from tests/unit/llms/vertex_ai/vertex_ai_partner_models/llama3/test_vertex_ai_partner_models_llama3_transformation.py rename to tests/unit/llms/vertex_ai/vertex_ai_partner_models/llama/test_vertex_ai_partner_models_llama_transformation.py index 05e4e36edd7..72efd7b66b0 100644 --- a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/llama3/test_vertex_ai_partner_models_llama3_transformation.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/llama/test_vertex_ai_partner_models_llama_transformation.py @@ -5,7 +5,7 @@ import httpx import pytest -from litellm.llms.vertex_ai.vertex_ai_partner_models.llama3.transformation import ( +from litellm.llms.vertex_ai.vertex_ai_partner_models.llama.transformation import ( VertexAILlama3Config, VertexAILlama3StreamingHandler, ) diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index a718426eea6..9ff2a6de9dd 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -3672,6 +3672,23 @@ class TestUpdateDBModelKeepsLegacyDropParams: assert decrypt_value_helper(value=stored["drop_params"], key="drop_params") == "true" +@pytest.mark.parametrize( + "payload", + [ + {"litellm_params_to_remove": ["model"]}, + {"litellm_params_to_remove": ["temperature"], "litellm_params": {"temperature": 0}}, + {"litellm_params_to_remove": ["temperature"], "litellm_params": {"temperature": None}}, + ], +) +def test_model_param_removal_rejects_required_and_conflicting_fields(payload: Mapping[str, object]) -> None: + from pydantic import ValidationError + + from litellm.types.proxy.management_endpoints.model_management_endpoints import PatchModelRequest + + with pytest.raises(ValidationError): + PatchModelRequest.model_validate(payload) + + def _build_db_model_with_pricing(): """Wildcard deployment with custom pricing in litellm_params; Deployment.__init__ mirrors SPECIAL_MODEL_INFO_PARAMS into model_info, so both blobs hold the rate.""" @@ -7843,6 +7860,55 @@ class TestTeamMemberAutoRouterWrites: "jev_classifier_config" if legacy else "opensource_classifier_config": classifier, } + @pytest.mark.asyncio + @pytest.mark.parametrize( + "stored_params,removed,denied_param", + [ + ({"litellm_credential_name": "stored-credential"}, "litellm_credential_name", "litellm_credential_name"), + ({"aws_session_tags": [{"Key": "scope", "Value": "restricted"}]}, "aws_session_tags", "aws_session_tags"), + ({"api_key": "oidc/file//var/run/token", "api_base": "https://provider.test"}, "api_base", "api_key"), + ({"anthropic_issuer_url": "https://issuer.test"}, "anthropic_issuer_url", "anthropic_issuer_url"), + ], + ) + async def test_team_admin_cannot_remove_admin_owned_model_params( + self, stored_params: Mapping[str, object], removed: str, denied_param: str + ) -> None: + from litellm.types.proxy.management_endpoints.model_management_endpoints import PatchModelRequest + + row: Final = self._row().model_copy( + update={"litellm_params": {"model": "openai/test-model", **stored_params}} + ) + team: Final = self._team().model_copy(update={"members_with_roles": [Member(user_id="owner", role="admin")]}) + database: Final = self._database(team, row) + database.db.litellm_credentialstable.find_unique = AsyncMock(return_value=None) + actor: Final = UserAPIKeyAuth(user_id="owner", user_role=LitellmUserRoles.INTERNAL_USER) + + with self._environment(database, row), pytest.raises(ProxyException) as denied: + await patch_model(row.model_id, PatchModelRequest(litellm_params_to_remove=(removed,)), actor) + + assert denied.value.code == "403" + assert denied.value.param == denied_param + database.db.litellm_proxymodeltable.update.assert_not_awaited() + + @pytest.mark.asyncio + async def test_param_removal_cannot_leave_a_strategy_router_without_required_configuration(self) -> None: + from litellm.types.proxy.management_endpoints.model_management_endpoints import PatchModelRequest + + row: Final = self._row() + database: Final = self._database(self._team(), row) + actor: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + + with self._environment(database, row), pytest.raises(ProxyException) as denied: + await patch_model( + row.model_id, + PatchModelRequest(litellm_params_to_remove=("complexity_router_config", "complexity_router_default_model")), + actor, + ) + + assert denied.value.code == "400" + assert "requires" in denied.value.message + database.db.litellm_proxymodeltable.update.assert_not_awaited() + @pytest.mark.asyncio @pytest.mark.parametrize("team_id", [None, "member-team"]) @pytest.mark.parametrize( diff --git a/tests/unit/proxy/management_endpoints/test_ptu_model_settings.py b/tests/unit/proxy/management_endpoints/test_ptu_model_settings.py index a319797a618..efb84ee8f3a 100644 --- a/tests/unit/proxy/management_endpoints/test_ptu_model_settings.py +++ b/tests/unit/proxy/management_endpoints/test_ptu_model_settings.py @@ -40,6 +40,7 @@ from litellm.proxy.management_endpoints.model_management_endpoints import ( from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR from litellm.repositories.chunked_in import IN_LIST_CHUNK_SIZE from litellm.router import Router +from litellm.types.proxy.management_endpoints.model_management_endpoints import PatchModelRequest from litellm.types.router import ( SPECIAL_MODEL_INFO_PARAMS, Deployment, @@ -1205,6 +1206,80 @@ class TestPtuDeploymentsAreNotBilledPerToken: for blob in ("model_info", "litellm_params"): assert json.loads(written[blob])["input_cost_per_token"] == 0, blob + @pytest.mark.parametrize("field,activating", [ + ("input_cost_per_token", False), ("cache_read_input_token_cost", False), + ("input_cost_per_second", False), ("tiered_pricing", False), ("search_context_cost_per_query", False), + ("input_cost_per_token", True), ("cache_read_input_token_cost", True), + ("tiered_pricing", True), ("search_context_cost_per_query", True), + ]) + def test_removing_ptu_managed_pricing_is_rejected(self, field: str, activating: bool) -> None: + original: Final = _deployment_without_ptu() if activating else _deployment_with_stored_ptu() + configured: Final = original.model_copy(update={ + "litellm_params": original.litellm_params.model_copy(update={"input_cost_per_second": 0.0}), + }) + with pytest.raises(HTTPException) as exc: + update_db_model( + db_model=configured if activating else _ptu_priced_deployment(configured), + updated_patch=PatchModelRequest( + litellm_params_to_remove=(field,), + model_info=_deployment_with_stored_ptu().model_info if activating else None, + ), + ) + assert exc.value.status_code == 400 + assert field in exc.value.detail + assert "ptu_count and cost_per_ptu_per_hour" in exc.value.detail + + @pytest.mark.parametrize("mode,field", [ + ("extra", "temperature"), ("disabled", "input_cost_per_token"), + ("disabled", "input_cost_per_second"), ("cleared", "input_cost_per_token"), + ("cleared", "input_cost_per_second"), ("activating", "input_cost_per_second"), + ("ordinary", "input_cost_per_second"), + ("ordinary", "input_cost_per_token_above_123k_tokens"), + ("ordinary", "regional_endpoint_uplift_multiplier"), + ]) + def test_allowed_ptu_removals_stay_absent_after_another_save( + self, monkeypatch: pytest.MonkeyPatch, mode: str, field: str + ) -> None: + original: Final = ( + _deployment_without_ptu() if mode in ("activating", "ordinary") + else _ptu_priced_deployment(_deployment_with_stored_ptu()) + ) + rate: Final = 0.5 if mode in ("activating", "ordinary") else 0.0 + price_field: Final = field if mode == "ordinary" else "input_cost_per_second" + configured: Final = original.model_copy(update={ + "litellm_params": original.litellm_params.model_copy(update={ + "temperature": 0.4, price_field: rate, + }), + "model_info": original.model_info.model_copy(update={price_field: rate}), + }) + if mode == "disabled": + monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR) + incoming_info: Final = ( + ModelInfo(ptu_count=None, cost_per_ptu_per_hour=None) if mode == "cleared" + else _deployment_with_stored_ptu().model_info if mode == "activating" else ModelInfo() + ) + patch: Final = PatchModelRequest( + litellm_params_to_remove=(field,), + model_info=incoming_info.model_copy(update={price_field: rate}), + ) + written: Final = update_db_model(db_model=configured, updated_patch=patch) + assert json.loads(written["model_info"]).get("ptu_count") == ( + None if mode in ("cleared", "ordinary") else 15 + ) + reloaded: Final = Deployment.model_validate({ + "model_name": written["model_name"], "litellm_params": json.loads(written["litellm_params"]), + "model_info": json.loads(written["model_info"]), + }) + saved_again: Final = update_db_model( + db_model=reloaded, updated_patch=updateDeployment( + model_name="renamed", model_info=ModelInfo.model_validate({price_field: rate}), + ) + ) + retried: Final = update_db_model(db_model=reloaded, updated_patch=patch) + for result in (written, saved_again, retried): + assert field not in json.loads(result["litellm_params"]) + assert field not in json.loads(result["model_info"]) + def test_a_patch_pricing_a_ptu_deployment_is_refused(self): with pytest.raises(HTTPException) as exc: update_db_model( diff --git a/tests/unit/proxy/proxy_server/test_routes_model_info.py b/tests/unit/proxy/proxy_server/test_routes_model_info.py index 2c0e0944edb..9993361d930 100644 --- a/tests/unit/proxy/proxy_server/test_routes_model_info.py +++ b/tests/unit/proxy/proxy_server/test_routes_model_info.py @@ -11,6 +11,7 @@ from __future__ import annotations import asyncio import copy +import json from collections.abc import Callable from contextlib import AbstractContextManager from typing import Final @@ -24,6 +25,7 @@ import litellm from litellm.caching.llm_caching_handler import LLMClientCache from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy import proxy_server +from litellm.proxy._types import LiteLLM_ProxyModelTable from litellm.utils import _invalidate_model_cost_lowercase_map from .conftest import normalize # type: ignore[import-not-found] @@ -206,6 +208,80 @@ def test_v2_model_info_invalid_page_returns_422(client, auth_as, empty_router): assert "detail" in response.json() +@pytest.mark.parametrize("encoded", [False, True]) +def test_deleted_model_params_stay_absent_after_reload_and_another_save( + client: TestClient, + auth_as: Callable[[], AbstractContextManager[object]], + monkeypatch: pytest.MonkeyPatch, + mock_prisma: MagicMock, + encoded: bool, +) -> None: + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + from litellm.types.proxy.management_endpoints.model_management_endpoints import PatchModelRequest + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo, updateDeployment, updateLiteLLMParams + + configured: Final = { + "model": "openai/test-model", "api_key": "synthetic-runtime-secret", + "use_litellm_proxy": True, "use_in_pass_through": True, "use_chat_completions_api": False, + "litellm_image_enabled": True, "timeout": 30, "input_cost_per_token": 0.003, + "team_id": "param-only-team", "temperature": 0, "metadata": {}, + } + original: Final = Deployment( + model_name="editor-model", litellm_params=LiteLLM_Params.model_validate(configured), + model_info=ModelInfo(id="editor-model", team_id="owner-team", access_groups=["retained-group"]), + ) + removals: Final = ("litellm_image_enabled", "use_litellm_proxy", "input_cost_per_token", "team_id") + deleted: Final = update_db_model( + db_model=original, + updated_patch=PatchModelRequest( + litellm_params=updateLiteLLMParams(max_retries=0), litellm_params_to_remove=removals, + model_info=ModelInfo(id="editor-model", input_cost_per_token=0.003), + ), + ) + saved_params: Final = json.loads(deleted["litellm_params"]) + saved_info: Final = json.loads(deleted["model_info"]) + assert saved_params == {**{key: value for key, value in configured.items() if key not in removals}, "max_retries": 0} + assert "input_cost_per_token" not in saved_info + assert saved_info["team_id"] == "owner-team" + assert saved_info["access_groups"] == ["retained-group"] + reloaded: Final = Deployment.model_validate( + {"model_name": original.model_name, "litellm_params": saved_params, "model_info": saved_info} + ) + router: Final = litellm.Router(model_list=[{ + "model_name": "editor-model", "litellm_params": saved_params, + "model_info": {"id": "editor-model", "db_model": True}, + }]) + stored_params: Final = {**saved_params, "api_key": "synthetic-stored-secret"} + mock_prisma.db.litellm_proxymodeltable.find_unique.return_value = LiteLLM_ProxyModelTable( + model_id="editor-model", model_name="editor-model", + litellm_params=json.dumps(stored_params) if encoded else stored_params, + model_info={"id": "editor-model"}, + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "prisma_client", mock_prisma) + + with auth_as(): + normal: Final = client.get("/v2/model/info", params={"modelId": "editor-model"}) + editor: Final = client.get("/v2/model/info", params={"modelId": "editor-model", "use_stored_params": "true"}) + + assert normal.status_code == editor.status_code == 200, editor.text + normal_params: Final = normal.json()["data"][0]["litellm_params"] + assert "use_litellm_proxy" in normal_params + assert editor.json()["data"][0]["litellm_params"] == { + key: value for key, value in normal_params.items() if key in saved_params + } + assert editor.json()["data"][0]["litellm_params"]["use_chat_completions_api"] is False + assert "synthetic-runtime-secret" not in editor.text + assert "synthetic-stored-secret" not in editor.text + assert "use_litellm_proxy" in router.model_list[0]["litellm_params"] + + saved_again: Final = update_db_model( + db_model=reloaded, updated_patch=updateDeployment(litellm_params=updateLiteLLMParams(rpm=10, timeout=None)) + ) + persisted: Final = json.loads(saved_again["litellm_params"]) + assert persisted == {**saved_params, "rpm": 10} + + def test_v2_model_info_in_openapi_schema(): """``GET /v2/model/info`` is published in the proxy OpenAPI/Swagger spec.""" from litellm.proxy.proxy_server import get_openapi_schema diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index c39fbeeb220..a44373b87c8 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -927,6 +927,222 @@ async def test_add_litellm_data_to_request_body_snapshot_excludes_proxy_server_r assert snapshot_body["messages"] == [{"role": "user", "content": "hello"}] +FORWARD_PROVIDER_AUTH_HEADERS: Final = { + "forward_llm_provider_auth_headers": True, + "forward_client_headers_to_llm_api": True, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "sent,request_headers,general_settings,expected", + [ + ({"headers": {"Authorization": "Bearer sk-caller"}}, {}, {}, ("headers",)), + ({"headers": {"X-Trace-Id": "trace-1"}}, {}, {}, ("headers",)), + ({"headers": {}}, {}, {}, ()), + ({"extra_headers": {"API-Key": "caller-azure-key"}}, {}, {}, ("extra_headers",)), + ({"extra_headers": {"OpenAI-Project": "proj_caller"}}, {}, {}, ("extra_headers",)), + ( + {"provider_specific_header": {"custom_llm_provider": "openai", "extra_headers": {"x-api-key": "caller"}}}, + {}, + {}, + ("provider_specific_header",), + ), + ( + { + "provider_specific_header": [ + {"custom_llm_provider": "anthropic", "extra_headers": {"anthropic-beta": "tools"}}, + {"custom_llm_provider": "gemini", "extra_headers": {"X-Goog-Api-Key": "caller"}}, + ] + }, + {}, + {}, + ("provider_specific_header",), + ), + ( + { + "provider_specific_header": { + "custom_llm_provider": "anthropic", + "extra_headers": {"x-caller-trace": "x"}, + } + }, + {}, + {}, + ("provider_specific_header",), + ), + ( + { + "provider_specific_header": { + "custom_llm_provider": "anthropic", + "extra_headers": {"anthropic-version": "2023-06-01", "anthropic-beta": "tools-2024-04-04"}, + } + }, + {}, + {}, + (), + ), + ({"headers": {"Anthropic-Version": "2023-06-01"}}, {}, {}, ()), + ({"organization": "org-caller"}, {}, {}, ("organization",)), + ({"headers": "authorization", "provider_specific_header": "x-api-key"}, {}, {}, ()), + ( + {"api_key": "sk-caller", "headers": {"Ocp-Apim-Subscription-Key": "caller"}}, + {}, + {}, + ("api_key", "headers"), + ), + ({}, {"anthropic-beta": "tools-2024-04-04"}, {}, ()), + ({}, {"anthropic-version": "2023-06-01", "anthropic-beta": "tools-2024-04-04"}, {}, ()), + ({}, {"openai-organization": "org-caller"}, {"forward_openai_org_id": True}, ("organization",)), + ( + {}, + {"x-litellm-api-key": "sk-virtual", "authorization": "Bearer sk-ant-oat01-caller"}, + {}, + ("provider_specific_header",), + ), + ( + {}, + {"authorization": "Bearer sk-virtual", "x-goog-api-key": "AIza-caller"}, + FORWARD_PROVIDER_AUTH_HEADERS, + ("headers",), + ), + ( + {}, + {"authorization": "Bearer sk-virtual", "x-portkey-api-key": "caller-gateway-key"}, + FORWARD_PROVIDER_AUTH_HEADERS, + ("headers",), + ), + ({}, {"authorization": "Bearer sk-virtual"}, FORWARD_PROVIDER_AUTH_HEADERS, ()), + ], + ids=[ + "authorization-in-headers", + "plain-headers", + "empty-headers", + "api-key-in-extra-headers", + "openai-project-in-extra-headers", + "credential-in-provider-specific-header", + "credential-in-one-provider-specific-header-entry", + "plain-provider-specific-header", + "anthropic-version-and-beta-in-provider-specific-header", + "anthropic-version-in-headers", + "organization", + "malformed-header-fields", + "api-key-and-credential-headers", + "proxy-forwarded-anthropic-beta", + "proxy-forwarded-anthropic-sdk-headers", + "proxy-forwarded-openai-organization", + "proxy-forwarded-anthropic-oauth-token", + "proxy-forwarded-provider-auth-header", + "proxy-forwarded-gateway-key", + "proxy-forwarding-nothing", + ], +) +async def test_caller_credential_fields_name_every_field_that_carries_a_caller_header_or_credential( + sent: dict[str, object], + request_headers: dict[str, str], + general_settings: dict[str, object], + expected: tuple[str, ...], +) -> None: + updated: Final = await _request_data_added_by_the_proxy(sent, request_headers, general_settings) + + assert updated["proxy_server_request"]["caller_credential_fields"] == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "request_headers,expected", + [({}, ()), ({"x-caller-trace": "trace-1"}, ("headers",))], + ids=["user-information-only", "user-information-and-caller-header"], +) +async def test_caller_credential_fields_ignore_the_user_information_the_proxy_adds_to_forwarded_headers( + monkeypatch: pytest.MonkeyPatch, request_headers: dict[str, str], expected: tuple[str, ...] +) -> None: + monkeypatch.setattr(litellm, "add_user_information_to_llm_headers", True) + + updated: Final = await _request_data_added_by_the_proxy({}, request_headers, FORWARD_PROVIDER_AUTH_HEADERS) + + assert updated["headers"]["x-litellm-user_api_key_user_id"] == "test-user" + assert updated["proxy_server_request"]["caller_credential_fields"] == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "sent,request_headers,general_settings,expected", + [ + ({"headers": {"Anthropic-Version": "2023-06-01"}}, {}, {}, ("headers",)), + ( + { + "provider_specific_header": { + "custom_llm_provider": "anthropic", + "extra_headers": {"anthropic-version": "2023-06-01"}, + } + }, + {}, + {}, + ("provider_specific_header",), + ), + ({"organization": "org-caller"}, {}, {}, ()), + ({"extra_headers": {"API-Key": "caller-azure-key"}}, {}, {}, ()), + ( + {"api_key": "sk-caller", "headers": {"Ocp-Apim-Subscription-Key": "caller"}}, + {}, + {}, + ("api_key", "headers"), + ), + ({}, {"openai-organization": "org-caller"}, {"forward_openai_org_id": True}, ()), + ( + {}, + {"authorization": "Bearer sk-virtual", "x-goog-api-key": "AIza-caller"}, + FORWARD_PROVIDER_AUTH_HEADERS, + (), + ), + ], + ids=[ + "anthropic-version-in-headers", + "anthropic-version-in-provider-specific-header", + "organization", + "api-key-in-extra-headers", + "api-key-and-credential-headers", + "proxy-forwarded-openai-organization", + "proxy-forwarded-provider-auth-header", + ], +) +async def test_credential_fields_name_the_stripped_fields_the_caller_sent( + sent: dict[str, object], + request_headers: dict[str, str], + general_settings: dict[str, object], + expected: tuple[str, ...], +) -> None: + updated: Final = await _request_data_added_by_the_proxy(sent, request_headers, general_settings) + + assert updated["proxy_server_request"]["credential_fields"] == expected + + +async def _request_data_added_by_the_proxy( + sent: dict[str, object], request_headers: dict[str, str], general_settings: dict[str, object] +) -> dict: + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request + + request_mock: Final = MagicMock(spec=Request) + request_mock.url = MagicMock() + request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions" + request_mock.url.path = "/v1/chat/completions" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = {"content-type": "application/json", **request_headers} + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + data: Final = {"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hello"}], **sent} + + return await add_litellm_data_to_request( + data=data, + request=request_mock, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key", user_id="test-user"), + proxy_config=MagicMock(), + general_settings=dict(general_settings), + version="test-version", + ) + + def test_initial_snapshot_refresh_clears_a_previous_guardrail_checkpoint() -> None: from litellm.integrations.shadow_eval_logger import GuardrailRequestSnapshot from litellm.litellm_core_utils.litellm_logging import Logging diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 8477cdc72cc..8c62b225992 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -15586,4 +15586,3 @@ async def test_spend_capture_rate_check_job_clears_the_gauge_once_the_setting_is call(api_provider="openai", capture_rate=0.97), call(api_provider="openai", capture_rate=None), ] - diff --git a/tests/unit/router_utils/test_reasoning_effort_capability.py b/tests/unit/router_utils/test_reasoning_effort_capability.py index d634e239ea9..d0564812c48 100644 --- a/tests/unit/router_utils/test_reasoning_effort_capability.py +++ b/tests/unit/router_utils/test_reasoning_effort_capability.py @@ -203,7 +203,7 @@ class TestNoneLevelPolarity: def test_azure_advertisement_matches_the_azure_request_gate(self, model_key): """AzureOpenAIGPT5Config raises UnsupportedParamsError on reasoning_effort='none' for models it does not flag, so advertising the level there would offer routing a 400.""" - from litellm.llms.azure.chat.gpt_5_transformation import AzureOpenAIGPT5Config + from litellm.llms.azure.chat.gpt_reasoning_transformation import AzureOpenAIGPT5Config from litellm.utils import get_model_info_helper model_info = dict(get_model_info_helper(model=model_key.split("/", 1)[1], custom_llm_provider="azure")) diff --git a/tests/unit/test_model_prices_and_context_window.py b/tests/unit/test_model_prices_and_context_window.py index 35658097880..1f84be7710f 100644 --- a/tests/unit/test_model_prices_and_context_window.py +++ b/tests/unit/test_model_prices_and_context_window.py @@ -12,7 +12,7 @@ from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter import litellm from litellm.litellm_core_utils.fallback_generalizations import match_capability_generalizations -from litellm.llms.openai.chat.gpt_5_transformation import is_gpt_reasoning_series_name +from litellm.llms.openai.chat.gpt_reasoning_transformation import is_gpt_reasoning_series_name from litellm.llms.openai_like.json_loader import JSONProviderRegistry from litellm.router_utils.reasoning_effort_capability import resolve_supported_reasoning_efforts from litellm.types.utils import SearchFreeRequests diff --git a/tests/unit/test_private_usage_aliases.py b/tests/unit/test_private_usage_aliases.py index 9a6fe8d2f49..924d1f1fc2b 100644 --- a/tests/unit/test_private_usage_aliases.py +++ b/tests/unit/test_private_usage_aliases.py @@ -2233,7 +2233,7 @@ LLMS_ALIAS_CASES: Final = ( ("litellm.llms.oci.chat.generic", "", "_normalize_oci_finish_reason", "normalize_oci_finish_reason", False), ("litellm.llms.oci.chat.generic", "", "_synthesize_oci_tool_call_id", "synthesize_oci_tool_call_id", False), ("litellm.llms.ollama.common_utils", "", "_convert_image", "convert_image", False), - ("litellm.llms.openai.chat.gpt_5_transformation", "", "_get_effort_level", "get_effort_level", False), + ("litellm.llms.openai.chat.gpt_reasoning_transformation", "", "_get_effort_level", "get_effort_level", False), ("litellm.llms.openai.completion.utils", "", "_transform_prompt", "transform_prompt", False), ( "litellm.llms.openai.cost_calculation", @@ -4867,7 +4867,7 @@ LLMS_FORWARDER_CASES: Final = ( False, ), ( - "litellm.llms.openai.chat.gpt_5_transformation", + "litellm.llms.openai.chat.gpt_reasoning_transformation", "OpenAIGPT5Config", "_supports_reasoning_effort_level", "supports_reasoning_effort_level", diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 6c7f6a9d869..38287b25169 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -576,7 +576,7 @@ def test_all_model_configs(): from litellm.llms.vertex_ai.vertex_ai_partner_models.ai21.transformation import ( VertexAIAi21Config, ) - from litellm.llms.vertex_ai.vertex_ai_partner_models.llama3.transformation import ( + from litellm.llms.vertex_ai.vertex_ai_partner_models.llama.transformation import ( VertexAILlama3Config, ) @@ -2610,7 +2610,7 @@ def test_priced_router_deployment_id_stays_out_of_wildcard_model_listing(): def test_anthropic_claude_4_invoke_chat_provider_config(): """Test that the Anthropic Claude 4 Invoke chat provider config is correct.""" - from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import ( + from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude_transformation import ( AmazonAnthropicClaudeConfig, ) @@ -5156,7 +5156,7 @@ def test_request_path_cost_map_walks_survive_a_registration_mid_walk( from litellm.llms.fal_ai.cost_calculator import _keyed_rows from litellm.llms.fireworks_ai.chat.transformation import FireworksAIConfig from litellm.llms.ollama.common_utils import OllamaModelInfo - from litellm.llms.openai.chat.gpt_5_transformation import _catalogue_declares_default_effort + from litellm.llms.openai.chat.gpt_reasoning_transformation import _catalogue_declares_default_effort from litellm.utils import _handle_new_key_with_scan, _invalidate_model_cost_lowercase_map walks: Final[Mapping[str, Callable[[], object]]] = { diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 32777b43504..c8231b90bfd 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -1730,9 +1730,6 @@ "src/components/model_group_alias_settings.tsx": { "local/filename-pascal-case": { "count": 1 - }, - "react-hooks/set-state-in-effect": { - "count": 1 } }, "src/components/model_info_view.tsx": { @@ -1740,7 +1737,7 @@ "count": 1 }, "prefer-const": { - "count": 5 + "count": 3 }, "react-hooks/set-state-in-effect": { "count": 1 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts index e9df939f220..76db710f5ba 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts @@ -125,10 +125,42 @@ describe("useModelsInfo", () => { undefined, false, undefined, + {}, ); expect(modelInfoCall).toHaveBeenCalledTimes(1); }); + it("keeps saved model parameters separate from hydrated parameters in the query cache", async () => { + vi.mocked(modelInfoCall).mockImplementation(async (...args) => ({ + ...mockPaginatedModelInfoResponse, + data: [{ litellm_params: args[15]?.useStoredParams ? {} : { litellm_image_enabled: false } }], + })); + const { result } = renderHook( + () => ({ + hydrated: useModelsInfo(1, 50, undefined, "model-1"), + stored: useModelsInfo( + 1, + 50, + undefined, + "model-1", + undefined, + undefined, + undefined, + false, + undefined, + undefined, + false, + undefined, + { useStoredParams: true }, + ), + }), + { wrapper }, + ); + await waitFor(() => expect(result.current.stored.isSuccess && result.current.hydrated.isSuccess).toBe(true)); + expect(result.current.stored.data?.data[0].litellm_params).toEqual({}); + expect(result.current.hydrated.data?.data[0].litellm_params).toEqual({ litellm_image_enabled: false }); + }); + it("should use custom page and size parameters", async () => { (modelInfoCall as any).mockResolvedValue(mockPaginatedModelInfoResponse); @@ -156,6 +188,7 @@ describe("useModelsInfo", () => { undefined, false, undefined, + {}, ); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts index 744db87ad11..f9304bab0d2 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts @@ -45,6 +45,7 @@ export const useModelsInfo = ( accessGroup?: string, wildcardOnly: boolean = false, blocked?: boolean, + options: { useStoredParams?: boolean } = {}, ) => { const { accessToken, userId, userRole } = useAuthorized(); return useQuery({ @@ -57,6 +58,7 @@ export const useModelsInfo = ( ...(search && { search }), ...(modelName && { modelName }), ...(modelId && { modelId }), + useStoredParams: Boolean(options.useStoredParams), ...(teamId && { teamId }), ...(sortBy && { sortBy }), ...(sortOrder && { sortOrder }), @@ -86,6 +88,7 @@ export const useModelsInfo = ( accessGroup, wildcardOnly, blocked, + options, ), enabled: Boolean(accessToken && userId && userRole), }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/ModelGroupAliasPanel.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/ModelGroupAliasPanel.integration.test.tsx index 227a34335a3..36a9f92dfd4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/ModelGroupAliasPanel.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/ModelGroupAliasPanel.integration.test.tsx @@ -10,6 +10,7 @@ vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ vi.mock("@/components/networking", () => ({ getCallbacksCall: vi.fn(), + getProxyBaseUrl: () => "https://proxy.test", getRouterSettingsCall: vi.fn(), setCallbacksCall: vi.fn(), })); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/ModelGroupAliasPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/ModelGroupAliasPanel.tsx index 025923e08e9..e6c720f38cf 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/ModelGroupAliasPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/ModelGroupAliasPanel.tsx @@ -2,7 +2,7 @@ import { useEffect, useState } from "react"; import ModelGroupAliasSettings from "@/components/model_group_alias_settings"; -import { getCallbacksCall, getRouterSettingsCall } from "@/components/networking"; +import { getCallbacksCall, getProxyBaseUrl, getRouterSettingsCall } from "@/components/networking"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; export default function ModelGroupAliasPanel() { @@ -22,6 +22,7 @@ export default function ModelGroupAliasPanel() { : null; const managedByConfig = currentAliasOwnership?.managedByConfig ?? false; const ownershipLoading = Boolean(accessToken && !currentAliasOwnership); + const preferenceScope = userID ? JSON.stringify([getProxyBaseUrl(), userID]) : null; useEffect(() => { if (!accessToken) { @@ -62,6 +63,8 @@ export default function ModelGroupAliasPanel() { return ( { vi.unstubAllGlobals(); }); - it("lists the key's decision models from the proxy that answers the request, sorted by name", async () => { + it("lists the key's models with decision modes or endpoints from the request proxy, sorted by name", async () => { mockFetch.mockResolvedValueOnce( modelGroupInfo([ { model_group: "pplx-decider", mode: "evaluation" }, { model_group: "gpt-5.5", mode: "chat" }, { model_group: "jev-latest", mode: "evaluation" }, { model_group: "unknown-mode", mode: null }, + { model_group: "native-decision", mode: "decisions" }, + { + model_group: "chat-decisions", + mode: "chat", + supported_endpoints: ["/v1/chat/completions", "/v1/decisions"], + }, + { model_group: "chat-systemone", mode: "chat", supported_endpoints: ["/v1/systemone"] }, + { model_group: "custom-decision", mode: "custom", supported_endpoints: ["/v1/decisions"] }, + { model_group: "endpoint-only", supported_endpoints: ["/v1/decisions"] }, + { model_group: "chat-only", mode: "chat", supported_endpoints: ["/v1/chat/completions"] }, + { model_group: "no-endpoints", mode: "chat", supported_endpoints: [] }, + { model_group: "null-endpoints", mode: "chat", supported_endpoints: null }, + { model_group: "missing-mode" }, + { model_group: "custom-mode", mode: "custom" }, ]), ); await expect(fetchDecisionModels("virtual-key", "https://custom.example.com/")).resolves.toEqual([ + "chat-decisions", + "chat-systemone", + "custom-decision", + "endpoint-only", "jev-latest", + "native-decision", "pplx-decider", ]); expect(mockFetch).toHaveBeenCalledWith( diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/system_one.ts b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/system_one.ts index a4030029903..035a05729ba 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/system_one.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/system_one.ts @@ -2,7 +2,7 @@ import { z } from "zod"; import { getGatewayBaseUrl, getGlobalLitellmHeaderName, getProxyBaseUrl } from "@/components/networking"; import { withRequiredHeaders } from "@/components/llm_calls/request_headers"; import { createApiClient } from "@/lib/http/client"; -import { isDecisionMode } from "@/lib/decisionModels"; +import { isDecisionEntry } from "@/lib/decisionModels"; import { systemOneResponseSchema, type PlaygroundRequest, @@ -26,7 +26,13 @@ export type PlaygroundDecisionRequest = | { endpoint: "/v1/decisions"; payload: OpenAIDecisionsRequest }; const modelGroupInfoSchema = z.object({ - data: z.array(z.object({ model_group: z.string(), mode: z.string().nullish() })), + data: z.array( + z.object({ + model_group: z.string(), + mode: z.string().nullish(), + supported_endpoints: z.array(z.string()).nullish(), + }), + ), }); function proxyClient(customBaseUrl?: string, getDefaultBaseUrl: () => string = getProxyBaseUrl) { @@ -56,7 +62,7 @@ export async function fetchDecisionModels( }); return modelGroupInfoSchema .parse(body) - .data.filter((group) => isDecisionMode(group.mode)) + .data.filter(isDecisionEntry) .map((group) => group.model_group) .sort((a, b) => a.localeCompare(b)); } diff --git a/ui/litellm-dashboard/src/components/ModelInfoEditForm.tsx b/ui/litellm-dashboard/src/components/ModelInfoEditForm.tsx index 70fe23202bf..085334567ed 100644 --- a/ui/litellm-dashboard/src/components/ModelInfoEditForm.tsx +++ b/ui/litellm-dashboard/src/components/ModelInfoEditForm.tsx @@ -34,7 +34,7 @@ import type { Tag } from "./tag_management/types"; import { ModelTeamSelect } from "./view_model/ModelTeamSelect"; import VectorStoreSelector from "./vector_store_management/VectorStoreSelector"; import { formatPtuUtcDisplay, utcIsoToPickerValue } from "../utils/ptuDatetime"; -import { isMaskedSecret } from "../utils/maskedSecretUtils"; +import { editableModelExtraParams, parseModelExtraParams } from "../utils/modelEditParams"; import { MAX_COST_PER_PTU_PER_HOUR, MAX_PTU_COUNT, @@ -127,22 +127,13 @@ const modelEditShape = { type ModelEditFieldName = keyof typeof modelEditShape; -const isJson = (value: string): boolean => { - try { - JSON.parse(value); - return true; - } catch { - return false; - } -}; - const buildSchema = (ptuEnabled: boolean, isFieldTouched: (field: TouchedPricingField) => boolean) => z.object(modelEditShape).superRefine((values, ctx) => { const reject = (path: ModelEditFieldName, message: string) => ctx.addIssue({ code: "custom", path: [path], message }); - if (values.litellm_extra_params && !isJson(values.litellm_extra_params)) { - reject("litellm_extra_params", "Please enter valid JSON"); + if (parseModelExtraParams(values.litellm_extra_params) === null) { + reject("litellm_extra_params", "Please enter a valid JSON object"); } // antd validates only mounted fields, and the PTU block does not render when the flag is off. @@ -245,15 +236,7 @@ export const toModelEditFormValues = (localModelData: any, isWildcardModel: bool // antd never mounted this field for a non-wildcard model, so the key must be absent, not null. ...(isWildcardModel ? { health_check_model: localModelData.model_info?.health_check_model } : {}), litellm_credential_name: localModelData.litellm_params?.litellm_credential_name ?? null, - litellm_extra_params: JSON.stringify( - Object.fromEntries( - Object.entries(localModelData.litellm_params || {}).filter( - ([key, value]) => key !== "litellm_credential_name" && !isMaskedSecret(value), - ), - ), - null, - 2, - ), + litellm_extra_params: JSON.stringify(editableModelExtraParams(localModelData.litellm_params ?? {}), null, 2), team_id: localModelData.model_info?.team_id ?? undefined, }); @@ -775,7 +758,7 @@ const ModelInfoEditForm: React.FC = ({ LiteLLM Params @@ -784,9 +767,10 @@ const ModelInfoEditForm: React.FC = ({ {({ value, ...control }) => (