diff --git a/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py b/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py index a4a40a7a3d7..1c433d04819 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py @@ -182,7 +182,7 @@ class ManagedObjectCreatorAttribution: team_id: Final = getattr(job, "team_id", None) rows: Final = _CreatorRows( key=await self._live_key_row(api_key, job.created_by), - team=await self._live_team_row(team_id), + team=await self._live_team_row(team_id, job.created_at), user=await self._user_row(object_id, _creator_user_id(job.created_by)), ) return ResponseCreator( @@ -269,10 +269,15 @@ class ManagedObjectCreatorAttribution: verbose_proxy_logger.error(f"{self._poller_name}: could not look up team {team_id} for {object_id}: {e}") return None - async def _live_team_row(self, team_id: str | None) -> prisma_models.LiteLLM_TeamTable | None: + async def _live_team_row( + self, team_id: str | None, created_before: datetime | None = None + ) -> prisma_models.LiteLLM_TeamTable | None: if not team_id: return None - return await _team_table(self._prisma_client).find_unique(where={"team_id": team_id}) + team: Final = await _team_table(self._prisma_client).find_unique(where={"team_id": team_id}) + if team is not None and created_before is not None and team.created_at > created_before: + return None + return team def _response_creator_metadata(job: ManagedObjectRow, rows: _CreatorRows) -> Mapping[str, object]: diff --git a/litellm-rust/crates/inference-messages/src/types.rs b/litellm-rust/crates/inference-messages/src/types.rs index ae346259559..ece17b3380c 100644 --- a/litellm-rust/crates/inference-messages/src/types.rs +++ b/litellm-rust/crates/inference-messages/src/types.rs @@ -144,6 +144,7 @@ mod tests { supports_reasoning: true, supports_adaptive_thinking: true, thinking_always_on: false, + thinking_on_by_default: false, supports_legacy_thinking: false, supports_output_config: true, supports_sampling_params: false, diff --git a/litellm-rust/crates/llms/src/anthropic/common_utils.rs b/litellm-rust/crates/llms/src/anthropic/common_utils.rs index 13653f438f7..0854820358b 100644 --- a/litellm-rust/crates/llms/src/anthropic/common_utils.rs +++ b/litellm-rust/crates/llms/src/anthropic/common_utils.rs @@ -1913,6 +1913,7 @@ mod tests { supports_reasoning: false, supports_adaptive_thinking: false, thinking_always_on: false, + thinking_on_by_default: false, supports_legacy_thinking: false, supports_output_config: false, supports_sampling_params: true, diff --git a/litellm-rust/crates/llms/src/anthropic/messages/thinking.rs b/litellm-rust/crates/llms/src/anthropic/messages/thinking.rs index ae3488a21f9..0ddbc915154 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/thinking.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/thinking.rs @@ -2,8 +2,8 @@ use litellm_llms_types::{ formats::{ chat_completions::ReasoningEffort, messages::{ - AdaptiveThinking, EffortLevel, MessagesOptionalParams, MessagesRequest, OutputConfig, - ThinkingConfig, ThinkingDisplay, + AdaptiveThinking, DisabledThinking, EffortLevel, MessagesOptionalParams, + MessagesRequest, OutputConfig, ThinkingConfig, ThinkingDisplay, }, }, recognized::Recognized, @@ -172,7 +172,11 @@ fn translate_reasoning_effort( ) else { return Ok(MessagesRequest { params: MessagesOptionalParams { - thinking: None, + thinking: (context.capabilities.thinking_on_by_default + && !context.capabilities.thinking_always_on) + .then(|| { + Recognized::Known(ThinkingConfig::Disabled(DisabledThinking::default())) + }), output_config: None, ..request.params }, @@ -748,6 +752,33 @@ mod tests { ); } + #[rstest] + #[case::default_thinking(true, false, Some(serde_json::json!({"type": "disabled"}))) ] + #[case::legacy(false, false, None)] + #[case::always_on(true, true, None)] + fn none_disables_only_optional_default_thinking( + #[case] thinking_on_by_default: bool, + #[case] thinking_always_on: bool, + #[case] thinking: Option, + ) { + let capabilities = MessagesModelCapabilities { + thinking_on_by_default, + thinking_always_on, + ..opus_4_7() + }; + let expected = match thinking { + Some(thinking) => serde_json::json!({"max_tokens": 1024, "thinking": thinking}), + None => serde_json::json!({"max_tokens": 1024}), + }; + assert_eq!( + translate( + capabilities, + serde_json::json!({"max_tokens": 1024, "reasoning_effort": "none"}) + ), + Ok(request(expected)) + ); + } + #[rstest] #[case::adaptive_model(opus_4_7())] #[case::effort_capable_model(opus_4_5())] diff --git a/litellm-rust/crates/llms/src/base_llm/messages/context.rs b/litellm-rust/crates/llms/src/base_llm/messages/context.rs index b514e975aa4..3d8b0f41d8c 100644 --- a/litellm-rust/crates/llms/src/base_llm/messages/context.rs +++ b/litellm-rust/crates/llms/src/base_llm/messages/context.rs @@ -32,6 +32,8 @@ pub struct MessagesModelCapabilities { #[serde(default)] pub thinking_always_on: bool, #[serde(default)] + pub thinking_on_by_default: bool, + #[serde(default)] pub supports_legacy_thinking: bool, #[serde(default)] pub supports_output_config: bool, @@ -55,6 +57,7 @@ impl Default for MessagesModelCapabilities { supports_reasoning: false, supports_adaptive_thinking: false, thinking_always_on: false, + thinking_on_by_default: false, supports_legacy_thinking: false, supports_output_config: false, supports_sampling_params: true, diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 09b787b2c94..0883dfa7d31 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -55,7 +55,9 @@ mod _native { #[pymodule_export] use crate::redis::store::sensitive_route_pins::NativeSensitiveRoutePins; #[pymodule_export] - use crate::redis::store::spend_counters::{NativeSpendCounters, PySpendRead}; + use crate::redis::store::spend_counters::{ + NativeSpendCounters, PyPostCallSettlement, PySpendRead, + }; #[pymodule_export] use crate::redis::store::{ NativePending, NativeRedisStoreError, router_cooldowns::NativeRouterCooldowns, @@ -154,6 +156,7 @@ mod tests { "NativeSensitiveRoutePins", "NativeRedisStoreError", "NativeSpendCounters", + "PostCallSettlement", "SpendRead", "LocalSizes", "native_redis_breaker_states", diff --git a/litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs b/litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs index ba016821e48..9da6a0e00e5 100644 --- a/litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs +++ b/litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs @@ -2,12 +2,12 @@ use std::{collections::HashMap, sync::Arc}; use litellm_redis_store::{ Number, - spend_counters::{LostSeed, Reading, SpendCounter, SpendCounters}, + spend_counters::{LostSeed, PostCallSettlement, Reading, SpendCounter, SpendCounters}, }; use pyo3::prelude::*; use super::{ - NativePending, PyData, PyNumber, kind_error, pending, redis_async, remote, + NativePending, PyData, PyNumber, duration, kind_error, pending, redis_async, remote, request::{NativeRequestRedis, NativeSend}, ttl, unwrap_all, }; @@ -29,6 +29,18 @@ impl PySpendRead { } } +/// Whether every increment one `increment_many_post_call` declared has its reply. +#[pyclass(frozen, name = "PostCallSettlement")] +pub(crate) struct PyPostCallSettlement(PostCallSettlement); + +#[pymethods] +impl PyPostCallSettlement { + #[getter] + fn settled(&self) -> bool { + self.0.is_settled() + } +} + fn increments(increments: Vec<(PyData, PyNumber)>) -> Vec<(SpendCounter, Number)> { increments .into_iter() @@ -73,6 +85,12 @@ impl NativeSpendCounters { self.0.redis_ttl().map(|ttl| ttl.as_secs_f64()) } + /// The least TTL a Redis write gives a committed spend twin. + fn set_committed_ttl(&self, ttl_seconds: f64) -> PyResult<()> { + self.0.set_committed_ttl(duration(ttl_seconds)?); + Ok(()) + } + /// `("redis" | "memory", value)`: where the answer came from, a Redis miss included. fn read<'py>( &self, @@ -125,14 +143,15 @@ impl NativeSpendCounters { }) } - /// Declares the increments on the request's post-call batch; `False` when there is no Redis. + /// Declares the increments on the request's post-call batch; `None` when there is no Redis. fn increment_many_post_call( &self, request: &NativeRequestRedis, increments: Vec<(PyData, PyNumber)>, - ) -> bool { + ) -> Option { self.0 .increment_many_post_call(&request.0, &self::increments(increments)) + .map(PyPostCallSettlement) } fn set_max<'py>( diff --git a/litellm-rust/crates/redis-store/AGENTS.md b/litellm-rust/crates/redis-store/AGENTS.md index cdd2bd7ca64..963f7ff168f 100644 --- a/litellm-rust/crates/redis-store/AGENTS.md +++ b/litellm-rust/crates/redis-store/AGENTS.md @@ -1,5 +1,5 @@ - One folder per Redis family (`router_cooldowns/`, `router_usage/`, `spend_counters/`, `auth_cache/`, `rate_limits/`), each with `mod.rs` (the store and its ops), `keys.rs` (key types and spellings) and `types.rs` (the family's values, with the `FromValue` / `IntoValue` codecs for what it stores in Redis). It is the only place that family's keys, value codec and scripts are written, and it takes domain ids, never keys. Shared pieces live once at the crate root: `Error` and `ErrorKind` (`error.rs`), `Memory` and `Clock` (`memory.rs`), `Pending` (`request.rs`), `Remote` and its attach-once `RemoteSlot`, which knows nothing of memory (`remote.rs`), `DualCache` (`dual_cache/`) and the Python codecs (`codec.rs`, `value.rs`) -- An op declared on a `Send` returns `Pending`, a lazy future built when the op is declared. An op that only means something on Redis (an `MGET` the caller falls back from, a script) returns `None` without one, and an op that replaced a `DualCache` call answers from memory. A post-call method returns `false` when there was no Redis to declare on +- An op declared on a `Send` returns `Pending`, a lazy future built when the op is declared. An op that only means something on Redis (an `MGET` the caller falls back from, a script) returns `None` without one, and an op that replaced a `DualCache` call answers from memory. A post-call method returns `false` when there was no Redis to declare on, or `None` when it returns a settlement the caller waits on, like `SpendCounterStore::increment_many_post_call` - A store that sends a pipeline of its own uses `RedisBatch::pipelined`, never `Pipeline` and raw `Value`s - Python and Rust pods share these keys during a rolling deploy, so a module writes the bytes `main`'s Python writes and reads everything `main` reads. Python formats go through `litellm-python-compat` (`repr`, `literal_eval`, `json`), never a hand-rolled codec - A family holds one `DualCache` per `DualCache` its Python held on `main`, and each `DualCache` method is the Python method it names, so a family op that was a `DualCache` call is one call here. Where `main`'s Python reached past `DualCache` to `in_memory_cache` or `redis_cache`, the family does the same through `memory()` and `redis()`, which keeps that reach visible. Caches a family's Python attached to one `RedisCache` share one `Arc`. Only a family that batch-reads builds its cache with `DualCache::batch_reading` from a `TierPolicy` the host passes in, and the Python host fills it from what `main` reads from the environment or `litellm` globals (`litellm/rust_bridge/host/settings.py`, `router_tiers`). Callers pass a `Remote` (or none), never a cache. Redis-only families are plain functions over `&RedisClient` diff --git a/litellm-rust/crates/redis-store/src/spend_counters/keys.rs b/litellm-rust/crates/redis-store/src/spend_counters/keys.rs index 061e19ed1de..5221f8ab66c 100644 --- a/litellm-rust/crates/redis-store/src/spend_counters/keys.rs +++ b/litellm-rust/crates/redis-store/src/spend_counters/keys.rs @@ -1,6 +1,9 @@ /// One budget's spend counter. `parse` and `key` round-trip `main`'s `spend:*` keys exactly. #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub enum SpendCounter { + /// The charges-only twin of a key, team, user, org or end user lifetime counter, which budget + /// holds never touch: `spend:committed:` and the twin's key past its `spend:`. + Committed(Box), Key { token: String, window: Option, @@ -33,6 +36,9 @@ impl SpendCounter { None => format!("spend:{scope}:{id}"), }; match self { + Self::Committed(counter) => { + format!("spend:committed:{}", &counter.key()["spend:".len()..]) + } Self::Key { token, window } => windowed("key", token, window), Self::Team { team_id, window } => windowed("team", team_id, window), Self::TeamMember { user_id, team_id } => { @@ -51,6 +57,16 @@ impl SpendCounter { pub fn parse(key: &str) -> Result { let invalid = || NotASpendCounter(key.to_owned()); let rest = key.strip_prefix("spend:").ok_or_else(invalid)?; + if let Some(twin) = rest.strip_prefix("committed:") { + return match Self::parse(&format!("spend:{twin}")).map_err(|_| invalid())? { + counter @ (Self::Key { window: None, .. } + | Self::Team { window: None, .. } + | Self::User(_) + | Self::Org(_) + | Self::EndUser(_)) => Ok(Self::Committed(Box::new(counter))), + _ => Err(invalid()), + }; + } let (scope, id) = rest.split_once(':').ok_or_else(invalid)?; let windowed = || match id.rsplit_once(":window:") { Some((id, window)) if !window.is_empty() => (id.to_owned(), Some(window.to_owned())), diff --git a/litellm-rust/crates/redis-store/src/spend_counters/mod.rs b/litellm-rust/crates/redis-store/src/spend_counters/mod.rs index 98f619cc501..425f83ef85c 100644 --- a/litellm-rust/crates/redis-store/src/spend_counters/mod.rs +++ b/litellm-rust/crates/redis-store/src/spend_counters/mod.rs @@ -1,13 +1,19 @@ mod keys; mod types; -use std::{sync::Arc, time::Duration}; +use std::{ + sync::{ + Arc, + atomic::{AtomicBool, AtomicU64, Ordering}, + }, + time::Duration, +}; use litellm_python_compat::repr::to_str; use litellm_redis_client::{Declared, Expiry, RedisBatch, commands}; pub use self::keys::{NotASpendCounter, SpendCounter}; -pub use self::types::{LostSeed, Reading}; +pub use self::types::{LostSeed, PostCallSettlement, Reading}; use crate::{ Error, Number, Pending, RequestRedis, Send, dual_cache::{DualCache, MemorySettings}, @@ -31,6 +37,7 @@ const SET_MAX: &str = concat!( /// `DualCache` to its Redis and its memory. pub struct SpendCounters { cache: DualCache, + committed_ttl_millis: AtomicU64, } impl SpendCounters { @@ -45,6 +52,33 @@ impl SpendCounters { MemorySettings::default(), now, ), + committed_ttl_millis: AtomicU64::new(0), + } + } + + /// The least TTL a Redis write gives a [`SpendCounter::Committed`] twin, so it outlives the DB + /// write that holds its charges and its next seed never reads a total still missing them. + pub fn set_committed_ttl(&self, ttl: Duration) { + let millis = u64::try_from(ttl.as_millis()).unwrap_or(u64::MAX); + self.committed_ttl_millis.store(millis, Ordering::Relaxed); + } + + fn ttl(&self, counter: &SpendCounter, default_ttl: Duration) -> Duration { + match counter { + SpendCounter::Committed(_) => default_ttl.max(Duration::from_millis( + self.committed_ttl_millis.load(Ordering::Relaxed), + )), + _ => default_ttl, + } + } + + /// A committed twin's expiry moves out on every charge, on a cluster too, so it never expires + /// while a charge it holds is still on its way to the DB. + fn post_call_expiry(&self, counter: &SpendCounter, default_ttl: Duration) -> Expiry { + let ttl = self.ttl(counter, default_ttl); + match counter { + SpendCounter::Committed(_) => Expiry::Refresh(ttl), + _ => Expiry::KeepExisting(ttl), } } @@ -127,7 +161,7 @@ impl SpendCounters { let increment = commands::incr_by_float( &key, amount.into(), - Some(Expiry::Refresh(remote.default_ttl)), + Some(Expiry::Refresh(self.ttl(counter, remote.default_ttl))), ); match remote.client.send(increment).await { Ok(total) => { @@ -159,7 +193,7 @@ impl SpendCounters { pipeline.add(commands::incr_by_float( &counter.key(), (*amount).into(), - Some(Expiry::Refresh(remote.default_ttl)), + Some(Expiry::Refresh(self.ttl(counter, remote.default_ttl))), )) }); let totals = Declared::all(totals.collect::>()).await; @@ -175,15 +209,14 @@ impl SpendCounters { } /// `increment_many` on the request's post-call batch. Each total reaches memory once the batch - /// goes out, and a counter whose increment failed is dropped everywhere. `false` without Redis. + /// goes out, and a counter whose increment failed is dropped everywhere. `None` without Redis, + /// else a settlement that reports when every reply is back, whoever flushed the batch. pub fn increment_many_post_call( self: &Arc, request: &RequestRedis, increments: &[(SpendCounter, Number)], - ) -> bool { - let Some(remote) = self.cache.redis().get() else { - return false; - }; + ) -> Option { + let remote = self.cache.redis().get()?; let batch = Send::PostCall(request).batch(&remote.client); let declared = increments .iter() @@ -191,12 +224,14 @@ impl SpendCounters { let total = batch.add(commands::incr_by_float( &counter.key(), (*amount).into(), - Some(Expiry::KeepExisting(remote.default_ttl)), + Some(self.post_call_expiry(counter, remote.default_ttl)), )); (counter.clone(), total) }) .collect::>(); let store = Arc::clone(self); + let settled = Arc::new(AtomicBool::new(false)); + let settlement = PostCallSettlement(Arc::clone(&settled)); request.spawn(async move { for (counter, total) in declared { match total.settled().await { @@ -215,8 +250,9 @@ impl SpendCounters { } } } + settled.store(true, Ordering::Release); }); - true + Some(settlement) } /// `_repair_stale_spend_counter`: raise memory and Redis to `spend`, never lowering either. @@ -243,7 +279,7 @@ impl SpendCounters { }; let args = ( to_str(&spend.to_value()), - remote.default_ttl.as_secs().to_string(), + self.ttl(counter, remote.default_ttl).as_secs().to_string(), ); let stored: Option = remote .client @@ -276,8 +312,9 @@ impl SpendCounters { }) })?); }; + let ttl = self.ttl(counter, *default_ttl); let current = if client - .send(commands::set_nx(&key, spend.to_json(), *default_ttl)) + .send(commands::set_nx(&key, spend.to_json(), ttl)) .await? { spend @@ -290,7 +327,7 @@ impl SpendCounters { let increment = commands::incr_by_float( &key, spend.into(), - Some(Expiry::KeepExisting(*default_ttl)), + Some(Expiry::KeepExisting(ttl)), ); Number::Float(client.send(increment).await?) } @@ -313,7 +350,7 @@ impl SpendCounters { let Some(remote) = self.cache.redis().get() else { return Ok(()); }; - let ttl = ttl.unwrap_or(remote.default_ttl); + let ttl = ttl.unwrap_or_else(|| self.ttl(counter, remote.default_ttl)); remote .client .send(commands::set(&key, spend.to_json(), ttl)) @@ -334,7 +371,10 @@ impl SpendCounters { }; Ok(remote .client - .send(commands::expire(&counter.key(), remote.default_ttl)) + .send(commands::expire( + &counter.key(), + self.ttl(counter, remote.default_ttl), + )) .await?) } diff --git a/litellm-rust/crates/redis-store/src/spend_counters/types.rs b/litellm-rust/crates/redis-store/src/spend_counters/types.rs index 61c211ec7e4..0509046d8c8 100644 --- a/litellm-rust/crates/redis-store/src/spend_counters/types.rs +++ b/litellm-rust/crates/redis-store/src/spend_counters/types.rs @@ -1,3 +1,8 @@ +use std::sync::{ + Arc, + atomic::{AtomicBool, Ordering}, +}; + use litellm_python_compat::Value; use super::keys::SpendCounter; @@ -32,6 +37,17 @@ impl Reading { } } +/// One `increment_many_post_call`: settled once every increment it declared has its reply and +/// memory mirrors it, whether `flush_post` or the post-call deadline sent the batch. +#[derive(Clone, Debug)] +pub struct PostCallSettlement(pub(super) Arc); + +impl PostCallSettlement { + pub fn is_settled(&self) -> bool { + self.0.load(Ordering::Acquire) + } +} + /// Who wins when a seed's `SET NX` lost and the winner's value is gone by the time it is read. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum LostSeed { diff --git a/litellm-rust/crates/redis-store/tests/spend_counters.rs b/litellm-rust/crates/redis-store/tests/spend_counters.rs index 75b7af2c321..7c89fb5b41e 100644 --- a/litellm-rust/crates/redis-store/tests/spend_counters.rs +++ b/litellm-rust/crates/redis-store/tests/spend_counters.rs @@ -16,6 +16,7 @@ use support::{COOLDOWN_POLICY, connect, connect_observed, remote, within}; use tokio::runtime::Handle; const TTL: Duration = Duration::from_secs(60); +const COMMITTED_TTL: Duration = Duration::from_secs(300); const LONG_DEADLINE: Duration = Duration::from_secs(60); const READS: [(&str, Number); 4] = [ @@ -45,6 +46,10 @@ fn team(id: &str) -> SpendCounter { } } +fn committed(counter: SpendCounter) -> SpendCounter { + SpendCounter::Committed(Box::new(counter)) +} + fn store(client: Option) -> Arc { Arc::new(SpendCounters::new(client.map(remote))) } @@ -114,6 +119,26 @@ fn a_reading_names_where_it_was_answered( "spend:model_access_group:gpt-group" )] #[case(SpendCounter::ShadowEval("job-7".to_owned()), "spend:shadow_eval:job-7")] +#[case::committed_key( + committed(SpendCounter::Key { + token: "88dc28d0".to_owned(), + window: None, + }), + "spend:committed:key:88dc28d0" +)] +#[case::committed_team(committed(team("team-a")), "spend:committed:team:team-a")] +#[case::committed_user( + committed(SpendCounter::User("user@example.com".to_owned())), + "spend:committed:user:user@example.com" +)] +#[case::committed_org( + committed(SpendCounter::Org("org-1".to_owned())), + "spend:committed:org:org-1" +)] +#[case::committed_end_user( + committed(SpendCounter::EndUser("customer é 1".to_owned())), + "spend:committed:end_user:customer é 1" +)] fn every_counter_formats_and_parses_back_to_mains_key( #[case] counter: SpendCounter, #[case] key: &str, @@ -128,6 +153,11 @@ fn every_counter_formats_and_parses_back_to_mains_key( #[case("spend:unknown:abc")] #[case("spend:team_member:no-team")] #[case("spend:")] +#[case::committed_window("spend:committed:key:abc:window:1d")] +#[case::committed_tag("spend:committed:tag:prod")] +#[case::committed_team_member("spend:committed:team_member:user:team-a")] +#[case::committed_twice("spend:committed:committed:key:abc")] +#[case::committed_nothing("spend:committed:")] fn keys_outside_the_spend_counter_keyspace_are_refused(#[case] key: &str) { assert!(SpendCounter::parse(key).is_err(), "{key}"); } @@ -307,7 +337,7 @@ async fn a_post_call_increment_from_a_thread_without_a_runtime_lands_in_redis_an .expect("a Python thread has no Tokio context, and declaring there must not panic"); request.flush_post().await; - assert!(declared); + assert!(declared.is_some()); let mirrored = litellm_redis_testing::eventually(async || store.get_local(&counter).unwrap()).await; assert_eq!(mirrored, Number::Float(3.5)); @@ -331,9 +361,12 @@ async fn post_call_increments_wait_for_the_post_batch_then_land_in_memory( let request = RequestRedis::new(Handle::current(), LONG_DEADLINE); let counter = team("t"); - assert!(store.increment_many_post_call(&request, &[(counter.clone(), Number::Float(3.5))])); + let settlement = store + .increment_many_post_call(&request, &[(counter.clone(), Number::Float(3.5))]) + .unwrap(); request.flush_pre().await; + assert!(!settlement.is_settled()); assert!(trips.seen().is_empty()); assert_eq!( raw.send(commands::get::(&counter.key())) @@ -345,6 +378,7 @@ async fn post_call_increments_wait_for_the_post_batch_then_land_in_memory( request.flush_post().await; tokio::time::sleep(Duration::from_millis(100)).await; + assert!(settlement.is_settled()); let increment = match topology { Topology::Standalone => Operation::Pipeline { ops: 1 }, Topology::Cluster => Operation::Command("INCRBYFLOAT"), @@ -557,13 +591,132 @@ async fn a_failed_post_call_increment_drops_the_counter_from_memory() { let request = RequestRedis::new(Handle::current(), LONG_DEADLINE); let counter = team("t"); store.increment(&counter, Number::Float(1.0)).await.unwrap(); - assert!(store.increment_many_post_call(&request, &[(counter.clone(), Number::Float(1.0))])); + let settlement = store + .increment_many_post_call(&request, &[(counter.clone(), Number::Float(1.0))]) + .unwrap(); deployment.stop().await.unwrap(); request.flush_post().await; tokio::time::sleep(Duration::from_millis(200)).await; assert_eq!(store.get_local(&counter).unwrap(), None); + assert!(settlement.is_settled()); +} + +#[rstest] +#[tokio::test] +async fn a_post_call_settlement_settles_when_the_deadline_sends_the_batch( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let store = store(Some(connect(&deployment))); + let request = RequestRedis::new(Handle::current(), Duration::from_millis(50)); + let counter = committed(team("t")); + + let settlement = store + .increment_many_post_call(&request, &[(counter.clone(), Number::Float(2.5))]) + .unwrap(); + + assert!(!settlement.is_settled()); + litellm_redis_testing::eventually(async || settlement.is_settled().then_some(())).await; + assert_eq!( + raw.send(commands::get::(&counter.key())) + .await + .unwrap(), + Some("2.5".to_owned()) + ); +} + +#[rstest] +#[case::committed_twin(committed, 300)] +#[case::enforcement_counter(std::convert::identity, 60)] +#[tokio::test] +async fn every_write_gives_a_committed_twin_the_committed_ttl_and_other_counters_the_redis_ttl( + #[case] counter_of: fn(SpendCounter) -> SpendCounter, + #[case] ttl_seconds: u64, + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let store = store(Some(connect(&deployment))); + store.set_committed_ttl(COMMITTED_TTL); + let request = RequestRedis::new(Handle::current(), LONG_DEADLINE); + let [ + incremented, + pipelined, + post_call, + seeded, + repaired, + refreshed, + ] = ["a", "b", "c", "d", "e", "f"].map(|id| counter_of(team(id))); + raw.send(commands::set( + &refreshed.key(), + "1.0", + Duration::from_secs(5), + )) + .await + .unwrap(); + + store + .increment(&incremented, Number::Float(1.0)) + .await + .unwrap(); + store + .increment_many(&[(pipelined.clone(), Number::Float(1.0))]) + .await + .unwrap(); + let settlement = store + .increment_many_post_call(&request, &[(post_call.clone(), Number::Float(1.0))]) + .unwrap(); + request.flush_post().await; + store + .seed(&seeded, Number::Float(1.0), LostSeed::Answer) + .await + .unwrap(); + store.set_max(&repaired, Number::Float(1.0)).await.unwrap(); + assert!(store.refresh_ttl(&refreshed).await.unwrap()); + + assert!(settlement.is_settled()); + let remaining = + async |counter: &SpendCounter| match raw.send(commands::ttl(&counter.key())).await.unwrap() + { + Ttl::Expires(remaining) => remaining, + other => panic!("{} has no TTL: {other:?}", counter.key()), + }; + assert!(within(remaining(&incremented).await, ttl_seconds)); + assert!(within(remaining(&pipelined).await, ttl_seconds)); + assert!(within(remaining(&post_call).await, ttl_seconds)); + assert!(within(remaining(&seeded).await, ttl_seconds)); + assert!(within(remaining(&repaired).await, ttl_seconds)); + assert!(within(remaining(&refreshed).await, ttl_seconds)); +} + +#[rstest] +#[tokio::test] +async fn a_post_call_charge_moves_a_committed_twins_expiry_out_to_the_committed_ttl( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let store = store(Some(connect(&deployment))); + store.set_committed_ttl(COMMITTED_TTL); + let request = RequestRedis::new(Handle::current(), LONG_DEADLINE); + let counter = committed(team("t")); + raw.send(commands::set(&counter.key(), "1.0", Duration::from_secs(5))) + .await + .unwrap(); + + let settlement = store + .increment_many_post_call(&request, &[(counter.clone(), Number::Float(1.0))]) + .unwrap(); + request.flush_post().await; + + assert!(settlement.is_settled()); + match raw.send(commands::ttl(&counter.key())).await.unwrap() { + Ttl::Expires(remaining) => assert!(within(remaining, COMMITTED_TTL.as_secs())), + other => panic!("{} has no TTL: {other:?}", counter.key()), + } } #[tokio::test] @@ -592,7 +745,11 @@ async fn without_redis_counters_live_in_memory() { store.read(&counter).await.unwrap(), Reading::Memory(Some(Number::Float(7.0))) ); - assert!(!store.increment_many_post_call(&request, &[(counter.clone(), Number::Int(1))])); + assert!( + store + .increment_many_post_call(&request, &[(counter.clone(), Number::Int(1))]) + .is_none() + ); assert!( store .read_many(RedisSend::PreCall(&request), &[counter]) diff --git a/litellm/accounting/metrics.py b/litellm/accounting/metrics.py index d4650813549..ee4b8b1c2b5 100644 --- a/litellm/accounting/metrics.py +++ b/litellm/accounting/metrics.py @@ -1,7 +1,9 @@ +import time +from collections.abc import Callable from dataclasses import dataclass from typing import TYPE_CHECKING, Final -from litellm.repositories.accounting_delivery_repository import DeliveryMetrics +from litellm.repositories.accounting_delivery_repository import DeliveryBacklog, DeliveryStorage if TYPE_CHECKING: from prometheus_client import Gauge, Histogram @@ -44,12 +46,32 @@ def _instruments() -> _Instruments | None: _INSTRUMENTS: Final = _instruments() -def record_deliveries(metrics: DeliveryMetrics, seconds: float) -> None: +def record_backlog(backlog: DeliveryBacklog, seconds: float) -> None: if _INSTRUMENTS is None: return - _INSTRUMENTS.pending.set(metrics.pending_count) - _INSTRUMENTS.parked.set(metrics.parked_count) - _INSTRUMENTS.oldest.set(max(0.0, metrics.oldest_pending_seconds)) + _INSTRUMENTS.pending.set(backlog.pending_count) + _INSTRUMENTS.parked.set(backlog.parked_count) + _INSTRUMENTS.oldest.set(max(0.0, backlog.oldest_pending_seconds)) _INSTRUMENTS.batch_time.observe(seconds) - _INSTRUMENTS.batch_count.set(metrics.batch_count) - _INSTRUMENTS.table_bytes.set(metrics.table_bytes) + + +def record_storage(storage: DeliveryStorage) -> None: + if _INSTRUMENTS is None: + return + _INSTRUMENTS.batch_count.set(storage.batch_count) + _INSTRUMENTS.table_bytes.set(storage.table_bytes) + + +class Cadence: + __slots__ = ("_clock", "_last") + + def __init__(self, clock: Callable[[], float] = time.monotonic) -> None: + self._clock = clock + self._last: float | None = None + + def due(self, interval_seconds: float) -> bool: + now: Final = self._clock() + if self._last is not None and now - self._last < interval_seconds: + return False + self._last = now + return True diff --git a/litellm/accounting/worker.py b/litellm/accounting/worker.py index be18d8ac6b1..0f7e80123c2 100644 --- a/litellm/accounting/worker.py +++ b/litellm/accounting/worker.py @@ -10,7 +10,7 @@ from typing import ( from litellm._logging import verbose_proxy_logger from litellm.accounting.application import apply_deliveries -from litellm.accounting.metrics import record_deliveries +from litellm.accounting.metrics import record_backlog, record_storage from litellm.accounting.projection import Contribution from litellm.repositories.accounting_delivery_repository import ( AccountingDeliveryRepository, @@ -96,24 +96,37 @@ async def flush_accounting_deliveries( or time.monotonic() >= cleanup_deadline ): break - metrics: Final = await repository.metrics() - record_deliveries(metrics, time.monotonic() - started) - if metrics.batch_count >= int( - os.environ.get("LITELLM_ACCOUNTING_BATCH_DELIVERY_ALERT_COUNT", "1000000") - ) or metrics.table_bytes >= int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_ALERT_BYTES", "1073741824")): - await _alert( - proxy_logging_obj, - message=f"Accounting delivery storage requires review. Retained batches={metrics.batch_count}, table bytes={metrics.table_bytes}.", - ) - if metrics.oldest_pending_seconds > 300: - await _alert( - proxy_logging_obj, - message=f"Accounting deliveries are delayed. Pending={metrics.pending_count}, oldest={metrics.oldest_pending_seconds:.0f}s.", - ) + await _report(writer, repository, proxy_logging_obj, time.monotonic() - started) finally: if acquired: await writer.pod_lock_manager.release_lock(cronjob_id=_JOB) +async def _report( + writer: "DBSpendUpdateWriter", + repository: AccountingDeliveryRepository, + proxy_logging_obj: _AccountingAlerts, + seconds: float, +) -> None: + backlog: Final = await repository.backlog() + record_backlog(backlog, seconds) + interval: Final = float(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_CHECK_INTERVAL_SECONDS", "900")) + if writer.accounting_storage_checks.due(interval): + storage: Final = await repository.storage() + record_storage(storage) + if storage.batch_count >= int( + os.environ.get("LITELLM_ACCOUNTING_BATCH_DELIVERY_ALERT_COUNT", "1000000") + ) or storage.table_bytes >= int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_ALERT_BYTES", "1073741824")): + await _alert( + proxy_logging_obj, + message=f"Accounting delivery storage requires review. Retained batches={storage.batch_count}, table bytes={storage.table_bytes}.", + ) + if backlog.oldest_pending_seconds > 300 and writer.accounting_delay_alerts.due(interval): + await _alert( + proxy_logging_obj, + message=f"Accounting deliveries are delayed. Pending={backlog.pending_count}, oldest={backlog.oldest_pending_seconds:.0f}s.", + ) + + async def _alert(alerts: _AccountingAlerts, message: str) -> None: await alerts.alerting_handler(message=message, level="High", alert_type=AlertType.failed_tracking_spend) diff --git a/litellm/caching/redis_batch.py b/litellm/caching/redis_batch.py index d58ddb9c629..b3e6d2eb291 100644 --- a/litellm/caching/redis_batch.py +++ b/litellm/caching/redis_batch.py @@ -452,6 +452,13 @@ class NativePending(Protocol[_T_co]): def resolve(self) -> Awaitable[_T_co]: ... +class PostCallSettlement(Protocol): + """Post-call writes declared together: settled once every reply is back, whoever sent their batch.""" + + @property + def settled(self) -> bool: ... + + class NativeBatchResult(Generic[_T]): """Awaitable handle for a read declared on the request's native pre-call batch; awaiting it runs the native flush hooks, flushes that batch, then answers from what it read.""" @@ -526,6 +533,7 @@ class RequestRedisBatches: "_native_misses", "_native_send_factory", "_post_call", + "_post_call_settlements", "post_call_deadline", "prefetched", ) @@ -538,6 +546,7 @@ class RequestRedisBatches: ) -> None: self._batches: Final[dict[object, RedisBatch]] = {} # mutable-ok: lazily filled per backend self._post_call: Final[dict[object, RedisBatch]] = {} # mutable-ok: lazily filled per backend + self._post_call_settlements: tuple[PostCallSettlement, ...] = () self.post_call_deadline: Final = post_call_deadline self._native_factory: Final = native_factory self._native_send_factory: Final = native_send_factory @@ -641,6 +650,14 @@ class RequestRedisBatches: if not any(batch.pending for batch in self._post_call.values()): _open_post_call.discard(self) + def track_post_call_settlement(self, settlement: PostCallSettlement) -> None: + """Post-call writes a later reader in this request must know have landed, such as its committed spend.""" + self._post_call_settlements = (*self._post_call_settlements, settlement) + + @property + def post_call_settlements(self) -> tuple[PostCallSettlement, ...]: + return self._post_call_settlements + @property def batches(self) -> tuple[RedisBatch, ...]: return tuple(self._batches.values()) diff --git a/litellm/constants.py b/litellm/constants.py index b43f6d9f358..3a49e8a7fc4 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1787,6 +1787,7 @@ AZURE_SPEECH_PASS_THROUGH_ROUTE_PREFIX: Final = "/azure_speech" AZURE_SPEECH_SHORT_AUDIO_PATH_PREFIX: Final = "/speech/" AZURE_SPEECH_BATCH_PATH_PREFIX: Final = "/speechtotext/" AZURE_SPEECH_FAST_TRANSCRIPTION_PATH: Final = "/speechtotext/transcriptions:transcribe" +OPENSEARCH_PASS_THROUGH_ROUTE_PREFIX: Final = "/opensearch" AZURE_SPEECH_STT_DOMAIN: Final = "stt.speech.microsoft.com" AZURE_SPEECH_COGNITIVE_SERVICES_DOMAIN: Final = "api.cognitive.microsoft.com" AZURE_SPEECH_SUBSCRIPTION_KEY_HEADER: Final = "Ocp-Apim-Subscription-Key" @@ -1940,6 +1941,7 @@ SPEND_ROLLUP_LOCK_TIMEOUT_MS: Final = max(1, get_env_int("SPEND_ROLLUP_LOCK_TIME SPEND_LOG_QUEUE_SIZE_THRESHOLD: Final = get_env_int("SPEND_LOG_QUEUE_SIZE_THRESHOLD", 100) SPEND_LOG_QUEUE_MAX_BYTES: Final = max(1, get_env_int("SPEND_LOG_QUEUE_MAX_BYTES", 64000000)) SPEND_LOG_QUEUE_POLL_INTERVAL: Final = get_env_float("SPEND_LOG_QUEUE_POLL_INTERVAL", 2.0) +SPEND_LOG_QUEUE_MAX_BACKOFF_SECONDS: Final = 30.0 RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS: Final = max(1, get_env_int("RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS", 3)) RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL: Final = get_env_float("RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL", 0.2) SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE: Final = get_env_int("SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE", 10000) diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index 5cd201a2f47..9bec220e073 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -34,7 +34,7 @@ from litellm.integrations.otel.model.db_endpoint import db_span_attributes from litellm.integrations.otel.model.metadata import flatten_metadata from litellm.integrations.otel.model.semconv import LiteLLM, Metric from litellm.integrations.otel.plumbing.otlp_tls import resolve_otlp_http_tls -from litellm.integrations.otel.routing import routing_decision_attributes +from litellm.integrations.otel.routing import routed_model_name, routing_decision_attributes from litellm.litellm_core_utils.core_helpers import iter_items, require_attributes from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call_from_params from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -2617,7 +2617,11 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): ) # The model used to generate the response. - response_model: Final = _response_field(response_obj, "model") + response_model: Final = routed_model_name( + _response_field(response_obj, "model"), + standard_logging_payload.get("model_group"), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # kwargs is untyped + decision, + ) if response_model: self.safe_set_attribute( span=span, diff --git a/litellm/integrations/otel/model/metadata.py b/litellm/integrations/otel/model/metadata.py index 28d6ad54607..f9739a4ba15 100644 --- a/litellm/integrations/otel/model/metadata.py +++ b/litellm/integrations/otel/model/metadata.py @@ -29,9 +29,10 @@ asks for a *model group* (e.g. ``gpt-4o``) that routes to a concrete deployment * ``hidden_params.litellm_model_name`` — a secondary source for the dispatched model (populated only on some call paths, e.g. files). -So ``gen_ai.request.model`` is the *group* (falling back to the call model on the -SDK path, which has no group), and ``litellm.provider.model`` is the *dispatched* -model. They coincide on the SDK path, which is correct. +So ``gen_ai.request.model`` is the *group* (the group an auto router picked, when +one routed the call; the call model on the SDK path, which has no group), and +``litellm.provider.model`` is the *dispatched* model. They coincide on the SDK +path, which is correct. """ from __future__ import annotations @@ -52,6 +53,7 @@ from litellm.constants import ( from litellm.integrations.otel.model.semconv import PASSTHROUGH_OPERATION_KEY, GenAIOperation, resolve_operation from litellm.integrations.otel.model.trace_controls import TraceControls, caller_trace_controls from litellm.integrations.otel.model.utils import as_str, as_str_mapping, to_seconds +from litellm.integrations.otel.routing import routed_model_name from litellm.types.utils import InternalCallOrigin if TYPE_CHECKING: @@ -170,7 +172,8 @@ class RequestIdentity: class RequestContext: """The fully-resolved view of a closed request, parsed once from the payload. - ``request_model`` is the user-facing requested model and ``provider_model`` + ``request_model`` is the user-facing requested model (or the group an auto + router picked for it) and ``provider_model`` (on :attr:`identity`) is the model litellm dispatched to the provider; the two differ on the proxy (group vs. deployment) and coincide on the SDK path. """ @@ -194,12 +197,15 @@ class RequestContext: raw_response: Final = payload.get("response") response: Final = cast(Mapping[str, object], raw_response if isinstance(raw_response, dict) else {}) model_group: Final = as_str(payload.get("model_group")) or as_str(raw_meta.get("model_group")) + decision: Final = raw_meta.get("routing_decision") return cls( - # The user asked for the group; fall back to the call model on the SDK - # path, which has no group. Empty string (never None) so the span name - # builder and the mapper see a plain string. - request_model=model_group or as_str(payload.get("model")) or "", - response_model=as_str(response.get("model")), + # The user asked for the group, or an auto router picked one for it (the + # proxy restamps the response model to the requested name, so that gets + # the pick too); fall back to the call model on the SDK path, which has no + # group. Empty string (never None) so the span name builder and the mapper + # see a plain string. + request_model=routed_model_name(model_group, model_group, decision) or as_str(payload.get("model")) or "", + response_model=routed_model_name(as_str(response.get("model")), model_group, decision), model_group=model_group, model_id=as_str(payload.get("model_id")) or _model_info_id(raw_meta.get("model_info")), api_base=as_str(payload.get("api_base")) or as_str(hidden.get("api_base")), diff --git a/litellm/integrations/otel/routing.py b/litellm/integrations/otel/routing.py index 339775855be..1143e6d2dbe 100644 --- a/litellm/integrations/otel/routing.py +++ b/litellm/integrations/otel/routing.py @@ -1,8 +1,11 @@ from collections.abc import Mapping from types import MappingProxyType -from typing import Final, TypeAlias +from typing import Final, TypeAlias, TypeVar + +from litellm.litellm_core_utils.core_helpers import as_str_mapping RoutingAttributeValue: TypeAlias = str | int | float | bool +_Model = TypeVar("_Model") # rebind-ok: a TypeVar cannot be Final _ROUTING_SCALAR_FIELDS: Final = frozenset( { @@ -60,3 +63,12 @@ def routing_decision_attributes(decision: Mapping[str, object] | None) -> Mappin if isinstance(value := decision.get(key), (str, int, float, bool)) } ) + + +def routed_model_name(model: _Model, requested_model: object, decision: object) -> _Model | str: + """The model an auto router picked when ``model`` is only the name the caller requested, else ``model``.""" + fields: Final = as_str_mapping(decision) + routed: Final = fields.get("routed_model") if fields is not None else None + if model and model == requested_model and isinstance(routed, str) and routed: + return routed + return model diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index fc1a601f9da..7984c7d04f3 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -21,6 +21,7 @@ from typing_extensions import ReadOnly, TypedDict, assert_never import litellm from litellm._internal_context import with_service_target from litellm._logging import print_verbose, verbose_logger +from litellm.caching.redis_batch import active_request_redis_batches from litellm.constants import ( AUTH_OBJECTS_TARGET, LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL, @@ -75,8 +76,10 @@ from litellm.types.integrations.prometheus import ( sanitize_prometheus_label_value, validate_prometheus_deployment_and_latency_caller_identity, ) -from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth, committed_budget_spend_of_row from litellm.types.proxy.carried_budget_state import ( + CommittedBudgetSpend, + EntitySpendCounterKeys, KeyBudgetSnapshot, OrgBudgetSnapshot, TeamBudgetSnapshot, @@ -281,6 +284,11 @@ def _logged_fields(value: object) -> Mapping[str, object]: return _NO_LOGGED_FIELDS +def _response_cost_of(logged: Mapping[str, object]) -> float: + cost: Final = logged.get("response_cost") + return float(cost) if isinstance(cost, (int, float)) else 0.0 + + _GOOGLE_PASS_THROUGH_PROVIDERS: Final = frozenset( (litellm.LlmProviders.GEMINI.value, litellm.LlmProviders.VERTEX_AI.value) ) @@ -365,6 +373,104 @@ def _get_budget_metrics_per_request_timeout() -> float: return parsed +def _stored_spend(value: object) -> float: + return float(value) if isinstance(value, (int, float)) else 0.0 + + +_OWN_CHARGE_MATCH_RATIO: Final = 1 - 1e-9 + + +def _committed_floor(stored: float, committed: float | None) -> float: + return stored if committed is None else max(stored, committed) + + +def _customer_spend_after_request(cached_spend: float, committed_spend: float | None, response_cost: float) -> float: + if committed_spend is None: + return cached_spend + response_cost + return max(committed_spend + response_cost, cached_spend) + + +def _budgeted_org_counter(org_id: str | None, org: OrgBudgetSnapshot | None) -> EntitySpendCounterKeys: + if org is None or org.max_budget is None: + return EntitySpendCounterKeys() + return EntitySpendCounterKeys.of(token=None, team_id=None, user_id=None, org_id=org_id, end_user_id=None) + + +async def _committed_spend_now(counter_keys: EntitySpendCounterKeys, timeout: float) -> CommittedBudgetSpend: + if not counter_keys.tracked(): + return CommittedBudgetSpend() + try: + from litellm.proxy.proxy_server import spend_counter_store + from litellm.proxy.spend_tracking.committed_spend import read_committed_budget_spend + + return await asyncio.wait_for( + read_committed_budget_spend(spend_counter_store, counter_keys), + timeout=timeout, + ) + except Exception as e: # noqa: BLE001 # a gauge read never fails the request; the gauges keep the stored spend + verbose_logger.debug("[Non-Blocking] Prometheus: committed spend read failed: %s", e) + return CommittedBudgetSpend() + + +OwnCharge: TypeAlias = Literal["unknown", "pending", "counted"] + + +def _own_charge() -> OwnCharge: + """Whether this request's committed spend charge is in Redis yet, known only inside a request with Redis.""" + try: + from litellm.proxy.proxy_server import spend_counter_store + except ImportError: + return "unknown" + batches: Final = active_request_redis_batches() + if batches is None or not spend_counter_store.has_redis: + return "unknown" + settlements: Final = batches.post_call_settlements + return "counted" if settlements and all(settlement.settled for settlement in settlements) else "pending" + + +def _committed_spend_before_request( + at_auth: CommittedBudgetSpend, now: CommittedBudgetSpend, response_cost: float, own_charge: OwnCharge +) -> CommittedBudgetSpend: + def before(at_auth_spend: float | None, now_spend: float | None) -> float | None: + if now_spend is None: + return at_auth_spend + match own_charge: + case "pending": + return now_spend + case "counted": + without_own_charge: Final = now_spend - response_cost + return without_own_charge if at_auth_spend is None else max(at_auth_spend, without_own_charge) + case "unknown": + if at_auth_spend is None: + return now_spend - response_cost + own_charge_counted: Final = now_spend - at_auth_spend >= response_cost * _OWN_CHARGE_MATCH_RATIO + return max(at_auth_spend, now_spend - response_cost) if own_charge_counted else now_spend + case _: + assert_never(own_charge) + + return CommittedBudgetSpend( + key=before(at_auth.key, now.key), + team=before(at_auth.team, now.team), + user=before(at_auth.user, now.user), + org=before(at_auth.org, now.org), + end_user=before(at_auth.end_user, now.end_user), + ) + + +async def _committed_spend_without_own_charge( + counter_keys: EntitySpendCounterKeys, + at_auth: CommittedBudgetSpend | None, + response_cost: float, + read_timeout: float, +) -> CommittedBudgetSpend: + own_charge_before_read: Final = _own_charge() + now: Final = await _committed_spend_now(counter_keys, read_timeout) + own_charge: Final = own_charge_before_read if _own_charge() == own_charge_before_read else "unknown" + return _committed_spend_before_request( + at_auth=now if at_auth is None else at_auth, now=now, response_cost=response_cost, own_charge=own_charge + ) + + def _get_proxy_llm_router() -> Router | None: try: from litellm.proxy.proxy_server import llm_router @@ -2229,6 +2335,10 @@ class PrometheusLogger(CustomLogger): _user_max_budget: Final = _metadata.get("user_api_key_user_max_budget", None) _user_email: Final = _metadata.get("user_api_key_user_email", None) _org_alias: Final = _metadata.get("user_api_key_org_alias", None) + _carried_metadata: Final = _REQUEST_METADATA_ADAPTER.validate_python(_metadata) + _committed_at_auth: Final = committed_budget_spend_of_row( + _carried_metadata, _REQUEST_METADATA_ADAPTER.validate_python(litellm_params) + ) # Bound the per-request budget-metric emission so that slow Redis/DB # lookups under load cannot consume the whole LoggingWorker watchdog @@ -2238,45 +2348,65 @@ class PrometheusLogger(CustomLogger): # so dropping one slow per-request emission only loses sub-cron real-time # detail, not correctness. budget_metrics_timeout: Final = _get_budget_metrics_per_request_timeout() - gather_coro: Final = asyncio.gather( - self._set_api_key_budget_metrics_after_api_request( - user_api_key=user_api_key, - user_api_key_alias=user_api_key_alias, + committed_read_timeout: Final = budget_metrics_timeout / 2 + + async def emit_budget_metrics() -> tuple[object, ...]: + committed: Final = await _committed_spend_without_own_charge( + counter_keys=_committed_at_auth.counter_keys + if _committed_at_auth is not None + else EntitySpendCounterKeys.of( + token=user_api_key, + team_id=user_api_team, + user_id=user_id, + org_id=user_api_key_org_id, + end_user_id=end_user_id, + ), + at_auth=_committed_at_auth, response_cost=response_cost, - key_max_budget=_api_key_max_budget, - key_spend=_api_key_spend, - carried=KeyBudgetSnapshot.from_metadata(_metadata), - ), - self._set_team_budget_metrics_after_api_request( - user_api_team=user_api_team, - user_api_team_alias=user_api_team_alias, - team_spend=_team_spend, - team_max_budget=_team_max_budget, - response_cost=response_cost, - carried=TeamBudgetSnapshot.from_metadata(_metadata), - ), - self._set_user_budget_metrics_after_api_request( - user_id=user_id, - user_spend=_user_spend, - user_max_budget=_user_max_budget, - response_cost=response_cost, - carried=UserBudgetSnapshot.from_metadata(_metadata), - user_email=_user_email if isinstance(_user_email, str) else None, - ), - self._set_org_budget_metrics_after_api_request( - org_id=user_api_key_org_id, - response_cost=response_cost, - carried=OrgBudgetSnapshot.from_metadata(_metadata), - org_alias=_org_alias if isinstance(_org_alias, str) else None, - ), - self._set_customer_budget_metrics_after_api_request( - end_user_id=end_user_id, - response_cost=response_cost, - ), - return_exceptions=True, - ) + read_timeout=committed_read_timeout, + ) + return await asyncio.gather( + self._set_api_key_budget_metrics_after_api_request( + user_api_key=user_api_key, + user_api_key_alias=user_api_key_alias, + response_cost=response_cost, + key_max_budget=_api_key_max_budget, + key_spend=_committed_floor(_stored_spend(_api_key_spend), committed.key), + carried=KeyBudgetSnapshot.from_metadata(_carried_metadata), + ), + self._set_team_budget_metrics_after_api_request( + user_api_team=user_api_team, + user_api_team_alias=user_api_team_alias, + team_spend=_committed_floor(_stored_spend(_team_spend), committed.team), + team_max_budget=_team_max_budget, + response_cost=response_cost, + carried=TeamBudgetSnapshot.from_metadata(_carried_metadata), + ), + self._set_user_budget_metrics_after_api_request( + user_id=user_id, + user_spend=_committed_floor(_stored_spend(_user_spend), committed.user), + user_max_budget=_user_max_budget, + response_cost=response_cost, + carried=UserBudgetSnapshot.from_metadata(_carried_metadata), + user_email=_user_email if isinstance(_user_email, str) else None, + ), + self._set_org_budget_metrics_after_api_request( + org_id=user_api_key_org_id, + response_cost=response_cost, + carried=OrgBudgetSnapshot.from_metadata(_carried_metadata), + org_alias=_org_alias if isinstance(_org_alias, str) else None, + committed_spend=committed.org, + ), + self._set_customer_budget_metrics_after_api_request( + end_user_id=end_user_id, + response_cost=response_cost, + committed_spend=committed.end_user, + ), + return_exceptions=True, + ) + try: - results: Final = await asyncio.wait_for(gather_coro, timeout=budget_metrics_timeout) + results: Final = await asyncio.wait_for(emit_budget_metrics(), timeout=budget_metrics_timeout) except asyncio.TimeoutError: verbose_logger.debug( "[Non-Blocking] Prometheus: per-request budget metric emission " @@ -2745,7 +2875,7 @@ class PrometheusLogger(CustomLogger): _labels, ) - async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: verbose_logger.debug( "prometheus Logging - Enters failure logging function (kwargs keys: %s)", list(kwargs.keys()) if isinstance(kwargs, dict) else type(kwargs).__name__, @@ -2799,9 +2929,27 @@ class PrometheusLogger(CustomLogger): label_context=PrometheusLabelFactoryContext(enum_values), ) self.set_llm_deployment_failure_metrics(kwargs) + failure_metadata: Final = _REQUEST_METADATA_ADAPTER.validate_python( + get_litellm_metadata_from_kwargs(kwargs) + ) + failure_org_alias: Final = failure_metadata.get("user_api_key_org_alias") + failure_org: Final = OrgBudgetSnapshot.from_metadata(failure_metadata) + failure_cost: Final = _response_cost_of(_LOGGED_FIELDS.validate_python(standard_logging_payload)) + failure_committed: Final = await _committed_spend_without_own_charge( + counter_keys=_budgeted_org_counter(user_api_key_org_id, failure_org), + at_auth=committed_budget_spend_of_row( + failure_metadata, _REQUEST_METADATA_ADAPTER.validate_python(litellm_params) + ) + or CommittedBudgetSpend(), + response_cost=failure_cost, + read_timeout=_get_budget_metrics_per_request_timeout(), + ) await self._set_org_budget_metrics_after_api_request( org_id=user_api_key_org_id, - response_cost=0, + response_cost=failure_cost, + carried=failure_org, + org_alias=failure_org_alias if isinstance(failure_org_alias, str) else None, + committed_spend=failure_committed.org, ) except Exception as e: verbose_logger.exception("prometheus Layer Error(): Exception occured - %s", e) @@ -4516,6 +4664,7 @@ class PrometheusLogger(CustomLogger): response_cost: float, carried: OrgBudgetSnapshot | None = None, org_alias: str | None = None, + committed_spend: float | None = None, ): """ Set org budget metrics after an LLM API request @@ -4533,7 +4682,7 @@ class PrometheusLogger(CustomLogger): self._set_org_budget_metrics( org_id=org_id, org_alias=org_alias or "", - spend=carried.spend + response_cost, + spend=_committed_floor(carried.spend, committed_spend) + response_cost, max_budget=carried.max_budget, budget_reset_at=None, ) @@ -4559,7 +4708,7 @@ class PrometheusLogger(CustomLogger): if org_info is None: return - _total_org_spend: Final = (org_info.spend or 0.0) + response_cost + _total_org_spend: Final = _committed_floor(org_info.spend or 0.0, committed_spend) + response_cost budget_table: Final = org_info.litellm_budget_table max_budget: Final = budget_table.max_budget if budget_table else None budget_reset_at: Final = getattr(budget_table, "budget_reset_at", None) if budget_table else None @@ -4626,6 +4775,7 @@ class PrometheusLogger(CustomLogger): self, end_user_id: str | None, response_cost: float, + committed_spend: float | None = None, ): if self._customer_budget_gauges_are_noop() or not _customer_budget_metrics_enabled(): return @@ -4651,7 +4801,9 @@ class PrometheusLogger(CustomLogger): budget_table: Final = cached_customer.litellm_budget_table self._set_customer_budget_metrics( end_user_id=end_user_id, - spend=cached_customer.spend + response_cost, + spend=_customer_spend_after_request( + cached_spend=cached_customer.spend, committed_spend=committed_spend, response_cost=response_cost + ), max_budget=budget_table.max_budget if budget_table is not None else None, budget_reset_at=None, ) diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index cfd07c9e2af..6c35b53c92d 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -17,7 +17,7 @@ from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass from datetime import datetime, timezone from functools import partial -from typing import TYPE_CHECKING, Final, Literal, TypeAlias, cast +from typing import TYPE_CHECKING, ClassVar, Final, Literal, TypeAlias, cast from urllib.parse import quote from uuid import uuid4 @@ -55,6 +55,7 @@ from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing +from litellm.llms.bedrock.web_identity_session_policy import S3_LOGGING_SESSION_POLICY, WebIdentitySessionPolicy from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, get_httpx_client, @@ -154,6 +155,7 @@ _in_flush: Final[contextvars.ContextVar[bool]] = contextvars.ContextVar("s3_v2_i class S3Logger(CustomBatchLogger, BaseAWSLLM): + _web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = S3_LOGGING_SESSION_POLICY preserve_events_added_during_flush = True _flush_retries: int = 0 _requeued_count: int = 0 diff --git a/litellm/integrations/sqs.py b/litellm/integrations/sqs.py index 83f85232182..ba7944dfb59 100644 --- a/litellm/integrations/sqs.py +++ b/litellm/integrations/sqs.py @@ -11,7 +11,7 @@ import base64 import json import re import traceback -from typing import Final +from typing import ClassVar, Final from pydantic import TypeAdapter @@ -25,6 +25,7 @@ from litellm.constants import ( ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing +from litellm.llms.bedrock.web_identity_session_policy import SQS_LOGGING_SESSION_POLICY, WebIdentitySessionPolicy from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -44,6 +45,8 @@ _BASE64_INLINE_PATTERN: Final = re.compile( class SQSLogger(CustomBatchLogger, BaseAWSLLM): """Batching logger that writes logs to an AWS SQS queue, optionally encrypting the payload.""" + _web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = SQS_LOGGING_SESSION_POLICY + def __init__( self, # --- Standard SQS params --- diff --git a/litellm/litellm_core_utils/allowed_tools_tool_choice.py b/litellm/litellm_core_utils/allowed_tools_tool_choice.py index 5c8f5d83ece..f4c4ecbbb26 100644 --- a/litellm/litellm_core_utils/allowed_tools_tool_choice.py +++ b/litellm/litellm_core_utils/allowed_tools_tool_choice.py @@ -8,6 +8,24 @@ from litellm.exceptions import BadRequestError ALLOWED_TOOLS_TOOL_CHOICE_PROVIDERS: Final = frozenset({"openai", "azure", "litellm_proxy"}) +HOSTED_TOOL_CHOICE_TYPES: Final = frozenset( + { + "file_search", + "tool_search", + "web_search", + "web_search_preview", + "web_search_preview_2025_03_11", + "computer", + "computer_use", + "computer_use_preview", + "code_interpreter", + "image_generation", + "mcp", + "shell", + "apply_patch", + } +) + CHAT_ALLOWED_TOOLS_LAYOUT: Final = ( '{"type": "allowed_tools", "allowed_tools": {"mode": "auto" | "required", "tools": []}}' ) @@ -70,6 +88,37 @@ def raise_if_allowed_tools_tool_choice_unsupported( ) +def hosted_tool_choice_type(tool_choice: object) -> str | None: + """The hosted tool a Responses API ``tool_choice`` forces, such as ``web_search``, or ``None`` for any other choice.""" + try: + choice_type: Final = _TOOL_CHOICE_TYPE.validate_python(tool_choice).type + except ValidationError: + return None + return choice_type if choice_type in HOSTED_TOOL_CHOICE_TYPES else None + + +def raise_if_hosted_tool_choice_unsupported( + tool_choice: object, + custom_llm_provider: str, + model: str, +) -> None: + """Anthropic drops a tool_choice dict it does not recognise and Bedrock sends a tool named "", so a forced hosted + tool stops here with a 400 on providers that cannot run it. ``drop_params`` is deliberately not consulted: the + request required the tool, and dropping the choice would let the model answer without it.""" + hosted: Final = hosted_tool_choice_type(tool_choice) + if hosted is None or custom_llm_provider in ALLOWED_TOOLS_TOOL_CHOICE_PROVIDERS: + return + raise BadRequestError( + message=( + f"tool_choice type '{hosted}' forces an OpenAI hosted tool and is only supported for " + f"{', '.join(sorted(ALLOWED_TOOLS_TOOL_CHOICE_PROVIDERS))}, got custom_llm_provider={custom_llm_provider}. " + "drop_params does not drop it, since that would let the model answer without the tool the request required." + ), + model=model, + llm_provider=custom_llm_provider, + ) + + def parse_chat_allowed_tools_tool_choice(tool_choice: object) -> ChatAllowedToolsToolChoice | None: try: return _CHAT_ALLOWED_TOOLS_TOOL_CHOICE.validate_python(tool_choice) diff --git a/litellm/litellm_core_utils/attempt_settlement.py b/litellm/litellm_core_utils/attempt_settlement.py new file mode 100644 index 00000000000..02a6497cff8 --- /dev/null +++ b/litellm/litellm_core_utils/attempt_settlement.py @@ -0,0 +1,90 @@ +""" +What a routed attempt holds (io-token reservations, least-busy in-flight counts) is released exactly +once, by whichever of the attempt itself or its logging callbacks gets there first, and the release +runs to completion even when the request is cancelled mid-release. +""" + +import asyncio +import contextlib +import sys +import uuid +from collections.abc import Coroutine, Generator, Sequence +from contextvars import ContextVar +from typing import Final + +from litellm._logging import verbose_router_logger + +ATTEMPT_SETTLEMENT_GRACE_SECONDS: Final = 5.0 + + +class RoutedAttempt: + __slots__ = ("answered", "attempt_id") + + def __init__(self) -> None: + self.attempt_id: Final = uuid.uuid4().hex + self.answered: bool = False + + +_ATTEMPT: Final[ContextVar[RoutedAttempt | None]] = ContextVar("routed_attempt", default=None) +_RELEASING: Final[set["asyncio.Task[None]"]] = set() # mutable-ok: a release leaves it when it finishes + + +def current_attempt_id() -> str | None: + attempt: Final = _ATTEMPT.get() + return None if attempt is None else attempt.attempt_id + + +def mark_attempt_answered() -> None: + """The provider answered the current routed attempt: what it reserved was spent, not just held.""" + attempt: Final = _ATTEMPT.get() + if attempt is not None: + attempt.answered = True + + +@contextlib.contextmanager +def routed_attempt() -> Generator[RoutedAttempt, None, None]: + attempt: Final = RoutedAttempt() + token: Final = _ATTEMPT.set(attempt) + try: + yield attempt + finally: + _ATTEMPT.reset(token) + + +def _finished(task: "asyncio.Task[None]") -> None: + _RELEASING.discard(task) + if task.cancelled(): + verbose_router_logger.warning("an attempt's release was cancelled before it finished") + return + error: Final = task.exception() + if error is not None: + verbose_router_logger.warning("an attempt's release failed: %s", error) + + +def _start(release: Coroutine[object, object, None]) -> "asyncio.Task[None]": + """Run the release up to its first real wait right away, so one that never waits (an in-memory + counter) is done before its caller goes on, as it was when the caller awaited it inline.""" + loop: Final = asyncio.get_running_loop() + if sys.version_info >= (3, 12): + return asyncio.Task(release, loop=loop, eager_start=True) + return loop.create_task(release) + + +def release_in_background(release: Coroutine[object, object, None]) -> "asyncio.Task[None]": + task: Final = _start(release) + _RELEASING.add(task) + task.add_done_callback(_finished) + return task + + +async def released(releases: Sequence["asyncio.Task[None]"], grace: float | None = None) -> None: + """Wait for ``releases`` without ever cancelling them: a cancelled waiter leaves them running.""" + if all(release.done() for release in releases): + return + _, pending = await asyncio.wait(releases, timeout=grace) + if pending: + verbose_router_logger.warning( + "%s attempt release(s) still running after %ss; they keep running until they finish", + len(pending), + grace, + ) diff --git a/litellm/litellm_core_utils/cli_token_utils.py b/litellm/litellm_core_utils/cli_token_utils.py index 37ed4efa0ca..5f45f942e4e 100644 --- a/litellm/litellm_core_utils/cli_token_utils.py +++ b/litellm/litellm_core_utils/cli_token_utils.py @@ -13,8 +13,10 @@ This module has no dependencies on proxy code and can be safely imported at the """ import math +import os import time from collections.abc import Mapping +from contextlib import suppress from dataclasses import dataclass from pathlib import Path from types import MappingProxyType @@ -42,6 +44,7 @@ from litellm.litellm_core_utils.private_json import ( ensure_private_dir, overwrite_private_json, stage_private_json, + write_private_bytes, write_private_json, ) from litellm.types.llms.base_model import LiteLLMBaseModel @@ -79,9 +82,21 @@ class CredentialNotCleared: detail: str +@dataclass(frozen=True, slots=True) +class LoginRecordNotMoved: + """The token file of a login the keychain kept could not be moved aside as the note. + + The file holds no secret, but it is the record a later logout pairs with the keychain entry to + revoke that login, so deleting it would leave the login live. It stays until its directory + takes the rename. + """ + + detail: str + + SecretSave: TypeAlias = SecretWrite | CredentialNotSaved | CredentialNotRecorded -SecretClear: TypeAlias = SecretErase | CredentialNotCleared +SecretClear: TypeAlias = SecretErase | CredentialNotCleared | LoginRecordNotMoved class CliTokenRecord(LiteLLMBaseModel): @@ -159,6 +174,7 @@ def save_cli_token(record: CliTokenRecord, *, vault: SecretVault = SYSTEM_KEYRIN keychain has already taken the new secret, and it reports itself as such rather than claiming the previous login survived. """ + replaced: Final = _read_token_file() stamped: Final = _stamped_past_every_stored_login(record, vault) staged: Final = _stage_token_file(_without_secret(stamped)) if isinstance(staged, CredentialNotSaved): @@ -167,7 +183,22 @@ def save_cli_token(record: CliTokenRecord, *, vault: SecretVault = SYSTEM_KEYRIN if isinstance(outcome, SecretStored): return outcome if _commit_token_file(staged) else CredentialNotRecorded() discard_staged_json(staged) - return _keep_the_secret_in_the_file(stamped, outcome) + kept: Final = _keep_the_secret_in_the_file(stamped, outcome) + if not isinstance(kept, CredentialNotSaved): + _note_the_login_left_in_the_keychain(replaced) + return kept + + +def _note_the_login_left_in_the_keychain(replaced: CliTokenRecord | None) -> None: + """Keep the login a save could not take out of the keychain revocable once the keychain answers. + + Its token file record is about to be overwritten by one that holds its own secret, so the note + is the only record left of the login whose refresh token the keychain still holds. + """ + if replaced is None or _holds_a_secret(replaced): + return + with suppress(OSError): + _write_keychain_not_cleared_note(_keychain_not_cleared_note_path(), replaced) def _stamped_past_every_stored_login(record: CliTokenRecord, vault: SecretVault) -> CliTokenRecord: @@ -220,27 +251,62 @@ def _keep_the_secret_in_the_file(record: CliTokenRecord, outcome: SecretWrite) - def clear_cli_token(*, vault: SecretVault = SYSTEM_KEYRING) -> SecretClear: """Remove the credential from both stores. Reports whether the keychain is now free of it. - A logout the keychain never answered keeps the token file, with its secret taken out, because - that file is the only remaining record that something may still be in there to remove. It is - what lets a later run tell a machine with a credential it cannot reach apart from one that never - had a login at all, and taking it away would leave the next logout answering the warning this - one just issued with a false all-clear. The secret goes either way, and a file that will give up - neither its copy nor itself is removed rather than kept, with the note written again afterwards - so the warning still outlives this run. + The token file goes whatever the keychain answers, so nothing left on disk reads as a login. A + keychain that could not be checked, or that would not give up its entry, may still hold a + credential of ours, so a logout that had something to end leaves a note beside the file, and the + next logout warns again instead of reading the missing file as a clean keychain. Only a keychain + that confirms the erase takes the note away. A file that will not go has its secret taken out + where it lies instead. """ outcome: Final = vault.erase() record: Final = _read_token_file() settled: Final = _nothing_left_behind(outcome, record) - if not settled and _keep_the_unchecked_keychain_on_record(outcome, record): - return outcome - removal: Final = _remove_token_file() - if removal is not None and record is not None and not _scrub_file_secret(record): - return removal - if removal is None and record is not None and _the_keychain_went_unchecked(outcome): - _write_the_note_the_removal_took_with_it(record) + left_on_disk: Final = _end_the_local_login(outcome, record) + if left_on_disk is not None: + return left_on_disk return SecretErased() if settled else outcome +def _end_the_local_login( + outcome: SecretErase, record: CliTokenRecord | None +) -> CredentialNotCleared | LoginRecordNotMoved | None: + if _names_a_login_the_keychain_kept(outcome, record): + return _move_token_file_to_note() + left_on_disk: Final = _remove_token_file_or_its_secret(record) + _note_what_the_keychain_answered(outcome, record) + return left_on_disk + + +def _names_a_login_the_keychain_kept(outcome: SecretErase, record: CliTokenRecord | None) -> bool: + match outcome: + case KeyringDisabled() | KeyringNotInstalled() | KeyringUnreachable(): + return record is not None and not _holds_a_secret(record) + case SecretErased() | SecretStranded(): + return False + + +def _move_token_file_to_note() -> LoginRecordNotMoved | None: + """Turn the token file of a login the keychain could not be asked to end into the note that ends it later. + + The file holds no secret, so it is already the record a later logout or login pairs with the + keychain entry to revoke the login, and renaming it needs no free space on a full disk. The file + never stays behind as a login that signs the user back in once the keychain unlocks: a note path + that will not take it leaves the file where it is, and the logout says so. + """ + try: + os.replace(get_cli_token_file_path(), _keychain_not_cleared_note_path()) + except OSError as error: + return LoginRecordNotMoved(str(error)) + return None + + +def _remove_token_file_or_its_secret(record: CliTokenRecord | None) -> CredentialNotCleared | None: + removal: Final = _remove_token_file() + if removal is None or record is None or _scrub_file_secret(record): + return None + return removal + + def _remove_token_file() -> CredentialNotCleared | None: try: Path(get_cli_token_file_path()).unlink(missing_ok=True) @@ -249,54 +315,87 @@ def _remove_token_file() -> CredentialNotCleared | None: return None -def _write_the_note_the_removal_took_with_it(record: CliTokenRecord) -> None: - """Put the secret-free note back after the file carrying it had to go to get the secret off disk. +def _keychain_not_cleared_note_path() -> Path: + return Path(get_cli_token_file_path()).with_name("keychain_not_cleared") - Reaching here means neither rewrite would take, so the file went instead, and its absence is - what the next logout would read as a keychain already known to be clean. Removing it is also - what frees the room the rewrite was refused for, so the note usually lands on this second try. - When it does not, the warning this logout printed is the only one the user gets. + +def load_cli_token_left_in_keychain(*, vault: SecretVault = SYSTEM_KEYRING) -> CliTokenRecord | None: + """The login an earlier logout could not erase from the keychain, once the keychain answers again. + + The token file went with that logout, so its note is the only record of the login, and the + refresh token that revokes it is only in the keychain. A keychain entry from any other sign-in + is never paired with the note, since the endpoints it names belong to that one login. """ - staged: Final = _stage_scrubbed_file(record) - if staged is not None: - _commit_token_file(staged) + note: Final = _read_keychain_not_cleared_note() + if note is None: + return None + secret: Final = _stored_secret(note.base_url, vault) + if secret is None or secret.timestamp != note.timestamp: + return None + return note.model_copy( + update=MappingProxyType( + {"key": secret.key, "jwt_token": secret.jwt_token, "refresh_token": secret.refresh_token} + ) + ) -def _keep_the_unchecked_keychain_on_record(outcome: SecretErase, record: CliTokenRecord | None) -> bool: - """Whether the token file, stripped of its secret, is worth keeping as the note that says so. - - Only a keychain that could not be reached leaves the question open. One that answered for itself - is remembered without any help from the file, and a file it can still pair a live entry with - would leave the machine signed in to the login that was just ended. A copy that will give up - its secret neither to a staged replacement nor to an overwrite is not kept either, because the - secret goes first. - """ - if record is None or not _the_keychain_went_unchecked(outcome): +def _holds_the_noted_login(blob: str) -> bool: + """Whether the keychain entry is the login an earlier logout could not end, which only it can revoke""" + note: Final = _read_keychain_not_cleared_note() + if note is None: return False - return _scrub_file_secret(record) + secret: Final = _decode_secret(blob, note.base_url) + return secret is not None and secret.timestamp == note.timestamp -def _the_keychain_went_unchecked(outcome: SecretErase) -> bool: - """Whether the keychain neither confirmed the erase nor answered that it still holds the secret""" - match outcome: - case SecretErased() | SecretStranded(): - return False - case KeyringDisabled() | KeyringNotInstalled() | KeyringUnreachable(): - return True +def _read_keychain_not_cleared_note() -> CliTokenRecord | None: + try: + return CliTokenRecord.model_validate_json(_keychain_not_cleared_note_path().read_bytes()) + except (OSError, ValidationError): + return None + + +def _note_what_the_keychain_answered(outcome: SecretErase, record: CliTokenRecord | None) -> None: + """Keep the note that a logout left the keychain uncleared for exactly as long as that is true. + + A keychain that would not give up its entry had already handed this logout the refresh token, + so its note holds nothing to revoke again. A login that kept its secret in the file never + reached the keychain, so the note an earlier logout left still names the login the keychain + holds and stays, and a logout with no file to end keeps it too. + + The secret is off disk by the time this runs, so a note that will not be written or removed is + no reason to fail the logout. The warning printed now is then the only one the user gets. + """ + note: Final = _keychain_not_cleared_note_path() + with suppress(OSError): + match outcome: + case SecretErased(): + note.unlink(missing_ok=True) + case SecretStranded(): + _write_keychain_not_cleared_note(note, None) + case KeyringDisabled() | KeyringNotInstalled() | KeyringUnreachable(): + if record is not None and not note.exists(): + _write_keychain_not_cleared_note(note, None) + + +def _write_keychain_not_cleared_note(note: Path, record: CliTokenRecord | None) -> None: + """Record the login a later command can still revoke, or only that the keychain was left uncleared""" + ensure_private_dir(note.parent) + if record is None: + write_private_bytes(str(note), b"") + return + write_private_json(str(note), _without_secret(record).model_dump(exclude_none=True)) def _nothing_left_behind(outcome: SecretErase, record: CliTokenRecord | None) -> bool: """Whether the keychain can be trusted to hold no credential of ours once the file is gone. - A machine with no token file has no stored login to end, and `clear_cli_token` keeps one behind - whenever the keychain is left unconfirmed, taking the secret out in place when it cannot stage a - replacement and writing the note again when the file holding it had to go, so a missing file is - real evidence rather than the absence of it. Past that, a - keychain that could not be reached is never trusted, whatever the file looks like. Even a file - holding its own secret says only that the login which wrote it had no keychain to write to, and - the login before it may well have had one: the entry that login left outlives both the - uninstalled package and the file that replaced it. `SecretStranded` is - the keychain answering for itself and outranks the file. + A keychain that could not be reached is trusted only on a machine with nothing to end: no token + file, and no note from an earlier logout that left the keychain uncleared. Even a file holding + its own secret says only that the login which wrote it had no keychain to write to, and the + login before it may well have had one: the entry that login left outlives both the uninstalled + package and the file that replaced it. `SecretStranded` is the keychain answering for itself and + outranks the file. """ match outcome: case SecretErased(): @@ -304,7 +403,7 @@ def _nothing_left_behind(outcome: SecretErase, record: CliTokenRecord | None) -> case SecretStranded(): return False case KeyringDisabled() | KeyringNotInstalled() | KeyringUnreachable(): - return record is None + return record is None and not os.path.exists(_keychain_not_cleared_note_path()) def get_litellm_gateway_api_key( @@ -407,7 +506,14 @@ def _apply_vault_secret(record: CliTokenRecord, blob: str, vault: SecretVault) - than counting as a clean sweep. Rolling the vault back the way a migration does is not the answer here, because the two stores hold different credentials and the rollback would hand the superseded one back out. + + A keychain entry an earlier logout could not end is left alone whatever the stamps say. Its note + pairs only with that entry, so the next login or logout can still revoke it, while serving it + would sign an ended login back in and moving the file's secret over it would leave its refresh + token live with nothing on this machine able to revoke it. """ + if _holds_a_secret(record) and _holds_the_noted_login(blob): + return record secret: Final = _decode_secret(blob, record.base_url) if secret is None or (_holds_a_secret(record) and secret.timestamp <= record.timestamp): return _migrate_file_secret(_rejoined(record, secret), vault, replacing=secret) diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index f31d4205b82..800973573c9 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -113,6 +113,27 @@ _Key = TypeVar("_Key") # rebind-ok: a TypeVar cannot be Final _Value = TypeVar("_Value") # rebind-ok: a TypeVar cannot be Final +INCLUDE_COST_IN_STREAMING_USAGE_HEADER: Final = "x-litellm-include-cost-in-streaming-usage" + + +def proxy_request_headers(container: Mapping[str, object] | None) -> Mapping[str, object] | None: + proxy_server_request: Final = as_str_mapping(container.get("proxy_server_request")) if container else None + return as_str_mapping(proxy_server_request.get("headers")) if proxy_server_request else None + + +def stream_cost_injection_requested(request_headers: Mapping[str, object] | None) -> bool: + """The global include_cost_in_streaming_usage flag or the per-request header asks for usage.cost on streams.""" + if bool(getattr(litellm, "include_cost_in_streaming_usage", False)): + return True + if not request_headers: + return False + header_value: Final = next( + (value for name, value in request_headers.items() if name.lower() == INCLUDE_COST_IN_STREAMING_USAGE_HEADER), + None, + ) + return isinstance(header_value, str) and header_value.strip().lower() == "true" + + def without( mapping: Mapping[_Key, _Value], *keys: object ) -> dict[_Key, _Value]: # mutable-ok: a fresh copy the caller owns diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index b2851345302..bbb1b19819f 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -194,11 +194,7 @@ def get_llm_provider( resolved_model, provider, dynamic_api_key, resolved_api_base = _get_llm_provider( model, custom_llm_provider, api_base, api_key, litellm_params ) - explicit_provider: Final = custom_llm_provider or (litellm_params.custom_llm_provider if litellm_params else None) - target: Final = ( - resolve_model_latest_alias(resolved_model, provider) if explicit_provider or "/" in model else resolved_model - ) - return target, provider, dynamic_api_key, resolved_api_base + return resolve_model_latest_alias(resolved_model, provider), provider, dynamic_api_key, resolved_api_base def _get_llm_provider( diff --git a/litellm/litellm_core_utils/internal_call_metadata.py b/litellm/litellm_core_utils/internal_call_metadata.py index 138a01c6524..de17bcbdfb9 100644 --- a/litellm/litellm_core_utils/internal_call_metadata.py +++ b/litellm/litellm_core_utils/internal_call_metadata.py @@ -46,7 +46,18 @@ ITPM_RESERVED_KEY: Final = "_litellm_itpm_reserved" OTPM_RESERVED_KEY: Final = "_litellm_otpm_reserved" ITPM_CACHE_KEY: Final = "_litellm_itpm_cache_key" OTPM_CACHE_KEY: Final = "_litellm_otpm_cache_key" -IO_TOKEN_RESERVATION_KEYS: Final = frozenset({ITPM_RESERVED_KEY, OTPM_RESERVED_KEY, ITPM_CACHE_KEY, OTPM_CACHE_KEY}) +IO_TOKEN_ATTEMPT_KEY: Final = "_litellm_io_token_attempt" +IO_TOKEN_RESERVATION_ID_KEY: Final = "_litellm_io_token_reservation_id" +IO_TOKEN_RESERVATION_KEYS: Final = frozenset( + { + ITPM_RESERVED_KEY, + OTPM_RESERVED_KEY, + ITPM_CACHE_KEY, + OTPM_CACHE_KEY, + IO_TOKEN_ATTEMPT_KEY, + IO_TOKEN_RESERVATION_ID_KEY, + } +) MODEL_ACCESS_GROUP_METADATA_KEY: Final = "user_api_key_matched_model_access_groups" """Where auth records the model access groups that authorized the request, for the spend writer. diff --git a/litellm/litellm_core_utils/model_latest_aliases.py b/litellm/litellm_core_utils/model_latest_aliases.py index e3b1a21dd5a..ee45ccc9b3c 100644 --- a/litellm/litellm_core_utils/model_latest_aliases.py +++ b/litellm/litellm_core_utils/model_latest_aliases.py @@ -1,33 +1,75 @@ from collections.abc import Mapping -from functools import cache -from importlib.resources import files from types import MappingProxyType -from typing import Final +from typing import Final, cast # noqa: TID251 # litellm.model_cost is an untyped module global -from pydantic import TypeAdapter - -_ALIASES_ADAPTER: Final = TypeAdapter(dict[str, dict[str, str]]) +import litellm +from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry + + +def _cost_map() -> Mapping[str, object]: + return cast("Mapping[str, object]", litellm.model_cost) # cast-ok: per-lookup catalog validation is too slow + + +def _alias_entry(key: str) -> AliasedCostEntry | None: + entry: Final = _cost_map().get(key) + return entry if isinstance(entry, AliasedCostEntry) and entry.names[:1] not in ((), (key,)) else None + + +def _alias_provider(key: str, entry: AliasedCostEntry) -> str | None: + """The provider an alias is requested through: its key prefix for a hosted alias such as + ``bedrock/claude-haiku-latest`` (whose target row belongs to ``bedrock_converse``), else the target's provider""" + prefix, separator, _ = key.partition("/") + if separator and prefix in litellm.provider_list: + return prefix + provider: Final = entry.get("litellm_provider") + return provider if isinstance(provider, str) else None + + +def _alias_target(key: str, provider: str) -> str | None: + entry: Final = _alias_entry(key) + if entry is None or _alias_provider(key, entry) != provider: + return None + return entry.names[0].removeprefix(f"{provider}/") -@cache def get_model_latest_aliases() -> Mapping[str, Mapping[str, str]]: - aliases: Final = _ALIASES_ADAPTER.validate_json(files("litellm").joinpath("model_latest_aliases.json").read_bytes()) - return MappingProxyType({provider: MappingProxyType(models) for provider, models in aliases.items()}) + """``{provider: {alias: target}}`` for every alias row in the cost map, with provider prefixes stripped""" + rows: Final = tuple( + (provider, key, entry.names[0]) + for key in _cost_map() + if (entry := _alias_entry(key)) is not None and (provider := _alias_provider(key, entry)) is not None + ) + return MappingProxyType( + { + provider: MappingProxyType( + { + alias.removeprefix(f"{provider}/"): target.removeprefix(f"{provider}/") + for row_provider, alias, target in rows + if row_provider == provider + } + ) + for provider in dict.fromkeys(provider for provider, _, _ in rows) + } + ) def resolve_model_latest_alias(model: str, provider: str) -> str: - aliases: Final = get_model_latest_aliases().get(provider) - return aliases.get(model, model) if aliases is not None else model + return _alias_target(f"{provider}/{model}", provider) or _alias_target(model, provider) or model def expand_model_latest_alias(model: str) -> tuple[str, ...]: provider, separator, name = model.partition("/") - target: Final = resolve_model_latest_alias(name, provider) if separator else name - return (model, f"{provider}/{target}") if separator and target != name else (model,) + if not separator: + entry: Final = _alias_entry(model) + return (model,) if entry is None else (model, entry.names[0]) + target: Final = resolve_model_latest_alias(name, provider) + return (model, f"{provider}/{target}") if target != name else (model,) def is_model_latest_alias(model: str, custom_llm_provider: str | None = None) -> bool: prefix, separator, name = model.partition("/") - provider: Final = custom_llm_provider or prefix + provider: Final = custom_llm_provider or (prefix if separator else None) + if provider is None: + return _alias_entry(model) is not None candidate: Final = name if separator and prefix == provider else model - return bool(custom_llm_provider or separator) and resolve_model_latest_alias(candidate, provider) != candidate + return resolve_model_latest_alias(candidate, provider) != candidate diff --git a/litellm/litellm_core_utils/total_deadline.py b/litellm/litellm_core_utils/total_deadline.py new file mode 100644 index 00000000000..f44eb97719d --- /dev/null +++ b/litellm/litellm_core_utils/total_deadline.py @@ -0,0 +1,197 @@ +""" +One wall-clock budget for a whole non-streaming proxy request (``total_timeout``): the proxy opens it, +and the router reads it to refuse a retry backoff that cannot finish in time. A request that turns out +to stream or run in the background is not bounded: a body ``total_timeout`` refuses it, a header one +(often set globally on a client) lets it go unbounded. +""" + +import asyncio +import contextlib +import math +import sys +from collections.abc import AsyncGenerator, Callable, Generator, Mapping +from contextvars import ContextVar +from dataclasses import dataclass +from typing import Final + +from litellm._logging import verbose_router_logger + +if sys.version_info >= (3, 11): + from asyncio import Timeout as _Scope + from asyncio import timeout_at as _timeout_at + + _ScopeTimeout: Final = TimeoutError + + def _expired(scope: _Scope) -> bool: + return scope.expired() + + def _disarm(scope: _Scope) -> None: + scope.reschedule(None) + + def _rearm(scope: _Scope, when: float) -> None: + scope.reschedule(when) + +else: + from async_timeout import Timeout as _Scope + from async_timeout import timeout_at as _timeout_at + + _ScopeTimeout: Final = asyncio.TimeoutError + + def _expired(scope: _Scope) -> bool: + return scope.expired + + def _disarm(scope: _Scope) -> None: + scope.reject() + + def _rearm(scope: _Scope, when: float) -> None: + scope.update(when) + + +TOTAL_TIMEOUT_HEADER: Final = "x-litellm-total-timeout" +TOTAL_TIMEOUT_FIELD: Final = "total_timeout" + + +@dataclass(frozen=True, slots=True) +class TotalTimeout: + seconds: float + from_body: bool + + +@dataclass(frozen=True, slots=True) +class InvalidTotalTimeout: + message: str + + +def _positive_finite_seconds(value: object) -> float | None: + if isinstance(value, bool) or not isinstance(value, int | float | str): + return None + try: + seconds: Final = float(value) + except (ValueError, OverflowError): + return None + return seconds if math.isfinite(seconds) and seconds > 0 else None + + +def _total_timeout(value: object, *, from_body: bool) -> TotalTimeout | InvalidTotalTimeout: + seconds: Final = _positive_finite_seconds(value) + if seconds is None: + name: Final = TOTAL_TIMEOUT_FIELD if from_body else TOTAL_TIMEOUT_HEADER + return InvalidTotalTimeout(f"{name} must be a positive, finite number of seconds, got {value!r}") + return TotalTimeout(seconds, from_body) + + +def total_timeout_from_request( + header: str | None, body: Mapping[str, object] +) -> TotalTimeout | InvalidTotalTimeout | None: + """The request's ``total_timeout``: the header wins over the body field.""" + if header is not None: + return _total_timeout(header, from_body=False) + if TOTAL_TIMEOUT_FIELD not in body: + return None + return _total_timeout(body[TOTAL_TIMEOUT_FIELD], from_body=True) + + +def bounded_by_total_timeout(call: Mapping[str, object]) -> bool: + """Only a request that neither streams nor runs in the background can be bounded.""" + return call.get("stream") is not True and call.get("background") is not True + + +class TotalTimeoutNotApplicable(Exception): + """A request bounded by a body ``total_timeout`` streams or runs in the background.""" + + +class TotalTimeoutExpired(TimeoutError): + """The request's ``total_timeout`` ran out.""" + + +class TotalDeadline: + __slots__ = ("_at", "_clock", "_released", "_scope", "_scope_at", "from_body") + + def __init__( + self, at: float, clock: Callable[[], float], scope: _Scope, scope_at: float, *, from_body: bool + ) -> None: + self._at: Final = at + self._clock: Final = clock + self._scope: Final = scope + self._scope_at: Final = scope_at + self._released: bool = False + self.from_body: Final = from_body + + def expired(self) -> bool: + return _expired(self._scope) + + def allows_wait(self, seconds: float) -> bool: + return self._released or self._clock() + seconds < self._at + + @contextlib.contextmanager + def held(self) -> Generator[None, None, None]: + """Run the block without interrupting it; a deadline that passed meanwhile fires right after it.""" + if self._released: + yield + return + _disarm(self._scope) + try: + yield + finally: + if not self._released: + _rearm(self._scope, self._scope_at) + + def release(self) -> None: + self._released = True + if not self.expired(): + _disarm(self._scope) + + +_CURRENT: Final[ContextVar[TotalDeadline | None]] = ContextVar("total_deadline", default=None) + + +def current_total_deadline() -> TotalDeadline | None: + return _CURRENT.get() + + +@contextlib.asynccontextmanager +async def total_timeout_held() -> AsyncGenerator[None, None]: + deadline: Final = _CURRENT.get() + if deadline is None: + yield + return + with deadline.held(): + yield + + +def enforce_total_timeout_scope(call: Mapping[str, object]) -> None: + """Raise ``TotalTimeoutNotApplicable`` for a body ``total_timeout`` on a streaming or background call, + and stop bounding it for a header one.""" + deadline: Final = _CURRENT.get() + if deadline is None or deadline._released or bounded_by_total_timeout(call): # pyright: ignore[reportPrivateUsage] # same module + return + if deadline.from_body: + raise TotalTimeoutNotApplicable( + "total_timeout applies only to non-streaming, non-background requests; drop it or the stream/background flag" + ) + verbose_router_logger.debug("x-litellm-total-timeout ignored: the request streams or runs in the background") + deadline.release() + + +@contextlib.asynccontextmanager +async def total_deadline( + seconds: float, *, from_body: bool, clock: Callable[[], float] | None = None +) -> AsyncGenerator[TotalDeadline, None]: + """Bound the block by ``seconds``: past it the block is cancelled and ``TotalTimeoutExpired`` raised, + while a ``TimeoutError`` raised inside the block passes through as it was.""" + loop: Final = asyncio.get_running_loop() + now: Final = clock if clock is not None else loop.time + scope_at: Final = loop.time() + seconds + scope: Final = _timeout_at(scope_at) + deadline: Final = TotalDeadline(now() + seconds, now, scope, scope_at, from_body=from_body) + try: + async with scope: + token: Final = _CURRENT.set(deadline) + try: + yield deadline + finally: + _CURRENT.reset(token) + except _ScopeTimeout as error: + if not deadline.expired(): + raise + raise TotalTimeoutExpired(f"the request did not finish within its total_timeout of {seconds}s") from error diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index e0d313175b7..48bccc7c001 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -1718,7 +1718,13 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): llm_provider: str = "anthropic", ) -> AnthropicThinkingParam | None: """Capability probes read the cost map under ``custom_llm_provider``; ``llm_provider`` only tags raised exceptions.""" - if reasoning_effort is None or reasoning_effort == "none": + if reasoning_effort is None: + return None + if reasoning_effort == "none": + if AnthropicModelInfo.thinks_by_default( + model, custom_llm_provider + ) and not AnthropicModelInfo._is_always_on_thinking_model(model, custom_llm_provider): + return AnthropicThinkingParam(type="disabled") return None if AnthropicConfig._is_adaptive_thinking_model(model, custom_llm_provider): # without display, Anthropic defaults adaptive thinking to @@ -2126,12 +2132,14 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): custom_llm_provider=self._resolved_provider, llm_provider=self._resolved_provider, ) - if mapped_thinking is None: + if mapped_thinking is None or mapped_thinking.get("type") == "disabled": optional_params.pop("thinking", None) optional_params.pop("output_config", None) - else: + if mapped_thinking is not None: optional_params["thinking"] = mapped_thinking - if AnthropicConfig._is_adaptive_thinking_model(model, self._resolved_provider): + if mapped_thinking.get("type") != "disabled" and AnthropicConfig._is_adaptive_thinking_model( + model, self._resolved_provider + ): mapped_effort = REASONING_EFFORT_TO_OUTPUT_CONFIG_EFFORT.get(effort_value) if mapped_effort is None: AnthropicConfig._raise_invalid_reasoning_effort( diff --git a/litellm/llms/anthropic/messages/native/transformation.py b/litellm/llms/anthropic/messages/native/transformation.py index 5281ad99620..aa0693d6f03 100644 --- a/litellm/llms/anthropic/messages/native/transformation.py +++ b/litellm/llms/anthropic/messages/native/transformation.py @@ -471,7 +471,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): """Map OpenAI-style ``reasoning_effort`` to native Anthropic params. Caller-supplied ``thinking`` / ``output_config`` win over the alias. - ``effort='none'`` clears both. Invalid efforts raise a 400. A mapped + ``effort='none'`` disables optional default thinking and clears effort. Invalid efforts raise a 400. A mapped thinking budget is capped below ``max_tokens`` and dropped when even the minimum budget cannot fit. """ @@ -485,6 +485,10 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): if not isinstance(reasoning_effort, str): return + if reasoning_effort == "none": + optional_params.pop("thinking", None) + optional_params.pop("output_config", None) + try: mapped_thinking: Final = AnthropicConfig.map_reasoning_effort( reasoning_effort=reasoning_effort, @@ -495,8 +499,6 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): raise AnthropicError(message=str(e.message), status_code=400) if mapped_thinking is None: - optional_params.pop("thinking", None) - optional_params.pop("output_config", None) return fitted_thinking: Final = AnthropicConfig.cap_thinking_budget_to_max_tokens(mapped_thinking, max_tokens) @@ -505,7 +507,9 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): return optional_params.setdefault("thinking", fitted_thinking) - if AnthropicModelInfo.is_adaptive_thinking_model(model, custom_llm_provider): + if fitted_thinking.get("type") != "disabled" and AnthropicModelInfo.is_adaptive_thinking_model( + model, custom_llm_provider + ): mapped_effort: Final = REASONING_EFFORT_TO_OUTPUT_CONFIG_EFFORT.get(reasoning_effort) if mapped_effort is None: raise AnthropicError( diff --git a/litellm/llms/anthropic/wif.py b/litellm/llms/anthropic/wif.py index ccb81c1244f..d2bf3e893af 100644 --- a/litellm/llms/anthropic/wif.py +++ b/litellm/llms/anthropic/wif.py @@ -552,9 +552,12 @@ def _resolve_assertion_ref(litellm_params: Mapping[str, object] | None) -> str | file_env: Final = _env_str("ANTHROPIC_IDENTITY_TOKEN_FILE") if file_env is not None: return f"oidc/file/{file_env}" - if _env_str(_INLINE_ENV_VAR) is not None: - return f"oidc/env/{_INLINE_ENV_VAR}" - return None + inline_env: Final = _env_str(_INLINE_ENV_VAR) + if inline_env is None: + return None + if inline_env.startswith(_ACCEPTED_REF_PREFIX): + return _validated_inline_ref(inline_env) + return f"oidc/env/{_INLINE_ENV_VAR}" def _validated_inline_ref(value: str) -> str: diff --git a/litellm/llms/aws_polly/text_to_speech/transformation.py b/litellm/llms/aws_polly/text_to_speech/transformation.py index 3b52a8df3f2..40a50e150f4 100644 --- a/litellm/llms/aws_polly/text_to_speech/transformation.py +++ b/litellm/llms/aws_polly/text_to_speech/transformation.py @@ -7,7 +7,7 @@ Reference: https://docs.aws.amazon.com/polly/latest/dg/API_SynthesizeSpeech.html import json from collections.abc import Coroutine -from typing import TYPE_CHECKING, Any, Final, Union +from typing import TYPE_CHECKING, Any, ClassVar, Final, Union import httpx @@ -18,6 +18,7 @@ from litellm.llms.base_llm.text_to_speech.transformation import ( TextToSpeechRequestData, ) from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.web_identity_session_policy import POLLY_SESSION_POLICY, WebIdentitySessionPolicy if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -35,6 +36,8 @@ class AWSPollyTextToSpeechConfig(BaseTextToSpeechConfig, BaseAWSLLM): Reference: https://docs.aws.amazon.com/polly/latest/dg/API_SynthesizeSpeech.html """ + _web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = POLLY_SESSION_POLICY + def __init__(self): BaseTextToSpeechConfig.__init__(self) BaseAWSLLM.__init__(self) diff --git a/litellm/llms/base_llm/auth/token_exchange.py b/litellm/llms/base_llm/auth/token_exchange.py index 5bf1f4dd0db..5b2ab86dd53 100644 --- a/litellm/llms/base_llm/auth/token_exchange.py +++ b/litellm/llms/base_llm/auth/token_exchange.py @@ -9,6 +9,7 @@ public exception contract. import asyncio import hashlib +import hmac import json import re import threading @@ -291,8 +292,12 @@ def _cache_key(spec: TokenExchangeSpec) -> str: def _shares_one_assertion_across_workers(spec: TokenExchangeSpec) -> bool: """The store exists so the workers reading one projected token file don't each spend that file's single-use ``jti``. A source that mints its own assertion per exchange shares nothing with another - worker, so it never reads the store, never finds a hit there, and keeps its minted token off disk.""" - return spec.assertion_source is None + worker, and a Google metadata-server token carries no ``jti`` (Anthropic only enforces single use on + assertions that do), so a stored token would only hand a mandatory refresh back the token it is + replacing. Neither reads the store, finds a hit there, or puts its minted token on disk.""" + from litellm.secret_managers.main import GOOGLE_OIDC_SECRET_PREFIX + + return spec.assertion_source is None and not spec.assertion_ref.startswith(GOOGLE_OIDC_SECRET_PREFIX) def _assertion_digest(assertion: SecretStr) -> str: @@ -538,10 +543,22 @@ class _Entry: """Single-flight state for one cache key; mutable by design, confined to the engine, and only ever mutated under the engine lock.""" - __slots__ = ("backoff_until", "done", "force_refresh", "in_flight", "last_error", "lifetime_seconds", "token") + __slots__ = ( + "backoff_until", + "done", + "force_refresh", + "in_flight", + "last_error", + "lifetime_seconds", + "previous_token", + "spec", + "token", + ) - def __init__(self, force_refresh: bool = False) -> None: + def __init__(self, spec: TokenExchangeSpec, force_refresh: bool = False) -> None: + self.spec: Final = spec self.token: MintedToken | None = None + self.previous_token: MintedToken | None = None self.lifetime_seconds: float | None = None self.in_flight: bool = False self.done: Final = threading.Event() @@ -562,7 +579,14 @@ class _Entry: self.in_flight = False self.done.set() + def issued(self, access_token: str) -> bool: + return any( + minted is not None and hmac.compare_digest(minted.access_token.get_secret_value(), access_token) + for minted in (self.token, self.previous_token) + ) + def _store(self, token: MintedToken, now: float) -> None: + self.previous_token = self.token self.token = token self.lifetime_seconds = None if token.expires_at is None else max(token.expires_at - now, 0.0) self.last_error = None @@ -684,11 +708,17 @@ class JwtBearerTokenExchangeEngine: async def aget_token(self, spec: TokenExchangeSpec) -> ExchangeResult: return await asyncio.to_thread(self.get_token, spec) + def spec_for_access_token(self, access_token: str) -> TokenExchangeSpec | None: + """Finds the identity a bearer was minted for, so resending an already built request can swap in + the current token instead of one that expired while the first attempt was in flight.""" + with self._lock: + return next((entry.spec for entry in self._entries.values() if entry.issued(access_token)), None) + def invalidate(self, spec: TokenExchangeSpec) -> None: key: Final = _cache_key(spec) with self._lock: if key in self._entries: - self._entries[key] = _Entry(force_refresh=True) + self._entries[key] = _Entry(spec, force_refresh=True) if self._shared_store is not None: self._shared_store.delete(key) @@ -699,7 +729,7 @@ class JwtBearerTokenExchangeEngine: return existing if len(self._entries) >= self._max_entries: self._evict_locked() - created: Final = _Entry() + created: Final = _Entry(spec) self._entries[key] = created return created diff --git a/litellm/llms/base_llm/batches/transformation.py b/litellm/llms/base_llm/batches/transformation.py index e7f3f1d277a..f25d077b2fe 100644 --- a/litellm/llms/base_llm/batches/transformation.py +++ b/litellm/llms/base_llm/batches/transformation.py @@ -47,8 +47,7 @@ class BaseBatchesConfig(ABC): return { k: v for k, v in class_attributes.items() - if not k.startswith("__") - and not k.startswith("_abc") + if not k.startswith("_") and not isinstance( v, ( diff --git a/litellm/llms/base_llm/text_to_speech/transformation.py b/litellm/llms/base_llm/text_to_speech/transformation.py index 6df83966bf8..15e6e515fe4 100644 --- a/litellm/llms/base_llm/text_to_speech/transformation.py +++ b/litellm/llms/base_llm/text_to_speech/transformation.py @@ -45,8 +45,7 @@ class BaseTextToSpeechConfig(ABC): return { k: v for k, v in class_attributes.items() - if not k.startswith("__") - and not k.startswith("_abc") + if not k.startswith("_") and not isinstance( v, ( diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 9591b962ab5..b0ba7470433 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -9,7 +9,7 @@ import urllib.parse from collections.abc import Callable, Mapping, MutableMapping, Sequence from concurrent.futures import ThreadPoolExecutor from datetime import datetime -from functools import partial +from functools import cache, partial from threading import Lock from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, ParamSpec, TypeVar, cast, get_args, overload @@ -33,6 +33,11 @@ from litellm.constants import ( from litellm.litellm_core_utils.aws_partition import contains_bedrock_arn, get_aws_dns_suffix from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.optional_imports import ensure_optional_import +from litellm.llms.bedrock.web_identity_session_policy import ( + BEDROCK_SESSION_POLICY, + WebIdentitySessionPolicy, + session_policy_document, +) from litellm.secret_managers.main import get_secret, get_secret_str from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.bedrock import AWS_AUTH_PARAM_KEYS, AwsAuthParams, AwsSessionTag, BearerPreparedRequest @@ -105,95 +110,6 @@ def _assume_role_params( ) -_SecureTransportBool: Final = TypedDict("_SecureTransportBool", {"aws:SecureTransport": ReadOnly[Literal["true"]]}) - -_CalledViaLastEquals: Final = TypedDict( - "_CalledViaLastEquals", {"aws:CalledViaLast": ReadOnly[Literal["aws-external-anthropic.amazonaws.com"]]} -) - - -class _SessionPolicyCondition(TypedDict): - Bool: ReadOnly[_SecureTransportBool] - StringEquals: NotRequired[ReadOnly[_CalledViaLastEquals]] - - -class _SessionPolicyStatement(TypedDict): - Sid: ReadOnly[str] - Effect: ReadOnly[Literal["Allow"]] - Action: ReadOnly[tuple[str, ...]] - Resource: ReadOnly[Literal["*"]] - Condition: ReadOnly[_SessionPolicyCondition] - - -class WebIdentitySessionPolicy(TypedDict): - Version: ReadOnly[Literal["2012-10-17"]] - Statement: ReadOnly[tuple[_SessionPolicyStatement, ...]] - - -_WEB_IDENTITY_SESSION_POLICY_ACTIONS: Final[Mapping[str, tuple[str, ...]]] = MappingProxyType( - { - "BedrockLiteLLM": ( - "bedrock:InvokeModel", - "bedrock:InvokeModelWithResponseStream", - "bedrock:CountTokens", - "bedrock:Rerank", - "bedrock:Retrieve", - "bedrock:ListKnowledgeBases", - "bedrock:InvokeAgent", - "bedrock:ApplyGuardrail", - "bedrock:GetGuardrail", - "bedrock:ListGuardrails", - ), - "BedrockAgentCoreLiteLLM": ( - "bedrock-agentcore:InvokeAgentRuntime", - "bedrock-agentcore:InvokeAgentRuntimeForUser", - "bedrock-agentcore:InvokeGateway", - ), - "ClaudePlatformLiteLLM": ( - "aws-external-anthropic:CreateInference", - "aws-external-anthropic:CreateBatchInference", - "aws-external-anthropic:CancelBatchInference", - "aws-external-anthropic:DeleteBatchInference", - "aws-external-anthropic:CountTokens", - "aws-external-anthropic:Get*", - "aws-external-anthropic:List*", - ), - "ClaudePlatformOutboundIdentityLiteLLM": ( - "sts:GetWebIdentityToken", - "sts:TagGetWebIdentityToken", - ), - "BedrockMantleLiteLLM": ("bedrock-mantle:CreateInference",), - } -) - -_SECURE_TRANSPORT_ONLY: Final = _SessionPolicyCondition(Bool=_SecureTransportBool({"aws:SecureTransport": "true"})) - -_WEB_IDENTITY_SESSION_POLICY_CONDITIONS: Final[Mapping[str, _SessionPolicyCondition]] = MappingProxyType( - { - "ClaudePlatformOutboundIdentityLiteLLM": _SessionPolicyCondition( - Bool=_SecureTransportBool({"aws:SecureTransport": "true"}), - StringEquals=_CalledViaLastEquals({"aws:CalledViaLast": "aws-external-anthropic.amazonaws.com"}), - ), - } -) - - -def build_web_identity_session_policy() -> WebIdentitySessionPolicy: - return WebIdentitySessionPolicy( - Version="2012-10-17", - Statement=tuple( - _SessionPolicyStatement( - Sid=sid, - Effect="Allow", - Action=actions, - Resource="*", - Condition=_WEB_IDENTITY_SESSION_POLICY_CONDITIONS.get(sid, _SECURE_TRANSPORT_ONLY), - ) - for sid, actions in _WEB_IDENTITY_SESSION_POLICY_ACTIONS.items() - ), - ) - - def pop_aws_auth_params( optional_params: MutableMapping[str, object], # mutable-ok: pops the aws_* keys out of the caller's mapping ) -> AwsAuthParams: @@ -246,7 +162,8 @@ class BaseAWSLLM(SignsRequestsWithAWS): # (600 seconds / 10 minutes); web identity STS credentials use # ``_get_default_ttl_for_boto3_credentials`` (~59 minutes); AssumeRole STS credentials expire with # the STS session itself (Expiration minus a safety margin). All are keyed on all aws_* credential - # args plus ssl_verify, so ``aws_session_name`` scopes an entry to one attributed identity. Profiles + # args plus ssl_verify, and web identity entries also on the caller class's session policy, so + # ``aws_session_name`` scopes an entry to one attributed identity. Profiles # and explicit session-token tuples are not cached — see ``get_credentials`` and # ``_get_or_set_cached_credentials``. The bound is larger than ``InMemoryCache``'s default because # per-user cost attribution puts one entry per attributed identity in this cache. @@ -263,6 +180,13 @@ class BaseAWSLLM(SignsRequestsWithAWS): Lock() for _ in range(BEDROCK_IAM_CACHE_FETCH_LOCK_STRIPES) ) + _web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = BEDROCK_SESSION_POLICY + + @classmethod + @cache + def _web_identity_session_policy_document(cls) -> str: + return session_policy_document(cls._web_identity_session_policy) + def __init__(self) -> None: self.iam_cache = BaseAWSLLM._shared_iam_cache super().__init__() @@ -478,7 +402,7 @@ class BaseAWSLLM(SignsRequestsWithAWS): aws_session_name, ): return self._get_or_set_cached_credentials( - args, + {**args, "web_identity_session_policy": self._web_identity_session_policy_document()}, lambda: self._auth_with_web_identity_token( aws_web_identity_token=cast(str, aws_web_identity_token), aws_role_name=cast(str, aws_role_name), @@ -1095,9 +1019,10 @@ class BaseAWSLLM(SignsRequestsWithAWS): # The session policy is an IAM PERMISSION CEILING — effective # permissions are the intersection of the role's identity policies # and this policy. Any action not listed here is silently denied - # even when the IAM role grants it. So every Bedrock route we - # support needs a matching action statement, or it 403s on OIDC - # auth only (static creds + IRSA take other code paths). + # even when the IAM role grants it. So each caller class's + # _web_identity_session_policy must list every action its routes + # call, or they 403 on OIDC auth only (static creds + IRSA take + # other code paths). # https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRoleWithWebIdentity.html # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sts/client/assume_role_with_web_identity.html assume_role_params: Final = { @@ -1105,7 +1030,7 @@ class BaseAWSLLM(SignsRequestsWithAWS): "RoleSessionName": aws_session_name, "WebIdentityToken": oidc_token, "DurationSeconds": 3600, - "Policy": json.dumps(build_web_identity_session_policy(), separators=(",", ":")), + "Policy": self._web_identity_session_policy_document(), } # Add ExternalId parameter if provided @@ -1787,6 +1712,24 @@ class BaseAWSLLM(SignsRequestsWithAWS): return request_headers_dict, request.body +def sign_aws_request( + method: str, + get_credentials: Callable[[], Credentials], + service_name: str, + aws_region_name: str | None, + url: str, + body: str | bytes, + headers: Mapping[str, str], +) -> AWSPreparedRequest: + ensure_optional_import("botocore") + from botocore.auth import SigV4Auth + from botocore.awsrequest import AWSRequest + + aws_request: Final = AWSRequest(method=method, url=url, data=body, headers=headers) + SigV4Auth(get_credentials(), service_name, aws_region_name).add_auth(aws_request) + return aws_request.prepare() + + def sign_aws_json_post( get_credentials: Callable[[], Credentials], service_name: str, @@ -1795,13 +1738,7 @@ def sign_aws_json_post( body: str, headers: Mapping[str, str], ) -> AWSPreparedRequest: - ensure_optional_import("botocore") - from botocore.auth import SigV4Auth - from botocore.awsrequest import AWSRequest - - aws_request: Final = AWSRequest(method="POST", url=url, data=body, headers=headers) - SigV4Auth(get_credentials(), service_name, aws_region_name).add_auth(aws_request) - return aws_request.prepare() + return sign_aws_request("POST", get_credentials, service_name, aws_region_name, url, body, headers) _SignParams = ParamSpec("_SignParams") diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index 84f40bec9ca..9daf67dd768 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -2,7 +2,7 @@ import os import re import time from collections.abc import Mapping -from typing import TYPE_CHECKING, Final, Literal, cast +from typing import TYPE_CHECKING, ClassVar, Final, Literal, cast from httpx import Headers, Response from pydantic import TypeAdapter, ValidationError @@ -14,6 +14,7 @@ from litellm.litellm_core_utils.cloud_storage_security import ( from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.bedrock.web_identity_session_policy import BEDROCK_BATCHES_SESSION_POLICY, WebIdentitySessionPolicy from litellm.types.llms.bedrock import ( BedrockCreateBatchRequest, BedrockCreateBatchResponse, @@ -102,9 +103,11 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): Config for Bedrock Batches - handles batch job creation and management for Bedrock """ + _web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = BEDROCK_BATCHES_SESSION_POLICY + def __init__(self): super().__init__() - self.common_utils = CommonBatchFilesUtils() + self.common_utils = CommonBatchFilesUtils(self) @property def custom_llm_provider(self) -> LlmProviders: diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 12c5031628b..5c82a8d6bc4 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -583,12 +583,14 @@ class AmazonConverseConfig(BaseConfig): custom_llm_provider="bedrock", llm_provider="bedrock_converse", ) - if mapped_thinking is None: + if mapped_thinking is None or mapped_thinking.get("type") == "disabled": optional_params.pop("thinking", None) optional_params.pop("output_config", None) - else: + if mapped_thinking is not None: optional_params["thinking"] = mapped_thinking - if AnthropicConfig.is_adaptive_thinking_model(model, "bedrock"): + if mapped_thinking.get("type") != "disabled" and AnthropicConfig.is_adaptive_thinking_model( + model, "bedrock" + ): mapped_effort = REASONING_EFFORT_TO_OUTPUT_CONFIG_EFFORT.get(reasoning_effort) if mapped_effort is None: AnthropicConfig.raise_invalid_reasoning_effort( diff --git a/litellm/llms/bedrock/chat/invoke_agent/transformation.py b/litellm/llms/bedrock/chat/invoke_agent/transformation.py index baef89912fa..41619f5d1b1 100644 --- a/litellm/llms/bedrock/chat/invoke_agent/transformation.py +++ b/litellm/llms/bedrock/chat/invoke_agent/transformation.py @@ -19,6 +19,7 @@ from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_a from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.common_utils import BedrockError +from litellm.types.llms.bedrock import AWS_AUTH_PARAM_KEYS from litellm.types.llms.bedrock_invoke_agents import ( InvokeAgentChunkPayload, InvokeAgentEvent, @@ -159,7 +160,7 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM): return { "inputText": query, "enableTrace": True, - **optional_params, + **{key: value for key, value in optional_params.items() if key not in AWS_AUTH_PARAM_KEYS}, } def _parse_aws_event_stream(self, raw_content: bytes) -> InvokeAgentEventList: diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 85e2c5fe1a1..6a05cb1d834 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -28,6 +28,8 @@ if TYPE_CHECKING: from litellm.types.llms.bedrock import BedrockCreateBatchRequest + from .base_aws_llm import BaseAWSLLM + import httpx from pydantic import ConfigDict, TypeAdapter, ValidationError @@ -2641,11 +2643,8 @@ class CommonBatchFilesUtils: Provides shared functionality to reduce code duplication between batches and files. """ - def __init__(self): - # Import here to avoid circular imports - from .base_aws_llm import BaseAWSLLM - - self._base_aws = BaseAWSLLM() + def __init__(self, aws: BaseAWSLLM) -> None: + self._base_aws = aws def get_bedrock_model_id_from_litellm_model(self, model: str) -> str: """ diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index f15a512077a..e28e39ef2d4 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -6,7 +6,7 @@ import copy import json import urllib.parse from collections.abc import Callable, Mapping -from typing import TYPE_CHECKING, Final, get_args, overload +from typing import TYPE_CHECKING, ClassVar, Final, get_args, overload import httpx from pydantic import ConfigDict, TypeAdapter @@ -39,6 +39,7 @@ from ..base_aws_llm import ( run_aws_signing, ) from ..common_utils import BedrockError, split_bedrock_region_path +from ..web_identity_session_policy import BEDROCK_EMBEDDING_SESSION_POLICY, WebIdentitySessionPolicy from .amazon_nova_transformation import AmazonNovaEmbeddingConfig from .amazon_titan_g1_transformation import AmazonTitanG1Config from .amazon_titan_multimodal_transformation import ( @@ -77,6 +78,8 @@ def _without_region_path(model: str) -> str: class BedrockEmbedding(BaseAWSLLM): + _web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = BEDROCK_EMBEDDING_SESSION_POLICY + @overload def _load_credentials( self, diff --git a/litellm/llms/bedrock/files/handler.py b/litellm/llms/bedrock/files/handler.py index 37f95ff7b9b..92839ee4752 100644 --- a/litellm/llms/bedrock/files/handler.py +++ b/litellm/llms/bedrock/files/handler.py @@ -1,6 +1,6 @@ import asyncio from collections.abc import Coroutine, Mapping -from typing import Final +from typing import ClassVar, Final import httpx @@ -10,6 +10,10 @@ from litellm.litellm_core_utils.cloud_storage_security import ( should_allow_legacy_cloud_file_ids, validate_managed_cloud_file_id, ) +from litellm.llms.bedrock.web_identity_session_policy import ( + S3_OBJECT_READ_SESSION_POLICY, + WebIdentitySessionPolicy, +) from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.llms.openai import ( FileContentRequest, @@ -27,6 +31,8 @@ class BedrockFilesHandler(BaseAWSLLM): stores batch output files. """ + _web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = S3_OBJECT_READ_SESSION_POLICY + def __init__(self): super().__init__() self.async_httpx_client = get_async_httpx_client( diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index e5f470fb87b..476d1395321 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -11,7 +11,7 @@ from datetime import datetime from email.utils import parsedate_to_datetime from itertools import chain from types import MappingProxyType -from typing import Any, Final, Literal, Protocol, TypeAlias, TypedDict +from typing import Any, ClassVar, Final, Literal, Protocol, TypeAlias, TypedDict from urllib.parse import quote, unquote, urlencode import httpx @@ -54,6 +54,7 @@ from litellm.llms.base_llm.files.transformation import ( BaseFilesConfig, LiteLLMLoggingObj, ) +from litellm.llms.bedrock.web_identity_session_policy import BEDROCK_FILES_SESSION_POLICY, WebIdentitySessionPolicy from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.bedrock import AwsAuthParams, BedrockBatchRecordKind @@ -591,6 +592,8 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): Config for Bedrock Files - handles S3 uploads for Bedrock batch processing """ + _web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = BEDROCK_FILES_SESSION_POLICY + def __init__(self): self.jsonl_transformation = BedrockJsonlFilesTransformation() super().__init__() diff --git a/litellm/llms/bedrock/model_listing.py b/litellm/llms/bedrock/model_listing.py index 409bf2c7b97..d64ca0cfad9 100644 --- a/litellm/llms/bedrock/model_listing.py +++ b/litellm/llms/bedrock/model_listing.py @@ -1,6 +1,6 @@ from collections.abc import Iterator, Mapping from types import MappingProxyType -from typing import Final, TypeAlias +from typing import ClassVar, Final, TypeAlias from urllib.parse import quote, urlencode import httpx @@ -9,6 +9,10 @@ from pydantic import BaseModel, ConfigDict, Field, TypeAdapter from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, bedrock_bearer_token from litellm.llms.bedrock.common_utils import BedrockError +from litellm.llms.bedrock.web_identity_session_policy import ( + BEDROCK_MODEL_LISTING_SESSION_POLICY, + WebIdentitySessionPolicy, +) from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.types.llms.bedrock import AwsAuthParams @@ -47,6 +51,8 @@ class ListInferenceProfilesResponse(BaseModel): class BedrockModelLister(BaseAWSLLM): + _web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = BEDROCK_MODEL_LISTING_SESSION_POLICY + def __init__(self, deployment: Mapping[str, object], client: HTTPHandler) -> None: super().__init__() self._aws_params: Final = dict(deployment) diff --git a/litellm/llms/bedrock/passthrough/transformation.py b/litellm/llms/bedrock/passthrough/transformation.py index edd30fff11b..4a2bb57ce5a 100644 --- a/litellm/llms/bedrock/passthrough/transformation.py +++ b/litellm/llms/bedrock/passthrough/transformation.py @@ -3,7 +3,7 @@ from collections.abc import Callable, Iterator, Mapping, Sequence from functools import partial from itertools import chain from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Optional, cast +from typing import TYPE_CHECKING, ClassVar, Final, Optional, cast import httpx from httpx import Response @@ -32,6 +32,7 @@ from ..request_metadata import ( merge_bedrock_invoke_headers, resolve_bedrock_request_metadata, ) +from ..web_identity_session_policy import BEDROCK_AGENT_RUNTIME_SESSION_POLICY, WebIdentitySessionPolicy if TYPE_CHECKING: from botocore.eventstream import EventStreamMessage @@ -379,6 +380,10 @@ class BedrockEventStreamCollector: return _build_logged_response(self._chunks.chunks(), litellm_logging_obj) +class BedrockAgentRuntimeAWS(BaseAWSLLM): + _web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = BEDROCK_AGENT_RUNTIME_SESSION_POLICY + + class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamDecoderBase, BasePassthroughConfig): def get_error_class( self, diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 9f4fb8cfecc..cb84277b7b7 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -13,7 +13,7 @@ from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Mappin from dataclasses import dataclass, field from enum import Enum, auto from types import MappingProxyType -from typing import TYPE_CHECKING, Final, NoReturn, Protocol, runtime_checkable +from typing import TYPE_CHECKING, ClassVar, Final, NoReturn, Protocol, runtime_checkable from pydantic import JsonValue, TypeAdapter @@ -54,6 +54,7 @@ from litellm.types.utils import StandardLoggingPayloadErrorInformation from ..base_aws_llm import BaseAWSLLM, run_aws_signing from ..common_utils import BedrockError +from ..web_identity_session_policy import BEDROCK_REALTIME_SESSION_POLICY, WebIdentitySessionPolicy from .transformation import BedrockRealtimeConfig, requested_modalities if TYPE_CHECKING: @@ -535,6 +536,8 @@ async def _close_bedrock_input( class BedrockRealtime(BaseAWSLLM): """Handler for Bedrock Nova Sonic realtime speech-to-speech API.""" + _web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = BEDROCK_REALTIME_SESSION_POLICY + def __init__( self, sdk_version_lookup: Callable[[], str | None] = _installed_sdk_version, diff --git a/litellm/llms/bedrock/web_identity_session_policy.py b/litellm/llms/bedrock/web_identity_session_policy.py new file mode 100644 index 00000000000..31187953dbc --- /dev/null +++ b/litellm/llms/bedrock/web_identity_session_policy.py @@ -0,0 +1,362 @@ +import json +from typing import Final, Literal, TypeAlias + +from typing_extensions import NotRequired, ReadOnly, TypedDict + +_SecureTransportBool: Final = TypedDict("_SecureTransportBool", {"aws:SecureTransport": ReadOnly[Literal["true"]]}) + +_CalledViaLastEquals: Final = TypedDict( + "_CalledViaLastEquals", + { + "aws:CalledViaLast": ReadOnly[ + Literal["aws-external-anthropic.amazonaws.com", "bedrock.amazonaws.com", "transcribe.amazonaws.com"] + ] + }, +) + +_PassedToServiceEquals: Final = TypedDict( + "_PassedToServiceEquals", + {"iam:PassedToService": ReadOnly[Literal["bedrock.amazonaws.com", "transcribe.amazonaws.com"]]}, +) + +_ServiceLinkedRoleEquals: Final = TypedDict( + "_ServiceLinkedRoleEquals", {"iam:AWSServiceName": ReadOnly[Literal["observability.aoss.amazonaws.com"]]} +) + +_KmsViaService: TypeAlias = Literal["s3.*", "sqs.*", "secretsmanager.*"] + +_ViaServiceLike: Final = TypedDict("_ViaServiceLike", {"kms:ViaService": ReadOnly[_KmsViaService]}) + +_StatementResource: TypeAlias = Literal["*", "arn:*:iam::*:role/litellm-bedrock-kb-*"] + + +class _SessionPolicyCondition(TypedDict): + Bool: ReadOnly[_SecureTransportBool] + StringEquals: NotRequired[ReadOnly[_CalledViaLastEquals | _PassedToServiceEquals | _ServiceLinkedRoleEquals]] + StringLike: NotRequired[ReadOnly[_ViaServiceLike]] + + +class _SessionPolicyStatement(TypedDict): + Sid: ReadOnly[str] + Effect: ReadOnly[Literal["Allow"]] + Action: ReadOnly[tuple[str, ...]] + Resource: ReadOnly[_StatementResource] + Condition: NotRequired[ReadOnly[_SessionPolicyCondition]] + + +class WebIdentitySessionPolicy(TypedDict): + Version: ReadOnly[Literal["2012-10-17"]] + Statement: ReadOnly[tuple[_SessionPolicyStatement, ...]] + + +_SECURE_TRANSPORT: Final = _SecureTransportBool({"aws:SecureTransport": "true"}) + +_SECURE_TRANSPORT_ONLY: Final = _SessionPolicyCondition(Bool=_SECURE_TRANSPORT) + + +def _allow( + sid: str, + actions: tuple[str, ...], + condition: _SessionPolicyCondition = _SECURE_TRANSPORT_ONLY, + resource: _StatementResource = "*", +) -> _SessionPolicyStatement: + return _SessionPolicyStatement(Sid=sid, Effect="Allow", Action=actions, Resource=resource, Condition=condition) + + +def _allow_over_any_transport(sid: str, actions: tuple[str, ...]) -> _SessionPolicyStatement: + return _SessionPolicyStatement(Sid=sid, Effect="Allow", Action=actions, Resource="*") + + +def _policy(*statements: _SessionPolicyStatement) -> WebIdentitySessionPolicy: + return WebIdentitySessionPolicy(Version="2012-10-17", Statement=statements) + + +def session_policy_document(policy: WebIdentitySessionPolicy) -> str: + return json.dumps(policy, separators=(",", ":")) + + +def _encryption_through(service: _KmsViaService, actions: tuple[str, ...]) -> _SessionPolicyStatement: + return _allow( + "KmsLiteLLM", + actions, + _SessionPolicyCondition(Bool=_SECURE_TRANSPORT, StringLike=_ViaServiceLike({"kms:ViaService": service})), + ) + + +_OBJECT_ENCRYPTION: Final = ("kms:GenerateDataKey", "kms:Decrypt") + +BEDROCK_SESSION_POLICY: Final = _policy( + _allow( + "BedrockLiteLLM", + ( + "bedrock:InvokeModel", + "bedrock:InvokeModelWithResponseStream", + "bedrock:CountTokens", + "bedrock:Rerank", + "bedrock:Retrieve", + "bedrock:ListKnowledgeBases", + "bedrock:InvokeAgent", + "bedrock:ApplyGuardrail", + "bedrock:GetGuardrail", + "bedrock:ListGuardrails", + ), + ), + _allow( + "BedrockAgentCoreLiteLLM", + ( + "bedrock-agentcore:InvokeAgentRuntime", + "bedrock-agentcore:InvokeAgentRuntimeForUser", + "bedrock-agentcore:InvokeGateway", + ), + ), + _allow( + "ClaudePlatformLiteLLM", + ( + "aws-external-anthropic:CreateInference", + "aws-external-anthropic:CreateBatchInference", + "aws-external-anthropic:CancelBatchInference", + "aws-external-anthropic:DeleteBatchInference", + "aws-external-anthropic:CountTokens", + "aws-external-anthropic:Get*", + "aws-external-anthropic:List*", + ), + ), + _allow( + "ClaudePlatformOutboundIdentityLiteLLM", + ("sts:GetWebIdentityToken", "sts:TagGetWebIdentityToken"), + _SessionPolicyCondition( + Bool=_SECURE_TRANSPORT, + StringEquals=_CalledViaLastEquals({"aws:CalledViaLast": "aws-external-anthropic.amazonaws.com"}), + ), + ), + _allow("BedrockMantleLiteLLM", ("bedrock-mantle:CreateInference", "bedrock-mantle:CountTokens")), +) + +BEDROCK_REALTIME_SESSION_POLICY: Final = _policy( + _allow_over_any_transport("BedrockRealtimeLiteLLM", ("bedrock:InvokeModel", "bedrock:ApplyGuardrail")), +) + +BEDROCK_EMBEDDING_SESSION_POLICY: Final = _policy( + _allow("BedrockEmbeddingLiteLLM", ("bedrock:InvokeModel", "bedrock:ApplyGuardrail", "bedrock:GetAsyncInvoke")), + _allow( + "BedrockS3LiteLLM", + ("s3:GetObject", "s3:PutObject"), + _SessionPolicyCondition( + Bool=_SECURE_TRANSPORT, + StringEquals=_CalledViaLastEquals({"aws:CalledViaLast": "bedrock.amazonaws.com"}), + ), + ), + _encryption_through("s3.*", _OBJECT_ENCRYPTION), +) + +BEDROCK_AGENT_RUNTIME_SESSION_POLICY: Final = _policy( + _allow( + "BedrockAgentRuntimeLiteLLM", + ( + "bedrock:InvokeAgent", + "bedrock:InvokeInlineAgent", + "bedrock:InvokeFlow", + "bedrock:StartFlowExecution", + "bedrock:StopFlowExecution", + "bedrock:Retrieve", + "bedrock:RetrieveAndGenerate", + "bedrock:GenerateQuery", + "bedrock:GetKnowledgeBase", + "bedrock:GetDocumentContent", + "bedrock:CheckIngestedDocumentAcl", + "bedrock:GetIngestedDocumentAcl", + "bedrock:Rerank", + "bedrock:OptimizePrompt", + "bedrock:InvokeModel", + "bedrock:ApplyGuardrail", + "bedrock:GetInferenceProfile", + ), + ), + _allow("SqlWorkbenchLiteLLM", ("sqlworkbench:GetSqlRecommendations",)), +) + +BEDROCK_MODEL_LISTING_SESSION_POLICY: Final = _policy( + _allow("BedrockModelListingLiteLLM", ("bedrock:ListFoundationModels", "bedrock:ListInferenceProfiles")), +) + +BEDROCK_BATCHES_SESSION_POLICY: Final = _policy( + _allow( + "BedrockBatchLiteLLM", + ( + "bedrock:CreateModelInvocationJob", + "bedrock:GetModelInvocationJob", + "bedrock:StopModelInvocationJob", + "bedrock:TagResource", + ), + ), + _allow( + "BedrockBatchServiceRoleLiteLLM", + ("iam:PassRole",), + _SessionPolicyCondition( + Bool=_SECURE_TRANSPORT, + StringEquals=_PassedToServiceEquals({"iam:PassedToService": "bedrock.amazonaws.com"}), + ), + ), +) + +BEDROCK_FILES_SESSION_POLICY: Final = _policy( + _allow("S3ObjectsLiteLLM", ("s3:PutObject", "s3:GetObject", "s3:DeleteObject", "s3:ListBucket")), + _encryption_through("s3.*", _OBJECT_ENCRYPTION), +) + +S3_OBJECT_READ_SESSION_POLICY: Final = _policy( + _allow("S3ObjectsLiteLLM", ("s3:GetObject",)), + _encryption_through("s3.*", ("kms:Decrypt",)), +) + +BEDROCK_KNOWLEDGE_BASE_INGESTION_SESSION_POLICY: Final = _policy( + _allow( + "BedrockKnowledgeBaseLiteLLM", + ( + "bedrock:ListDataSources", + "bedrock:GetDataSource", + "bedrock:CreateKnowledgeBase", + "bedrock:GetKnowledgeBase", + "bedrock:CreateDataSource", + "bedrock:StartIngestionJob", + "bedrock:GetIngestionJob", + ), + ), + _allow("S3ObjectsLiteLLM", ("s3:CreateBucket", "s3:PutObject")), + _encryption_through("s3.*", _OBJECT_ENCRYPTION), + _allow( + "OpenSearchServerlessLiteLLM", + ( + "aoss:CreateSecurityPolicy", + "aoss:CreateAccessPolicy", + "aoss:CreateCollection", + "aoss:BatchGetCollection", + "aoss:APIAccessAll", + ), + ), + _allow( + "OpenSearchServerlessServiceLinkedRoleLiteLLM", + ("iam:CreateServiceLinkedRole",), + _SessionPolicyCondition( + Bool=_SECURE_TRANSPORT, + StringEquals=_ServiceLinkedRoleEquals({"iam:AWSServiceName": "observability.aoss.amazonaws.com"}), + ), + ), + _allow( + "KnowledgeBaseRoleLiteLLM", + ("iam:CreateRole", "iam:PutRolePolicy"), + resource="arn:*:iam::*:role/litellm-bedrock-kb-*", + ), + _allow( + "KnowledgeBaseServiceRoleLiteLLM", + ("iam:PassRole",), + _SessionPolicyCondition( + Bool=_SECURE_TRANSPORT, + StringEquals=_PassedToServiceEquals({"iam:PassedToService": "bedrock.amazonaws.com"}), + ), + resource="arn:*:iam::*:role/litellm-bedrock-kb-*", + ), +) + +SAGEMAKER_SESSION_POLICY: Final = _policy( + _allow("SageMakerLiteLLM", ("sagemaker:InvokeEndpoint", "sagemaker:InvokeEndpointWithResponseStream")), +) + +POLLY_SESSION_POLICY: Final = _policy(_allow("PollyLiteLLM", ("polly:SynthesizeSpeech",))) + +OPENSEARCH_SESSION_POLICY: Final = _policy(_allow("OpenSearchLiteLLM", ("es:ESHttpPost", "aoss:APIAccessAll"))) + +OPENSEARCH_INDEX_SESSION_POLICY: Final = _policy( + _allow( + "OpenSearchIndexLiteLLM", + ( + "es:ESHttpGet", + "es:ESHttpHead", + "es:ESHttpPost", + "es:ESHttpPut", + "es:ESHttpDelete", + "es:ESHttpPatch", + "aoss:APIAccessAll", + ), + ), +) + +S3_VECTORS_QUERY_SESSION_POLICY: Final = _policy( + _allow("S3VectorsLiteLLM", ("s3vectors:QueryVectors", "s3vectors:GetVectors")), +) + +S3_VECTORS_INGESTION_SESSION_POLICY: Final = _policy( + _allow( + "S3VectorsLiteLLM", + ( + "s3vectors:GetVectorBucket", + "s3vectors:CreateVectorBucket", + "s3vectors:GetIndex", + "s3vectors:CreateIndex", + "s3vectors:PutVectors", + ), + ), +) + +S3_LOGGING_SESSION_POLICY: Final = _policy( + _allow("S3ObjectsLiteLLM", ("s3:PutObject", "s3:GetObject")), + _encryption_through("s3.*", _OBJECT_ENCRYPTION), +) + +SQS_LOGGING_SESSION_POLICY: Final = _policy( + _allow("SqsLiteLLM", ("sqs:SendMessage",)), + _encryption_through("sqs.*", _OBJECT_ENCRYPTION), +) + +SECRETS_MANAGER_SESSION_POLICY: Final = _policy( + _allow( + "SecretsManagerLiteLLM", + ( + "secretsmanager:GetSecretValue", + "secretsmanager:DescribeSecret", + "secretsmanager:CreateSecret", + "secretsmanager:PutSecretValue", + "secretsmanager:UpdateSecret", + "secretsmanager:RestoreSecret", + "secretsmanager:TagResource", + "secretsmanager:ReplicateSecretToRegions", + "secretsmanager:DeleteSecret", + ), + ), + _encryption_through("secretsmanager.*", _OBJECT_ENCRYPTION), +) + +COMPREHEND_MEDICAL_SESSION_POLICY: Final = _policy( + _allow( + "ComprehendMedicalLiteLLM", + ( + "comprehendmedical:DetectEntitiesV2", + "comprehendmedical:DetectPHI", + "comprehendmedical:InferICD10CM", + "comprehendmedical:InferRxNorm", + "comprehendmedical:InferSNOMEDCT", + ), + ), +) + +TRANSCRIBE_SESSION_POLICY: Final = _policy( + _allow("TranscribeLiteLLM", ("transcribe:*",)), + _allow("TranscribeMediaLiteLLM", ("s3:GetObject",)), + _allow( + "TranscribeOutputLiteLLM", + ("s3:PutObject",), + _SessionPolicyCondition( + Bool=_SECURE_TRANSPORT, + StringEquals=_CalledViaLastEquals({"aws:CalledViaLast": "transcribe.amazonaws.com"}), + ), + ), + _allow( + "TranscribeDataAccessRoleLiteLLM", + ("iam:PassRole",), + _SessionPolicyCondition( + Bool=_SECURE_TRANSPORT, + StringEquals=_PassedToServiceEquals({"iam:PassedToService": "transcribe.amazonaws.com"}), + ), + ), + _encryption_through("s3.*", _OBJECT_ENCRYPTION), +) diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 3a6fa5f51da..0cfdbc74227 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -44,6 +44,8 @@ from litellm.litellm_core_utils.logging_utils import track_llm_api_timing from litellm.litellm_core_utils.request_timeout_resolver import ( get_configured_request_timeout, ) +from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine, default_token_exchange_engine +from litellm.llms.base_llm.auth.types import MintedToken from litellm.llms.custom_httpx.aiohttp_pool_diagnostics import build_pool_wait_trace_config from litellm.llms.custom_httpx.http2_connection_spread import AsyncHTTP2ConnectionSpread from litellm.types.llms.custom_http import * @@ -586,6 +588,28 @@ def _raise_masked_sync_error(e: httpx.HTTPStatusError, stream: bool) -> NoReturn raise MaskedHTTPStatusError(e, message=_text, text=_text) from None +def _same_origin(first: httpx.URL, second: httpx.URL) -> bool: + return (first.scheme, first.host, first.port) == (second.scheme, second.host, second.port) + + +async def _with_current_bearer( + headers: Mapping[str, str] | None, url: str, engine: JwtBearerTokenExchangeEngine +) -> Mapping[str, str] | None: + auth_key: Final = next((key for key in headers or () if key.lower() == "authorization"), None) + if headers is None or auth_key is None: + return headers + auth_value: Final = headers[auth_key] + if not auth_value.startswith("Bearer "): + return headers + spec: Final = engine.spec_for_access_token(auth_value.removeprefix("Bearer ")) + if spec is None or not _same_origin(httpx.URL(spec.token_url), httpx.URL(url)): + return headers + current: Final = await engine.aget_token(spec) + if not isinstance(current, MintedToken): + return headers + return MappingProxyType({**headers, auth_key: f"Bearer {current.access_token.get_secret_value()}"}) + + async def _raise_masked_async_error(e: httpx.HTTPStatusError, stream: bool) -> NoReturn: """Raise a MaskedHTTPStatusError for async HTTP handlers.""" if stream: @@ -639,7 +663,7 @@ class HTTPResponseLimitError(ValueError): class MaskedHTTPStatusError(httpx.HTTPStatusError): - def __init__(self, original_error, message: str | None = None, text: str | None = None): + def __init__(self, original_error: httpx.HTTPStatusError, message: str | None = None, text: str | None = None): # Create a new error with the masked URL masked_url: Final = mask_sensitive_info(str(original_error.request.url)) # Mask the original exception message too (it contains the full URL) @@ -702,12 +726,14 @@ class AsyncHTTPHandler: transport: httpx.AsyncBaseTransport | None = None, follow_redirects: bool = True, http2: bool = False, + token_exchange_engine: JwtBearerTokenExchangeEngine = default_token_exchange_engine, ): self.timeout = timeout self.event_hooks = event_hooks self.ssl_verify = ssl_verify self.shared_session = shared_session self.transport = transport + self.token_exchange_engine = token_exchange_engine self.follow_redirects = follow_redirects self.http2 = http2 self._owns_client = True @@ -1037,7 +1063,7 @@ class AsyncHTTPHandler: data=retry_data, json=json, params=params, - headers=headers, + headers=await _with_current_bearer(headers, url, self.token_exchange_engine), timeout=timeout, content=retry_content, ) @@ -1193,7 +1219,7 @@ class AsyncHTTPHandler: data=request_data, json=json, params=params, - headers=headers, + headers=await _with_current_bearer(headers, url, self.token_exchange_engine), content=request_content, ) response: Final = await client.send(req, stream=stream) diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index a173cdb50f2..7917a74247e 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -629,12 +629,14 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): custom_llm_provider="databricks", llm_provider="databricks", ) - if mapped_thinking is None: + if mapped_thinking is None or mapped_thinking.get("type") == "disabled": optional_params.pop("thinking", None) optional_params.pop("output_config", None) - else: + if mapped_thinking is not None: optional_params["thinking"] = mapped_thinking - if AnthropicConfig._is_adaptive_thinking_model(model, "databricks"): + if mapped_thinking.get("type") != "disabled" and AnthropicConfig._is_adaptive_thinking_model( + model, "databricks" + ): mapped_effort: str | None = None if isinstance(reasoning_effort_value, str): mapped_effort = REASONING_EFFORT_TO_OUTPUT_CONFIG_EFFORT.get(reasoning_effort_value) diff --git a/litellm/llms/opensearch/vector_stores/index_passthrough.py b/litellm/llms/opensearch/vector_stores/index_passthrough.py new file mode 100644 index 00000000000..9432af4c1b3 --- /dev/null +++ b/litellm/llms/opensearch/vector_stores/index_passthrough.py @@ -0,0 +1,379 @@ +import json +from collections.abc import Iterable, Iterator, Mapping, Sequence +from dataclasses import dataclass +from itertools import accumulate, chain, repeat, takewhile +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Literal, TypeAlias +from urllib.parse import quote + +from pydantic import TypeAdapter + +from litellm.exceptions import AuthenticationError, BadRequestError, InternalServerError +from litellm.llms.bedrock.base_aws_llm import run_aws_signing + +if TYPE_CHECKING: + from litellm.llms.opensearch.vector_stores.transformation import OpenSearchVectorStoreConfig + +IndexHttpMethod: TypeAlias = Literal["GET", "HEAD", "POST", "PUT", "DELETE", "PATCH"] +IndexPermission: TypeAlias = Literal["read", "write"] +BodyKind: TypeAlias = Literal["query", "mget", "bulk", "update", "unscanned"] + +DEFAULT_CONTENT_TYPE: Final = "application/json" +NDJSON_CONTENT_TYPE: Final = "application/x-ndjson" +ACCEPTED_MEDIA_TYPES: Final = frozenset({DEFAULT_CONTENT_TYPE, NDJSON_CONTENT_TYPE}) +_ACCEPTED_MEDIA_TYPE_PARAMETERS: Final = frozenset({"", "charset=utf-8"}) +_INDEX_KEYS: Final = frozenset({"index", "_index", "indices"}) +_OUT_OF_INDEX_KEYS: Final = _INDEX_KEYS | {"indexed_shape", "wrapper", "search_pipeline"} +_BULK_METADATA_KEYS: Final = _INDEX_KEYS | {"pipeline"} +_OUT_OF_INDEX_QUERY_PARAMETERS: Final = frozenset({"pipeline", "source", "source_content_type"}) +_BULK_ACTIONS: Final = frozenset({"index", "create", "update", "delete"}) +_DOCUMENT_READING_KEYS: Final = frozenset({"_source", "script"}) +_JSON_BLANK: Final = b" \t\r\n" +_PATH_HELP: Final = "Use /opensearch//, for example /opensearch//_search." + + +@dataclass(frozen=True, slots=True) +class IndexRefusal: + status_code: int + message: str + + +@dataclass(frozen=True, slots=True) +class IndexPath: + index_name: str + api: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class IndexOperation: + permission: IndexPermission + methods: tuple[IndexHttpMethod, ...] + api: str + takes_id: bool + body: BodyKind + + def matches(self, method: str, api: tuple[str, ...]) -> bool: + named: Final = (self.api,) if self.api else () + return method in self.methods and api[: len(named)] == named and len(api) == len(named) + self.takes_id + + +@dataclass(frozen=True, slots=True) +class SignedIndexRequest: + url: str + media_type: str + headers: Mapping[str, str] + body: bytes + + +INDEX_OPERATIONS: Final = ( + IndexOperation("read", ("GET", "HEAD"), "", False, "unscanned"), + IndexOperation("write", ("POST",), "_doc", False, "unscanned"), + IndexOperation("write", ("PUT", "POST", "DELETE"), "_doc", True, "unscanned"), + IndexOperation("write", ("PUT", "POST"), "_create", True, "unscanned"), + IndexOperation("write", ("POST",), "_update", True, "update"), + IndexOperation("write", ("POST", "PUT"), "_bulk", False, "bulk"), + IndexOperation("write", ("POST",), "_delete_by_query", False, "query"), + IndexOperation("write", ("POST",), "_update_by_query", False, "query"), + IndexOperation("read", ("GET", "POST"), "_search", False, "query"), + IndexOperation("read", ("GET", "POST"), "_count", False, "query"), + IndexOperation("read", ("GET", "HEAD"), "_doc", True, "unscanned"), + IndexOperation("read", ("GET", "HEAD"), "_source", True, "unscanned"), + IndexOperation("read", ("GET", "POST"), "_mget", False, "mget"), + IndexOperation("read", ("GET", "POST"), "_explain", True, "query"), + IndexOperation("read", ("GET",), "_mapping", False, "unscanned"), +) + + +def index_operation(method: str, api: tuple[str, ...]) -> IndexOperation | None: + return next((operation for operation in INDEX_OPERATIONS if operation.matches(method, api)), None) + + +def parse_index_path(endpoint: str) -> IndexPath | IndexRefusal: + segments: Final = tuple(endpoint.split("/")) + if any(segment in ("", ".", "..") for segment in segments): + return IndexRefusal(400, f"The OpenSearch path '{endpoint}' has an empty, '.' or '..' segment. {_PATH_HELP}") + return IndexPath(index_name=segments[0], api=segments[1:]) + + +@dataclass(frozen=True, slots=True) +class _JsonObject: + pairs: tuple[tuple[str, object], ...] + + +@dataclass(frozen=True, slots=True) +class _Json: + value: object + + +@dataclass(frozen=True, slots=True) +class _BulkAction: + name: str + metadata: _JsonObject + has_source: bool + + +@dataclass(frozen=True, slots=True) +class _IndexScope: + name: str + provider_name: str + can_read: bool + + +def _json_object(pairs: Sequence[tuple[str, object]]) -> _JsonObject: + return _JsonObject(tuple(pairs)) + + +_DECODER: Final = json.JSONDecoder(object_pairs_hook=_json_object) +_JSON_ARRAY: Final = TypeAdapter(list[object]) +_JSON_VALUE: Final = TypeAdapter(object) + + +def _parse(raw: bytes) -> _Json | None: + try: + return _Json(_JSON_VALUE.validate_python(_DECODER.decode(raw.decode("utf-8")))) + except ValueError: + return None + + +def _own_keys(value: object) -> tuple[str, ...]: + return tuple(key for key, _ in value.pairs) if isinstance(value, _JsonObject) else () + + +def _children(value: object) -> Sequence[object]: + if isinstance(value, _JsonObject): + return tuple(child for _, child in value.pairs) + if isinstance(value, list): + return _JSON_ARRAY.validate_python(value) + return () + + +def _next_level(level: Sequence[object], _: None) -> Sequence[object]: + return tuple(chain.from_iterable(map(_children, level))) + + +def _forbidden_key(value: object, forbidden: frozenset[str]) -> str | None: + levels: Final = takewhile(bool, accumulate(repeat(None), _next_level, initial=(value,))) + keys: Final[Iterable[str]] = chain.from_iterable(map(_own_keys, chain.from_iterable(levels))) + return next((key for key in keys if key in forbidden), None) + + +def _plain(value: _JsonObject) -> Mapping[str, object]: + return dict(value.pairs) + + +def _object_json(pairs: Sequence[tuple[str, object]]) -> str: + fields: Final = ",".join( + f"{json.dumps(key)}:{json.dumps(child, default=_plain, separators=(',', ':'))}" for key, child in pairs + ) + return f"{{{fields}}}" + + +def _unparseable(api: str) -> IndexRefusal: + return IndexRefusal(400, f"LiteLLM could not parse the {api} request body as JSON, so it did not send it.") + + +def _index_reference(value: object, forbidden: frozenset[str]) -> IndexRefusal | None: + key: Final = _forbidden_key(value, forbidden) + if key is None: + return None + return IndexRefusal( + 400, + f"The request body uses '{key}', which can reach beyond the OpenSearch index in the URL path. " + f"Remove '{key}' from the body.", + ) + + +def _read_needed(cause: str, index_name: str) -> IndexRefusal: + return IndexRefusal( + 403, + f"{cause}, so the key needs read on index '{index_name}' as well as write. " + "Ask your administrator to add read to the key or its team.", + ) + + +def _document_reading_payload(api: str, payload: bytes, index: _IndexScope) -> IndexRefusal | None: + if index.can_read: + return None + parsed: Final = _parse(payload) + if parsed is None: + return _unparseable(api) + key: Final = next((key for key in _own_keys(parsed.value) if key in _DOCUMENT_READING_KEYS), None) + return None if key is None else _read_needed(f"'{key}' in an {api} body reads stored documents", index.name) + + +def _line_refusal(api: str, line: bytes, forbidden: frozenset[str]) -> IndexRefusal | None: + parsed: Final = _parse(line) + if parsed is None: + return _unparseable(api) + return _index_reference(parsed.value, forbidden) + + +def _whole_body_refusal(api: str, body: bytes) -> IndexRefusal | None: + if not body.strip(_JSON_BLANK): + return None + return _line_refusal(api, body, _OUT_OF_INDEX_KEYS) + + +def _pinned_index(metadata: _JsonObject, index: _IndexScope) -> _JsonObject | IndexRefusal | None: + unnamed: Final = _JsonObject(tuple(pair for pair in metadata.pairs if pair[0] != "_index")) + refusal: Final = _index_reference(unnamed, _BULK_METADATA_KEYS) + if refusal is not None: + return refusal + named: Final = tuple(value for name, value in metadata.pairs if name == "_index") + if not named: + return None + if any(value != index.name for value in named): + return IndexRefusal( + 400, + f"A _bulk action line sets '_index' to another index than '{index.name}' in the URL path. " + f"Remove '_index' or set it to '{index.name}'.", + ) + return _JsonObject( + tuple((name, index.provider_name if name == "_index" else value) for name, value in metadata.pairs) + ) + + +def _bulk_action(line: bytes) -> _BulkAction | None: + parsed: Final = _parse(line) + if parsed is None: + return None + match parsed.value: + case _JsonObject(pairs=((name, _JsonObject() as metadata),)) if name in _BULK_ACTIONS: + return _BulkAction(name=name, metadata=metadata, has_source=name != "delete") + case _: + return None + + +def _bulk_action_line(action: _BulkAction | None, line: bytes, index: _IndexScope) -> bytes | IndexRefusal: + if action is None: + return IndexRefusal( + 400, + "A _bulk action line must be one JSON object naming index, create, update, or delete. " + "LiteLLM did not send the request.", + ) + if not index.can_read and "_source" in _own_keys(action.metadata): + return _read_needed("'_source' in a _bulk action line returns stored documents", index.name) + pinned: Final = _pinned_index(action.metadata, index) + if pinned is None: + return line + if isinstance(pinned, IndexRefusal): + return pinned + return f"{{{json.dumps(action.name)}:{_object_json(pinned.pairs)}}}".encode() + + +def _bulk_entry(line: bytes, lines: Iterator[bytes], index: _IndexScope) -> tuple[bytes | IndexRefusal, ...]: + if not line.strip(_JSON_BLANK): + return (line,) + action: Final = _bulk_action(line) + rendered: Final = _bulk_action_line(action, line, index) + if isinstance(rendered, IndexRefusal) or action is None or not action.has_source: + return (rendered,) + source: Final = next(lines, None) + if source is None: + return (rendered,) + if action.name != "update": + return (rendered, source) + return (rendered, _document_reading_payload("_bulk update", source, index) or source) + + +def _bulk_lines(lines: Iterator[bytes], index: _IndexScope) -> Iterator[bytes | IndexRefusal]: + for line in lines: + yield from (entry := _bulk_entry(line, lines, index)) + if any(isinstance(rendered, IndexRefusal) for rendered in entry): + return + + +def _rendered_body(lines: Iterator[bytes | IndexRefusal]) -> bytes | IndexRefusal: + rendered: Final = tuple(lines) + refusal: Final = next((line for line in rendered if isinstance(line, IndexRefusal)), None) + if refusal is not None: + return refusal + return b"\n".join(line for line in rendered if isinstance(line, bytes)) + + +def _checked_body(operation: IndexOperation, body: bytes, index: _IndexScope) -> bytes | IndexRefusal: + match operation.body: + case "query" if operation.permission == "write" and not index.can_read: + return _read_needed(f"{operation.api} searches stored documents", index.name) + case "query" | "mget": + return _whole_body_refusal(operation.api, body) or body + case "update": + return _document_reading_payload(operation.api, body, index) or body + case "bulk": + return _rendered_body(_bulk_lines(iter(body.split(b"\n")), index)) + case "unscanned": + return body + + +def accepted_media_type(content_type: str | None) -> str | None: + media, _, parameters = (content_type or DEFAULT_CONTENT_TYPE).partition(";") + normalized: Final = media.strip().lower() + if normalized in ACCEPTED_MEDIA_TYPES and parameters.strip().lower() in _ACCEPTED_MEDIA_TYPE_PARAMETERS: + return normalized + return None + + +def _checked_request_body( + operation: IndexOperation | None, + query: Sequence[tuple[str, str]], + body: bytes, + index: _IndexScope, +) -> bytes | IndexRefusal: + if operation is None: + return body + parameter: Final = next((key for key, _ in query if key in _OUT_OF_INDEX_QUERY_PARAMETERS), None) + if parameter is not None: + return IndexRefusal( + 400, f"The {parameter} query parameter can reach beyond the OpenSearch index in the URL path. Remove it." + ) + source_parameter: Final = next((key for key, _ in query if key.startswith("_source")), None) + if operation.permission == "write" and not index.can_read and source_parameter is not None: + return _read_needed(f"The {source_parameter} query parameter returns stored documents", index.name) + try: + return _checked_body(operation, body, index) + except RecursionError: + return IndexRefusal(400, f"The {operation.api} request body is nested too deeply for LiteLLM to check.") + + +def upstream_url(api_base: str, provider_index: str, api: tuple[str, ...], query: Sequence[tuple[str, str]]) -> str: + path: Final = "/".join(quote(segment, safe="") for segment in (provider_index, *api)) + query_string: Final = "&".join(f"{quote(key, safe='')}={quote(value, safe='')}" for key, value in query) + return f"{api_base}/{path}?{query_string}" if query_string else f"{api_base}/{path}" + + +async def prepare_index_request( + config: "OpenSearchVectorStoreConfig", + method: str, + path: IndexPath, + operation: IndexOperation | None, + provider_index: str, + content_type: str | None, + query: Sequence[tuple[str, str]], + body: bytes, + store_params: Mapping[str, object], + can_read: bool, +) -> SignedIndexRequest | IndexRefusal: + media: Final = accepted_media_type(content_type) + if media is None: + return IndexRefusal( + 400, "The OpenSearch pass-through takes application/json or application/x-ndjson bodies in UTF-8." + ) + checked: Final = _checked_request_body( + operation, query, body, _IndexScope(name=path.index_name, provider_name=provider_index, can_read=can_read) + ) + if isinstance(checked, IndexRefusal): + return checked + api_base: Final = store_params.get("api_base") + try: + url: Final = upstream_url( + config.get_complete_url(api_base if isinstance(api_base, str) else None, store_params), + provider_index, + path.api, + query, + ) + signed: Final = await run_aws_signing(config.sign_http_request, method, url, checked, media, store_params) + except (AuthenticationError, BadRequestError, InternalServerError) as error: + return IndexRefusal(error.status_code, error.message) + return SignedIndexRequest( + url=url, media_type=media, headers=MappingProxyType(dict(signed.headers.items())), body=checked + ) diff --git a/litellm/llms/opensearch/vector_stores/transformation.py b/litellm/llms/opensearch/vector_stores/transformation.py index 2d862bc8ea6..f7ba28e7f3f 100644 --- a/litellm/llms/opensearch/vector_stores/transformation.py +++ b/litellm/llms/opensearch/vector_stores/transformation.py @@ -3,14 +3,14 @@ import json import re from collections.abc import Mapping, Sequence from functools import partial -from typing import TYPE_CHECKING, Final, Literal, NoReturn +from typing import TYPE_CHECKING, ClassVar, Final, Literal, NoReturn from urllib.parse import quote, urlsplit import httpx from pydantic import ConfigDict, Field, TypeAdapter, ValidationError from litellm._logging import verbose_logger -from litellm.exceptions import AuthenticationError, BadRequestError +from litellm.exceptions import AuthenticationError, BadRequestError, InternalServerError from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.optional_imports import ensure_optional_import from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -18,7 +18,12 @@ from litellm.llms.base_llm.vector_store.transformation import ( BaseQueryEmbeddingVectorStoreConfig, VectorStoreEmbeddingExecutor, ) -from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, SignsRequestsWithAWS, sign_aws_json_post +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, SignsRequestsWithAWS, sign_aws_request +from litellm.llms.bedrock.web_identity_session_policy import ( + OPENSEARCH_INDEX_SESSION_POLICY, + OPENSEARCH_SESSION_POLICY, + WebIdentitySessionPolicy, +) from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.bedrock import AwsAuthParams from litellm.types.router import GenericLiteLLMParams @@ -34,6 +39,7 @@ from litellm.types.vector_stores import ( ) if TYPE_CHECKING: + from botocore.awsrequest import AWSPreparedRequest from botocore.credentials import Credentials from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -80,6 +86,16 @@ def _unauthenticated(reason: str) -> AuthenticationError: ) +def _unavailable_credentials(error: Exception) -> InternalServerError: + verbose_logger.warning("OpenSearch vector store AWS credentials could not be loaded or refreshed: %s", error) + return InternalServerError( + message="LiteLLM could not load or refresh the AWS credentials that sign the OpenSearch request, " + "see the LiteLLM logs.", + llm_provider="opensearch", + model=None, + ) + + def _usable_credentials(credentials: object) -> "Credentials": from botocore.credentials import Credentials @@ -208,9 +224,17 @@ def _search_result(hit: _Hit, params: _OpenSearchParams) -> VectorStoreSearchRes ) +class _OpenSearchAWS(BaseAWSLLM): + _web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = OPENSEARCH_SESSION_POLICY + + +class OpenSearchIndexAWS(BaseAWSLLM): + _web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = OPENSEARCH_INDEX_SESSION_POLICY + + class OpenSearchVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, SignsRequestsWithAWS): def __init__(self, aws: BaseAWSLLM | None = None) -> None: - self._aws: Final = aws if aws is not None else BaseAWSLLM() + self._aws: Final = aws if aws is not None else _OpenSearchAWS() def get_auth_credentials(self, litellm_params: Mapping[str, object]) -> BaseVectorStoreAuthCredentials: return BaseVectorStoreAuthCredentials() @@ -238,6 +262,7 @@ class OpenSearchVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, SignsRequ if bool(auth.aws_access_key_id) != bool(auth.aws_secret_access_key): raise _unauthenticated("aws_access_key_id and aws_secret_access_key must be set together") ensure_optional_import("botocore") + from botocore.credentials import Credentials from botocore.exceptions import ( CredentialRetrievalError, NoCredentialsError, @@ -246,12 +271,21 @@ class OpenSearchVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, SignsRequ ) try: - return _usable_credentials(self._aws.resolve_credentials(auth, region)) + credentials: Final = _usable_credentials(self._aws.resolve_credentials(auth, region)) except CredentialRetrievalError as error: verbose_logger.warning("OpenSearch vector store AWS credential source failed: %s", error) raise _unauthenticated("the AWS credential source failed, see the LiteLLM logs") from None + except OSError as error: + raise _unavailable_credentials(error) from None except (NoCredentialsError, PartialCredentialsError, ProfileNotFound) as error: raise _unauthenticated(str(error)) from None + try: + frozen: Final = credentials.get_frozen_credentials() + except (CredentialRetrievalError, OSError, RuntimeError) as error: + raise _unavailable_credentials(error) from None + if frozen.access_key is None or frozen.secret_key is None: + raise _unauthenticated("none were found") + return Credentials(frozen.access_key, frozen.secret_key, frozen.token, account_id=frozen.account_id) def get_complete_url(self, api_base: str | None, litellm_params: Mapping[str, object]) -> str: parsed: Final = urlsplit(api_base or "") @@ -363,25 +397,28 @@ class OpenSearchVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, SignsRequ api_base: str, api_key: str | None = None, ) -> tuple[dict[str, str], bytes]: # mutable-ok: the shared HTTP handler requires writable headers - region: Final = self._signing_region(optional_params, api_base) - body: Final = json.dumps(request_data) - signed: Final = sign_aws_json_post( - get_credentials=partial(self._signing_credentials, optional_params, region), - service_name=aws_signing_target(api_base)[0], - aws_region_name=region, - url=api_base, - body=body, - headers={ - "Content-Type": "application/json", - "x-amz-content-sha256": hashlib.sha256(body.encode()).hexdigest(), - }, - ) + body: Final = json.dumps(request_data).encode() + signed: Final = self.sign_http_request("POST", api_base, body, "application/json", optional_params) signed_headers: Final = dict(signed.headers.items()) signed_names: Final = frozenset(name.lower() for name in signed_headers) return { **{name: value for name, value in headers.items() if name.lower() not in signed_names}, **signed_headers, - }, body.encode() + }, body + + def sign_http_request( + self, method: str, url: str, body: bytes, content_type: str, litellm_params: Mapping[str, object] + ) -> "AWSPreparedRequest": + region: Final = self._signing_region(litellm_params, url) + return sign_aws_request( + method=method, + get_credentials=partial(self._signing_credentials, litellm_params, region), + service_name=aws_signing_target(url)[0], + aws_region_name=region, + url=url, + body=body, + headers={"Content-Type": content_type, "x-amz-content-sha256": hashlib.sha256(body).hexdigest()}, + ) def transform_search_vector_store_response( self, response: httpx.Response, litellm_logging_obj: "LiteLLMLoggingObj" diff --git a/litellm/llms/s3_vectors/vector_stores/transformation.py b/litellm/llms/s3_vectors/vector_stores/transformation.py index 2b0de3dc9ea..1b208542993 100644 --- a/litellm/llms/s3_vectors/vector_stores/transformation.py +++ b/litellm/llms/s3_vectors/vector_stores/transformation.py @@ -1,5 +1,5 @@ from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, ClassVar, Final import httpx @@ -9,6 +9,7 @@ from litellm.llms.base_llm.vector_store.transformation import ( VectorStoreEmbeddingExecutor, ) from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.web_identity_session_policy import S3_VECTORS_QUERY_SESSION_POLICY, WebIdentitySessionPolicy from litellm.types.rag import RAGIngestEmbeddingOptions from litellm.types.router import GenericLiteLLMParams from litellm.types.vector_stores import ( @@ -77,6 +78,8 @@ def s3_vectors_ingest_embedding_options( class S3VectorsVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, BaseAWSLLM): """Vector store configuration for AWS S3 Vectors.""" + _web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = S3_VECTORS_QUERY_SESSION_POLICY + def __init__(self) -> None: BaseQueryEmbeddingVectorStoreConfig.__init__(self) BaseAWSLLM.__init__(self) diff --git a/litellm/llms/sagemaker/chat/handler.py b/litellm/llms/sagemaker/chat/handler.py index 3e671ce5d07..19c612d49be 100644 --- a/litellm/llms/sagemaker/chat/handler.py +++ b/litellm/llms/sagemaker/chat/handler.py @@ -1,13 +1,14 @@ import json from collections.abc import Callable from copy import deepcopy -from typing import Final +from typing import ClassVar, Final import httpx from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.optional_imports import ensure_optional_import from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, pop_aws_auth_params +from litellm.llms.bedrock.web_identity_session_policy import SAGEMAKER_SESSION_POLICY, WebIdentitySessionPolicy from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.utils import ModelResponse, get_secret @@ -16,6 +17,8 @@ from .transformation import SagemakerChatConfig class SagemakerChatHandler(BaseAWSLLM): + _web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = SAGEMAKER_SESSION_POLICY + def _load_credentials( self, optional_params: dict, diff --git a/litellm/llms/sagemaker/chat/transformation.py b/litellm/llms/sagemaker/chat/transformation.py index 76eafcd6bb6..475c3ba213d 100644 --- a/litellm/llms/sagemaker/chat/transformation.py +++ b/litellm/llms/sagemaker/chat/transformation.py @@ -8,7 +8,7 @@ Huggingface Docs: https://huggingface.co/docs/text-generation-inference/en/messa """ from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, ClassVar, Final, cast import httpx from httpx._models import Headers @@ -18,6 +18,7 @@ from litellm.litellm_core_utils.logging_utils import track_llm_api_timing from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.web_identity_session_policy import SAGEMAKER_SESSION_POLICY, WebIdentitySessionPolicy from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, @@ -39,6 +40,8 @@ else: class SagemakerChatConfig(OpenAIGPTConfig, BaseAWSLLM): + _web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = SAGEMAKER_SESSION_POLICY + def __init__(self, **kwargs): OpenAIGPTConfig.__init__(self, **kwargs) BaseAWSLLM.__init__(self, **kwargs) diff --git a/litellm/llms/sagemaker/completion/handler.py b/litellm/llms/sagemaker/completion/handler.py index b89c59a0e5e..a852ae8b1a1 100644 --- a/litellm/llms/sagemaker/completion/handler.py +++ b/litellm/llms/sagemaker/completion/handler.py @@ -2,7 +2,7 @@ import json import types from collections.abc import Callable, Mapping from copy import deepcopy -from typing import Final, cast +from typing import ClassVar, Final, cast import httpx from pydantic import TypeAdapter @@ -15,6 +15,7 @@ from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.optional_imports import ensure_optional_import from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, pop_aws_auth_params +from litellm.llms.bedrock.web_identity_session_policy import SAGEMAKER_SESSION_POLICY, WebIdentitySessionPolicy from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, get_httpx_client, @@ -57,6 +58,8 @@ os.environ['AWS_SECRET_ACCESS_KEY'] = "" # set os.environ['AWS_REGION_NAME'] = class SagemakerLLM(BaseAWSLLM): + _web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = SAGEMAKER_SESSION_POLICY + def _load_credentials( self, optional_params: dict, diff --git a/litellm/main.py b/litellm/main.py index f3620d7a8c2..febf9503a7e 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -76,6 +76,7 @@ from litellm.exceptions import LiteLLMUnknownProvider from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.allowed_tools_tool_choice import ( raise_if_allowed_tools_tool_choice_unsupported, + raise_if_hosted_tool_choice_unsupported, ) from litellm.litellm_core_utils.asyncify import asyncify, run_async_function from litellm.litellm_core_utils.audio_utils.utils import ( @@ -5915,6 +5916,13 @@ def completion( custom_llm_provider=custom_llm_provider, model=model, ) + raise_if_hosted_tool_choice_unsupported( + tool_choice=cast( # cast-ok: widening an untyped dict to object is always sound, the helper validates it with pydantic + object, tool_choice + ), + custom_llm_provider=custom_llm_provider, + model=model, + ) if provider_specific_header is not None: headers.update( diff --git a/litellm/model_latest_aliases.json b/litellm/model_latest_aliases.json deleted file mode 100644 index 7b0fa5d3d39..00000000000 --- a/litellm/model_latest_aliases.json +++ /dev/null @@ -1,21 +0,0 @@ -{ - "openai": { - "gpt-luna-latest": "gpt-6-luna", - "gpt-sol-latest": "gpt-6.1-sol", - "gpt-terra-latest": "gpt-6.1-terra", - "gpt-astra-latest": "gpt-6-astra", - "gpt-mini-latest": "gpt-5.4-mini" - }, - "anthropic": { - "claude-opus-latest": "claude-opus-5-5", - "claude-sonnet-latest": "claude-sonnet-5-5", - "claude-haiku-latest": "claude-haiku-5-5", - "claude-fable-latest": "claude-fable-5-1" - }, - "xai": {"grok-latest": "grok-4.7"}, - "moonshot": {"kimi-latest": "kimi-k3"}, - "zai": {"glm-latest": "glm-5.3", "glm-flash-latest": "glm-5.3-flash"}, - "deepseek": {"deepseek-pro-latest": "deepseek-v4-pro", "deepseek-flash-latest": "deepseek-v4-flash"}, - "dashscope": {"qwen-max-latest": "qwen3.8-max", "qwen-flash-latest": "qwen3.8-flash"}, - "minimax": {"minimax-latest": "MiniMax-M3"} -} diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index ac5e99c9313..79e4d92aa9d 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -1969,6 +1969,9 @@ "prompt_cache_min_tokens": 512, "source": "https://aws.amazon.com/bedrock/pricing/" }, + "bedrock/claude-fable-latest": { + "alias_of": "global.anthropic.claude-fable-5-1" + }, "us.anthropic.claude-fable-5": { "cache_creation_input_token_cost": 1.375e-05, "cache_creation_input_token_cost_above_1hr": 2.2e-05, @@ -2305,6 +2308,9 @@ "region": "us-east-1" } }, + "bedrock/claude-opus-latest": { + "alias_of": "global.anthropic.claude-opus-5-5" + }, "us.anthropic.claude-opus-5": { "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, @@ -4063,6 +4069,9 @@ "deprecation_date": "2027-12-05", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement" }, + "azure_ai/claude-fable-latest": { + "alias_of": "azure_ai/claude-fable-5-1" + }, "azure_ai/claude-mythos-5-1": { "input_cost_per_token": 1e-05, "output_cost_per_token": 5e-05, @@ -4165,6 +4174,9 @@ "prompt_cache_min_tokens": 512, "source": "https://management.azure.com/subscriptions/c873328e-b572-4770-8dff-aaeb6f1f0e79/providers/Microsoft.CognitiveServices/locations/eastus2/models?api-version=2024-10-01" }, + "azure_ai/claude-opus-latest": { + "alias_of": "azure_ai/claude-opus-5-5" + }, "azure_ai/claude-opus-4-8": { "deprecation_date": "2027-09-01", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement", @@ -4382,6 +4394,9 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure_ai/gpt-astra-latest": { + "alias_of": "azure_ai/gpt-6-astra" + }, "azure_ai/gpt-6-luna": { "cache_creation_input_token_cost": 1.25e-07, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-07, @@ -4430,6 +4445,9 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure_ai/gpt-luna-latest": { + "alias_of": "azure_ai/gpt-6-luna" + }, "azure_ai/gpt-6-sol": { "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, @@ -4526,6 +4544,9 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure_ai/gpt-sol-latest": { + "alias_of": "azure_ai/gpt-6.1-sol" + }, "azure_ai/gpt-6.1-terra": { "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-06, @@ -4574,6 +4595,9 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure_ai/gpt-terra-latest": { + "alias_of": "azure_ai/gpt-6.1-terra" + }, "azure_ai/gpt-5.5": { "deprecation_date": "2027-10-26", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement", @@ -5036,6 +5060,9 @@ "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false }, + "azure_ai/gpt-mini-latest": { + "alias_of": "azure_ai/gpt-5.4-mini" + }, "azure_ai/gpt-5.4-mini-2026-03-17": { "deprecation_date": "2027-09-21", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement", @@ -5316,7 +5343,7 @@ "supports_tool_choice": true }, "azure/eu/gpt-4o-realtime-preview-2024-12-17": { - "cache_read_input_audio_token_cost": 2.5e-06, + "cache_read_input_audio_token_cost": 2.75e-06, "cache_read_input_token_cost": 2.75e-06, "input_cost_per_audio_token": 4.4e-05, "input_cost_per_token": 5.5e-06, @@ -5326,7 +5353,7 @@ "mode": "realtime", "output_cost_per_audio_token": 8.8e-05, "output_cost_per_token": 2.2e-05, - "source": "https://azure.microsoft.com/en-us/pricing/details/azure-openai/", + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_modalities": [ "text", "audio" @@ -6664,6 +6691,7 @@ "supports_tool_choice": true }, "azure/gpt-realtime": { + "cache_read_input_image_token_cost": 5e-07, "cache_creation_input_audio_token_cost": 4e-06, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, @@ -6699,9 +6727,11 @@ "supports_parallel_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure" + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/" }, "azure/gpt-realtime-2025-08-28": { + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "cache_read_input_image_token_cost": 5e-07, "cache_creation_input_audio_token_cost": 4e-06, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, @@ -6739,6 +6769,7 @@ "supports_tool_choice": true }, "azure/gpt-realtime-1.5": { + "cache_read_input_image_token_cost": 5e-07, "cache_creation_input_audio_token_cost": 4e-06, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, @@ -6775,9 +6806,11 @@ "supports_parallel_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure" + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/" }, "azure/gpt-realtime-1.5-2026-02-23": { + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "cache_read_input_image_token_cost": 5e-07, "cache_creation_input_audio_token_cost": 4e-06, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, @@ -6898,6 +6931,8 @@ "supports_tool_choice": true }, "azure/gpt-realtime-mini": { + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "cache_read_input_image_token_cost": 8e-08, "cache_creation_input_audio_token_cost": 3e-07, "cache_read_input_audio_token_cost": 3e-07, "cache_read_input_token_cost": 6e-08, @@ -7006,6 +7041,7 @@ ] }, "azure/gpt-4o-realtime-preview-2024-12-17": { + "cache_read_input_audio_token_cost": 2.5e-06, "cache_read_input_token_cost": 2.5e-06, "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 5e-06, @@ -7015,7 +7051,7 @@ "mode": "realtime", "output_cost_per_audio_token": 8e-05, "output_cost_per_token": 2e-05, - "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure", + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_modalities": [ "text", "audio" @@ -9080,6 +9116,9 @@ "cache_read_input_token_cost_flex": 5e-07, "cache_creation_input_token_cost_flex": 6.25e-06 }, + "azure/gpt-astra-latest": { + "alias_of": "azure/gpt-6-astra" + }, "azure/gpt-6-astra-2026-09-03": { "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, @@ -9184,6 +9223,9 @@ "output_cost_per_token_above_272k_tokens_flex": 3.75e-07, "output_cost_per_token_flex": 2.5e-07 }, + "azure/gpt-luna-latest": { + "alias_of": "azure/gpt-6-luna" + }, "azure/gpt-6-luna-2026-09-22": { "cache_creation_input_token_cost": 1.25e-07, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-07, @@ -9403,6 +9445,9 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure/gpt-terra-latest": { + "alias_of": "azure/gpt-6.1-terra" + }, "azure/gpt-6.1-sol": { "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, @@ -9455,6 +9500,9 @@ "cache_read_input_token_cost_above_272k_tokens_flex": 1e-07, "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06 }, + "azure/gpt-sol-latest": { + "alias_of": "azure/gpt-6.1-sol" + }, "azure/gpt-6.1-sol-2026-09-29": { "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, @@ -11146,6 +11194,9 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_xhigh_reasoning_effort": true }, + "azure/gpt-mini-latest": { + "alias_of": "azure/gpt-5.4-mini" + }, "azure/gpt-5.4-mini-2026-03-17": { "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_flex": 3.75e-08, @@ -12388,7 +12439,7 @@ "supports_tool_choice": true }, "azure/us/gpt-4o-realtime-preview-2024-12-17": { - "cache_read_input_audio_token_cost": 2.5e-06, + "cache_read_input_audio_token_cost": 2.75e-06, "cache_read_input_token_cost": 2.75e-06, "input_cost_per_audio_token": 4.4e-05, "input_cost_per_token": 5.5e-06, @@ -12398,7 +12449,7 @@ "mode": "realtime", "output_cost_per_audio_token": 8.8e-05, "output_cost_per_token": 2.2e-05, - "source": "https://azure.microsoft.com/en-us/pricing/details/azure-openai/", + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_modalities": [ "text", "audio" @@ -13212,6 +13263,9 @@ "supports_tool_choice": true, "supports_vision": true }, + "azure_ai/kimi-latest": { + "alias_of": "azure_ai/FW-Kimi-K3" + }, "azure_ai/FW-MiniMax-M2.5": { "deprecation_date": "2027-07-01", "cache_read_input_token_cost": 3.3e-08, @@ -13251,6 +13305,9 @@ "supports_tool_choice": true, "supports_vision": true }, + "azure_ai/minimax-latest": { + "alias_of": "azure_ai/FW-MiniMax-M3" + }, "azure_ai/FW-Nemotron-Lightning-3.5-30B-A3B": { "cache_read_input_token_cost": 1e-08, "input_cost_per_token": 6e-08, @@ -13656,6 +13713,9 @@ "cache_read_input_token_cost": 1.45e-07, "supports_prompt_caching": true }, + "azure_ai/deepseek-pro-latest": { + "alias_of": "azure_ai/deepseek-v4-pro" + }, "azure_ai/deepseek-v4-flash": { "deprecation_date": "2028-02-20", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement", @@ -13674,6 +13734,9 @@ "cache_read_input_token_cost": 2.8e-08, "supports_prompt_caching": true }, + "azure_ai/deepseek-flash-latest": { + "alias_of": "azure_ai/deepseek-v4-flash" + }, "azure_ai/DeepSeek-V4-Flash-0731": { "cache_read_input_token_cost": 1.4e-08, "deprecation_date": "2026-12-03", @@ -16623,6 +16686,9 @@ "source": "https://platform.claude.com/docs/en/about-claude/pricing", "supports_web_search": true }, + "claude-sonnet-latest": { + "alias_of": "claude-sonnet-5-5" + }, "claude-sonnet-4-6": { "supports_web_search": true, "supports_anthropic_compaction": true, @@ -17024,6 +17090,9 @@ "supports_native_structured_output": true, "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, + "claude-fable-latest": { + "alias_of": "claude-fable-5-1" + }, "claude-opus-5-5": { "supports_web_search": true, "supports_anthropic_compaction": true, @@ -17069,6 +17138,9 @@ "prompt_cache_min_tokens": 512, "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, + "claude-opus-latest": { + "alias_of": "claude-opus-5-5" + }, "claude-opus-5": { "supports_web_search": true, "supports_anthropic_compaction": true, @@ -18868,6 +18940,9 @@ "supports_tool_choice": true, "supports_vision": true }, + "dashscope/qwen-max-latest": { + "alias_of": "dashscope/qwen3.8-max" + }, "dashscope/qwen3.8-flash": { "cache_creation_input_token_cost": 2e-07, "cache_read_input_token_cost": 1.6e-08, @@ -18887,6 +18962,9 @@ "supports_vision": true, "supports_web_search": true }, + "dashscope/qwen-flash-latest": { + "alias_of": "dashscope/qwen3.8-flash" + }, "dashscope/qwen3.8-omni-flash": { "cache_read_input_token_cost": 1.6e-08, "input_cost_per_token": 1.5e-07, @@ -21767,6 +21845,9 @@ "supports_vision": true, "thinking_always_on": true }, + "databricks/claude-fable-latest": { + "alias_of": "databricks/databricks-claude-fable-5-1" + }, "databricks/databricks-claude-haiku-4-5": { "cache_creation_input_token_cost": 1.24999e-06, "cache_creation_input_token_cost_above_1hr": 1.99997e-06, @@ -22007,6 +22088,9 @@ "supports_xhigh_reasoning_effort": true, "thinking_always_on": true }, + "databricks/claude-opus-latest": { + "alias_of": "databricks/databricks-claude-opus-5-5" + }, "databricks/databricks-claude-sonnet-4": { "cache_creation_input_token_cost": 3.74997e-06, "cache_read_input_token_cost": 3.0002e-07, @@ -22724,6 +22808,9 @@ "cache_read_input_token_cost_priority": 4.55e-07, "cache_creation_input_token_cost_priority": 2.45e-06 }, + "databricks/glm-latest": { + "alias_of": "databricks/databricks-glm-5-3" + }, "databricks/databricks-glm-5-3-flash": { "cache_creation_input_token_cost": 1.5001e-07, "cache_read_input_token_cost": 3.003e-08, @@ -22757,6 +22844,9 @@ "cache_read_input_token_cost_priority": 5.25e-08, "cache_creation_input_token_cost_priority": 2.625e-07 }, + "databricks/glm-flash-latest": { + "alias_of": "databricks/databricks-glm-5-3-flash" + }, "databricks/databricks-gpt-5": { "cache_creation_input_token_cost": 1.24999e-06, "cache_read_input_token_cost": 1.2502e-07, @@ -22919,6 +23009,9 @@ "cache_read_input_token_cost_priority": 1.5001e-07, "cache_creation_input_token_cost_priority": 1.50003e-06 }, + "databricks/gpt-mini-latest": { + "alias_of": "databricks/databricks-gpt-5-4-mini" + }, "databricks/databricks-gpt-5-4-nano": { "cache_creation_input_token_cost": 1.9999e-07, "cache_read_input_token_cost": 2.002e-08, @@ -23315,6 +23408,9 @@ "cache_read_input_token_cost_priority": 5.25e-07, "cache_creation_input_token_cost_priority": 5.25e-06 }, + "databricks/kimi-latest": { + "alias_of": "databricks/databricks-kimi-k3" + }, "databricks/system.ai.kimi-k3": { "litellm_provider": "databricks", "mode": "chat", @@ -28843,6 +28939,9 @@ "supports_tool_choice": true, "supports_vision": true }, + "fireworks_ai/minimax-latest": { + "alias_of": "fireworks_ai/accounts/fireworks/models/minimax-m3" + }, "fireworks_ai/accounts/fireworks/models/mixtral-8x22b-instruct-hf": { "input_cost_per_token": 1.2e-06, "litellm_provider": "fireworks_ai", @@ -29273,52 +29372,6 @@ "comment": "Open flagship GLM for long-horizon coding agents and million-token context work", "source": "https://api.friendli.ai/serverless/v1/models" }, - "friendliai/deepseek-ai/DeepSeek-V3.2": { - "litellm_provider": "friendliai", - "max_input_tokens": 163840, - "max_output_tokens": 163840, - "input_cost_per_token": 5e-07, - "output_cost_per_token": 1.5e-06, - "cache_read_input_token_cost": 2.5e-07, - "supports_prompt_caching": true, - "supports_reasoning": true, - "reasoning_effort_levels": [], - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_native_structured_output": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": false, - "supports_image_input": false, - "supports_video_input": false, - "mode": "chat", - "comment": "DeepSeek chat model for instruction following, coding, and analysis", - "source": "https://api.friendli.ai/serverless/v1/models" - }, - "friendliai/MiniMaxAI/MiniMax-M2.5": { - "litellm_provider": "friendliai", - "max_input_tokens": 196608, - "max_output_tokens": 196608, - "input_cost_per_token": 3e-07, - "output_cost_per_token": 1.2e-06, - "cache_read_input_token_cost": 6e-08, - "supports_prompt_caching": true, - "supports_reasoning": true, - "reasoning_effort_levels": [], - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_native_structured_output": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": false, - "supports_image_input": false, - "supports_video_input": false, - "mode": "chat", - "comment": "Prior MiniMax coding model for agent workflows, office edits, and automation", - "source": "https://api.friendli.ai/serverless/v1/models" - }, "friendliai/zai-org/GLM-5.1": { "litellm_provider": "friendliai", "max_input_tokens": 202752, @@ -30485,7 +30538,6 @@ "successor_model": "gemini-3.8-flash", "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "cache_read_input_token_cost_above_200k_tokens_priority": 4.5e-07, "cache_read_input_token_cost_flex": 1.25e-07, "cache_read_input_token_cost_priority": 2.25e-07, @@ -30544,7 +30596,6 @@ "output_cost_per_token_priority": 1.8e-05 }, "gemini-3.1-pro-preview": { - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, @@ -30608,7 +30659,6 @@ "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, "input_cost_per_token_batches": 1e-06, @@ -30660,7 +30710,6 @@ "vertex_ai/gemini-3-pro-preview": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, "input_cost_per_token_batches": 1e-06, @@ -31104,7 +31153,6 @@ "supports_web_search": true }, "vertex_ai/gemini-3.1-pro-preview": { - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, @@ -31168,7 +31216,6 @@ "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, "input_cost_per_token_batches": 1e-06, @@ -32947,6 +32994,7 @@ "input_cost_per_token_flex": 1e-06, "output_cost_per_token_flex": 6e-06, "cache_read_input_token_cost_above_200k_tokens_batches": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_flex": 4e-07, "input_cost_per_token_above_200k_tokens_batches": 2e-06, "input_cost_per_token_above_200k_tokens_flex": 2e-06, "output_cost_per_token_above_200k_tokens_batches": 9e-06, @@ -33015,6 +33063,7 @@ "input_cost_per_token_flex": 1e-06, "output_cost_per_token_flex": 6e-06, "cache_read_input_token_cost_above_200k_tokens_batches": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_flex": 4e-07, "input_cost_per_token_above_200k_tokens_batches": 2e-06, "input_cost_per_token_above_200k_tokens_flex": 2e-06, "output_cost_per_token_above_200k_tokens_batches": 9e-06, @@ -33789,6 +33838,9 @@ "supports_vision": true, "thinking_always_on": true }, + "github_copilot/claude-fable-latest": { + "alias_of": "github_copilot/claude-fable-5.1" + }, "github_copilot/claude-haiku-4.5": { "cache_creation_input_token_cost": 1.25e-06, "cache_read_input_token_cost": 1e-07, @@ -33832,6 +33884,9 @@ "supports_vision": true, "thinking_always_on": false }, + "github_copilot/claude-haiku-latest": { + "alias_of": "github_copilot/claude-haiku-5.5" + }, "github_copilot/claude-opus-4.8": { "cache_creation_input_token_cost": 6.25e-06, "cache_read_input_token_cost": 5e-07, @@ -33920,6 +33975,9 @@ "supports_vision": true, "thinking_always_on": true }, + "github_copilot/claude-opus-latest": { + "alias_of": "github_copilot/claude-opus-5.5" + }, "github_copilot/claude-sonnet-5": { "cache_creation_input_token_cost": 2.5e-06, "cache_read_input_token_cost": 2e-07, @@ -33963,6 +34021,9 @@ "supports_vision": true, "thinking_always_on": true }, + "github_copilot/claude-sonnet-latest": { + "alias_of": "github_copilot/claude-sonnet-5.5" + }, "github_copilot/gemini-3.7-flash": { "cache_read_input_token_cost": 7.5e-08, "deprecation_date": "2026-10-19", @@ -34308,6 +34369,9 @@ "supports_response_schema": true, "supports_vision": true }, + "github_copilot/gpt-astra-latest": { + "alias_of": "github_copilot/gpt-6-astra" + }, "github_copilot/gpt-6-luna": { "cache_creation_input_token_cost": 1.25e-07, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-07, @@ -34331,6 +34395,9 @@ "supports_response_schema": true, "supports_vision": true }, + "github_copilot/gpt-luna-latest": { + "alias_of": "github_copilot/gpt-6-luna" + }, "github_copilot/gpt-6-sol": { "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, @@ -34377,6 +34444,9 @@ "supports_response_schema": true, "supports_vision": true }, + "github_copilot/gpt-sol-latest": { + "alias_of": "github_copilot/gpt-6.1-sol" + }, "github_copilot/kimi-k3": { "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, @@ -34394,6 +34464,9 @@ "supports_response_schema": true, "supports_vision": true }, + "github_copilot/kimi-latest": { + "alias_of": "github_copilot/kimi-k3" + }, "github_copilot/mai-code-1.1-flash": { "cache_read_input_token_cost": 2e-08, "input_cost_per_token": 2e-07, @@ -37687,6 +37760,9 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "gpt-astra-latest": { + "alias_of": "gpt-6-astra" + }, "gpt-6-sol": { "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_12hr": 5e-06, @@ -37884,6 +37960,9 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "gpt-luna-latest": { + "alias_of": "gpt-6-luna" + }, "gpt-5.6": { "cache_creation_input_token_cost": 5e-06, "cache_creation_input_token_cost_above_12hr": 1e-05, @@ -39067,6 +39146,9 @@ "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false }, + "gpt-mini-latest": { + "alias_of": "gpt-5.4-mini" + }, "gpt-5.4-mini-2026-03-17": { "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_batches": 3.75e-08, @@ -43025,6 +43107,9 @@ "max_input_tokens": 1000000, "max_output_tokens": 128000 }, + "minimax/minimax-latest": { + "alias_of": "minimax/MiniMax-M3" + }, "minimax/MiniMax-H3": { "litellm_provider": "minimax", "mode": "video_generation", @@ -44535,6 +44620,9 @@ "supports_video_input": true, "supports_vision": true }, + "moonshot/kimi-latest": { + "alias_of": "moonshot/kimi-k3" + }, "moonshot/moonshot-v1-128k": { "deprecation_date": "2026-08-31", "input_cost_per_token": 2e-06, @@ -50047,13 +50135,13 @@ "supports_reasoning": true }, "together_ai/Qwen/Qwen3.7-Max": { - "cache_read_input_token_cost": 3e-07, - "input_cost_per_token": 1.5e-06, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 2.5e-06, "litellm_provider": "together_ai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "mode": "chat", - "output_cost_per_token": 4.5e-06, + "output_cost_per_token": 7.5e-06, "source": "https://www.together.ai/pricing", "supports_prompt_caching": true }, @@ -50136,6 +50224,9 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "together_ai/deepseek-pro-latest": { + "alias_of": "together_ai/deepseek-ai/DeepSeek-V4-Pro-0813" + }, "together_ai/meta-models/Muse-Glimmer-30B": { "cache_read_input_token_cost": 4e-08, "input_cost_per_token": 3.5e-07, @@ -50182,6 +50273,9 @@ "supports_tool_choice": true, "supports_vision": true }, + "together_ai/kimi-latest": { + "alias_of": "together_ai/moonshotai/Kimi-K3" + }, "together_ai/thinkingmachines/Inkling": { "cache_read_input_token_cost": 1.7e-07, "deprecation_date": "2026-11-06", @@ -50232,6 +50326,9 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "together_ai/glm-latest": { + "alias_of": "together_ai/zai-org/GLM-5.3" + }, "together_ai/zai-org/GLM-5.3-Flash": { "cache_read_input_token_cost": 3e-08, "input_cost_per_token": 1.5e-07, @@ -50249,6 +50346,9 @@ "supports_tool_choice": true, "supports_vision": true }, + "together_ai/glm-flash-latest": { + "alias_of": "together_ai/zai-org/GLM-5.3-Flash" + }, "tts-1": { "deprecation_date": "2027-01-06", "deprecation_source": "https://developers.openai.com/api/docs/deprecations", @@ -53019,6 +53119,9 @@ "output_cost_per_token_batches": 2.5e-05, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" }, + "vertex_ai/claude-fable-latest": { + "alias_of": "vertex_ai/claude-fable-5-1" + }, "vertex_ai/claude-fable-5@default": { "deprecation_date": "2027-06-08", "regional_endpoint_uplift_multiplier": 1.1, @@ -53217,6 +53320,9 @@ "prompt_cache_min_tokens": 512, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" }, + "vertex_ai/claude-opus-latest": { + "alias_of": "vertex_ai/claude-opus-5-5" + }, "vertex_ai/claude-opus-5-5@default": { "regional_endpoint_uplift_multiplier": 1.1, "supports_mid_conversation_system": true, @@ -55765,6 +55871,9 @@ "supports_vision": true, "supports_web_search": true }, + "xai/grok-latest": { + "alias_of": "xai/grok-4.7" + }, "zai.glm-4.7": { "deprecation_date": "2027-03-30", "deprecation_source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-lifecycle-legacy.html", @@ -55866,6 +55975,9 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "zai/glm-latest": { + "alias_of": "zai/glm-5.3" + }, "zai/glm-5.3-flash": { "cache_creation_input_token_cost": 0, "cache_read_input_token_cost": 3e-08, @@ -55883,6 +55995,9 @@ "source": "https://docs.z.ai/guides/overview/pricing", "supports_vision": true }, + "zai/glm-flash-latest": { + "alias_of": "zai/glm-5.3-flash" + }, "zai/glm-5v-turbo": { "cache_read_input_token_cost": 2.4e-07, "input_cost_per_token": 1.2e-06, @@ -61519,6 +61634,9 @@ "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-astra.html", "supports_web_search": true }, + "bedrock_mantle/gpt-astra-latest": { + "alias_of": "bedrock_mantle/openai.gpt-6-astra" + }, "bedrock_mantle/openai.gpt-6-sol": { "input_cost_per_token": 2.2e-06, "input_cost_per_token_above_272k_tokens": 4.4e-06, @@ -61599,6 +61717,9 @@ "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-cards-openai.html", "supports_web_search": true }, + "bedrock_mantle/gpt-luna-latest": { + "alias_of": "bedrock_mantle/openai.gpt-6-luna" + }, "us.openai.gpt-6-astra": { "supports_bedrock_runtime_chat_completions_response_format": true, "input_cost_per_token": 1.1e-05, @@ -61767,6 +61888,9 @@ "/v1/responses" ] }, + "bedrock/gpt-astra-latest": { + "alias_of": "global.openai.gpt-6-astra" + }, "openai.gpt-6-sol": { "input_cost_per_token": 2e-06, "input_cost_per_token_above_272k_tokens": 4e-06, @@ -61909,6 +62033,9 @@ "/v1/responses" ] }, + "bedrock/gpt-luna-latest": { + "alias_of": "global.openai.gpt-6-luna" + }, "bedrock_mantle/openai.gpt-5.5": { "input_cost_per_token": 5.5e-06, "input_cost_per_token_above_272k_tokens": 1.1e-05, @@ -62186,6 +62313,9 @@ "supports_forced_tool_use": false, "token_count_model": {"model":"anthropic.claude-opus-5-5","litellm_provider":"bedrock_mantle","region":"us-east-1"} }, + "bedrock_mantle/claude-opus-latest": { + "alias_of": "bedrock_mantle/anthropic.claude-opus-5-5" + }, "bedrock_mantle/anthropic.claude-sonnet-5-5": { "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, @@ -62226,6 +62356,9 @@ "prompt_cache_min_tokens": 512, "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" }, + "bedrock_mantle/claude-sonnet-latest": { + "alias_of": "bedrock_mantle/anthropic.claude-sonnet-5-5" + }, "us.xai.grok-4.6": { "supports_regex_lookaround": false, "input_cost_per_token": 2.2e-06, @@ -63796,6 +63929,9 @@ "token_count_framing": 23, "token_count_tool_framing": 12 }, + "deepseek/deepseek-flash-latest": { + "alias_of": "deepseek/deepseek-v4-flash" + }, "deepseek/deepseek-v4-flash-vision-exp": { "cache_creation_input_token_cost": 0.0, "cache_read_input_token_cost": 6e-09, @@ -63904,6 +64040,9 @@ "token_count_framing": 76, "token_count_tool_framing": 75 }, + "deepseek/deepseek-pro-latest": { + "alias_of": "deepseek/deepseek-v4-pro" + }, "tencent/deepseek-v4-pro": { "cache_creation_input_token_cost": 0.0, "cache_read_input_token_cost": 1.45e-07, @@ -65327,6 +65466,9 @@ "supports_tool_choice": true, "supports_vision": true }, + "fireworks_ai/kimi-latest": { + "alias_of": "fireworks_ai/accounts/fireworks/models/kimi-k3" + }, "fireworks_ai/deepseek-v4-flash-0731": { "cache_read_input_token_cost": 7e-09, "cache_read_input_token_cost_priority": 8.75e-09, @@ -67576,7 +67718,7 @@ "deepinfra/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B": { "max_output_tokens": 262144, "max_input_tokens": 262144, - "input_cost_per_token": 8.5e-08, + "input_cost_per_token": 7.5e-08, "output_cost_per_token": 4e-07, "litellm_provider": "deepinfra", "mode": "chat", @@ -69000,6 +69142,9 @@ "output_cost_per_token": 4.7e-07, "source": "https://api.together.ai/v1/models" }, + "together_ai/qwen-flash-latest": { + "alias_of": "together_ai/Qwen/Qwen3.8-Flash" + }, "together_ai/moonshotai/Kimi-K2.5-fp4": { "input_cost_per_token": 5e-07, "output_cost_per_token": 2.8e-06, @@ -69216,6 +69361,9 @@ "supports_reasoning": true, "supports_tool_choice": true }, + "azure_ai/glm-latest": { + "alias_of": "azure_ai/FW-GLM-5.3" + }, "azure_ai/FW-GLM-5.3-Flash": { "cache_read_input_token_cost": 3.8e-08, "input_cost_per_token": 1.88e-07, @@ -69231,6 +69379,9 @@ "supports_reasoning": true, "supports_tool_choice": true }, + "azure_ai/glm-flash-latest": { + "alias_of": "azure_ai/FW-GLM-5.3-Flash" + }, "azure_ai/FW-GPT-OSS-120B": { "deprecation_date": "2027-07-01", "cache_read_input_token_cost": 8.2e-08, @@ -71236,6 +71387,9 @@ "supports_tool_choice": true, "supports_vision": true }, + "baseten/glm-latest": { + "alias_of": "baseten/zai-org/GLM-5.3" + }, "baseten/zai-org/GLM-5.3-Fast": { "cache_read_input_token_cost": 2.1e-07, "input_cost_per_token": 2.1e-06, @@ -71579,6 +71733,13 @@ "output_cost_per_token": 1.7e-06, "source": "https://api.together.ai/v1/models" }, + "together_ai/essentialai/rnj-1-instruct": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 1.5e-07, + "source": "https://www.together.ai/models/rnj-1-instruct" + }, "together_ai/google/gemma-4-31B-it": { "deprecation_date": "2026-09-15", "deprecation_source": "https://docs.together.ai/docs/deprecations", @@ -72720,6 +72881,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/gpt-realtime-2": { + "cache_read_input_image_token_cost": 5e-07, "cache_creation_input_audio_token_cost": 4e-07, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, @@ -72738,7 +72900,7 @@ "provider_specific_entry": { "realtime_ga_only": 1 }, - "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_endpoints": [ "/v1/realtime" ], @@ -74425,6 +74587,9 @@ "supports_tool_choice": true, "supports_vision": true }, + "bedrock/kimi-latest": { + "alias_of": "global.moonshotai.kimi-k3" + }, "us.moonshotai.kimi-k3": { "supports_regex_lookaround": false, "cache_creation_input_token_cost": 4.125e-06, @@ -74975,6 +75140,9 @@ "supports_tool_choice": true, "supports_vision": true }, + "baseten/kimi-latest": { + "alias_of": "baseten/moonshotai/Kimi-K3" + }, "baseten/nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B": { "cache_read_input_token_cost": 1.2e-07, "input_cost_per_token": 6e-07, @@ -75057,6 +75225,9 @@ "supports_tool_choice": true, "supports_vision": true }, + "baseten/glm-flash-latest": { + "alias_of": "baseten/zai-org/GLM-5.3-Flash" + }, "baseten/deepseek-ai/DeepSeek-V4-Flash-0731": { "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 1.3e-07, @@ -75073,6 +75244,9 @@ "supports_tool_choice": true, "supports_vision": false }, + "baseten/deepseek-flash-latest": { + "alias_of": "baseten/deepseek-ai/DeepSeek-V4-Flash-0731" + }, "baseten/deepseek-ai/DeepSeek-V4-Pro": { "deprecation_date": "2026-10-16", "cache_read_input_token_cost": 1.45e-07, @@ -75106,6 +75280,9 @@ "supports_tool_choice": true, "supports_vision": false }, + "baseten/deepseek-pro-latest": { + "alias_of": "baseten/deepseek-ai/DeepSeek-V4-Pro-0813" + }, "baseten/zai-org/GLM-5.2-Fast": { "cache_read_input_token_cost": 2.1e-07, "input_cost_per_token": 2.1e-06, @@ -76097,6 +76274,9 @@ "prompt_cache_min_tokens": 512, "source": "https://management.azure.com/subscriptions/c873328e-b572-4770-8dff-aaeb6f1f0e79/providers/Microsoft.CognitiveServices/locations/eastus2/models?api-version=2024-10-01" }, + "azure_ai/claude-sonnet-latest": { + "alias_of": "azure_ai/claude-sonnet-5-5" + }, "bedrock/us-gov-east-1/anthropic.claude-sonnet-5-5": { "bedrock_converse_supports_strict_tools": false, "bedrock_output_config_effort_ceiling": "xhigh", @@ -76261,6 +76441,9 @@ "prompt_cache_min_tokens": 512, "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" }, + "bedrock/claude-sonnet-latest": { + "alias_of": "global.anthropic.claude-sonnet-5-5" + }, "jp.anthropic.claude-sonnet-5-5": { "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, @@ -76520,6 +76703,9 @@ "prompt_cache_min_tokens": 512, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" }, + "vertex_ai/claude-sonnet-latest": { + "alias_of": "vertex_ai/claude-sonnet-5-5" + }, "vertex_ai/claude-sonnet-5-5@default": { "regional_endpoint_uplift_multiplier": 1.1, "supports_mid_conversation_system": true, @@ -76632,6 +76818,9 @@ "supports_tool_choice": true, "supports_vision": true }, + "bedrock/grok-latest": { + "alias_of": "global.xai.grok-4.7" + }, "us.xai.grok-4.7": { "supports_regex_lookaround": false, "supported_endpoints": ["/v1/chat/completions"], @@ -76804,6 +76993,9 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "gpt-sol-latest": { + "alias_of": "gpt-6.1-sol" + }, "global.openai.gpt-6.1-sol": { "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, @@ -76850,6 +77042,9 @@ "supports_sampling_params": false, "supports_xhigh_reasoning_effort": true }, + "bedrock/gpt-sol-latest": { + "alias_of": "global.openai.gpt-6.1-sol" + }, "bedrock_mantle/openai.gpt-6.1-sol": { "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_272k_tokens": 5.5e-06, @@ -76898,6 +77093,9 @@ "supports_xhigh_reasoning_effort": true, "use_openai_responses_path": true }, + "bedrock_mantle/gpt-sol-latest": { + "alias_of": "bedrock_mantle/openai.gpt-6.1-sol" + }, "us.openai.gpt-6.1-sol": { "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_272k_tokens": 5.5e-06, @@ -76989,6 +77187,9 @@ "supports_tool_choice": true, "supports_vision": true }, + "vertex_ai/grok-latest": { + "alias_of": "vertex_ai/xai/grok-4.7" + }, "azure_ai/kimi-k2-thinking": { "deprecation_date": "2026-03-29", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", @@ -77046,6 +77247,9 @@ "supports_tool_choice": true, "supports_vision": false }, + "bedrock/glm-latest": { + "alias_of": "global.zai.glm-5.3" + }, "us.zai.glm-5.3": { "cache_creation_input_token_cost": 2.31e-06, "cache_creation_input_token_cost_batches": 1.155e-06, @@ -77269,6 +77473,9 @@ "supports_vision": true, "supports_web_search": true }, + "azure_ai/grok-latest": { + "alias_of": "azure_ai/grok-4.7" + }, "in.moonshotai.kimi-k3": { "cache_creation_input_token_cost": 4.125e-06, "cache_creation_input_token_cost_flex": 2.0625e-06, @@ -77355,6 +77562,9 @@ "cache_creation_input_token_cost_above_1hr_above_100k_tokens": 1e-06, "cache_read_input_token_cost_above_100k_tokens": 5e-08 }, + "claude-haiku-latest": { + "alias_of": "claude-haiku-5-5" + }, "bedrock_mantle/anthropic.claude-haiku-5-5": { "input_cost_per_token_above_100k_tokens": 5.5e-07, "output_cost_per_token_above_100k_tokens": 2.75e-06, @@ -77401,6 +77611,9 @@ "prompt_cache_min_tokens": 512, "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-haiku-5-5.html" }, + "bedrock_mantle/claude-haiku-latest": { + "alias_of": "bedrock_mantle/anthropic.claude-haiku-5-5" + }, "bedrock_mantle/us-gov-west-1/anthropic.claude-haiku-5-5": { "input_cost_per_token_above_100k_tokens": 6e-07, "output_cost_per_token_above_100k_tokens": 3e-06, @@ -77632,6 +77845,9 @@ "prompt_cache_min_tokens": 512, "source": "https://platform.claude.com/docs/en/models/haiku-5-5/migration-guide" }, + "azure_ai/claude-haiku-latest": { + "alias_of": "azure_ai/claude-haiku-5-5" + }, "bedrock/us-gov-east-1/anthropic.claude-haiku-5-5": { "input_cost_per_token_above_100k_tokens": 6e-07, "output_cost_per_token_above_100k_tokens": 3e-06, @@ -77822,6 +78038,9 @@ "prompt_cache_min_tokens": 512, "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" }, + "bedrock/claude-haiku-latest": { + "alias_of": "global.anthropic.claude-haiku-5-5" + }, "jp.anthropic.claude-haiku-5-5": { "input_cost_per_token_above_100k_tokens": 5.5e-07, "output_cost_per_token_above_100k_tokens": 2.75e-06, @@ -78029,6 +78248,9 @@ "prompt_cache_min_tokens": 512, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" }, + "vertex_ai/claude-haiku-latest": { + "alias_of": "vertex_ai/claude-haiku-5-5" + }, "vertex_ai/claude-haiku-5-5@default": { "input_cost_per_token_above_100k_tokens": 5e-07, "output_cost_per_token_above_100k_tokens": 2.5e-06, @@ -78933,6 +79155,9 @@ "supports_response_schema": true, "supports_vision": true }, + "github_copilot/grok-latest": { + "alias_of": "github_copilot/grok-4.7" + }, "together_ai/deepseek-ai/DeepSeek-V4.1-Flash-1017": { "cache_read_input_token_cost": 6e-09, "input_cost_per_token": 3e-07, @@ -80356,6 +80581,9 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "gpt-terra-latest": { + "alias_of": "gpt-6.1-terra" + }, "chatgpt/gpt-6.1-terra": { "litellm_provider": "chatgpt", "max_input_tokens": 922000, @@ -80414,6 +80642,9 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true }, + "bedrock/gpt-terra-latest": { + "alias_of": "global.openai.gpt-6.1-terra" + }, "us.openai.gpt-6.1-terra": { "cache_creation_input_token_cost": 1.375e-06, "cache_creation_input_token_cost_above_272k_tokens": 2.75e-06, @@ -80493,6 +80724,9 @@ "supports_response_schema": true, "supports_vision": true }, + "github_copilot/gpt-terra-latest": { + "alias_of": "github_copilot/gpt-6.1-terra" + }, "refuseless/glm-5.3-abliterated": { "cache_read_input_token_cost": 1.5e-07, "input_cost_per_token": 1e-06, @@ -80647,6 +80881,9 @@ "supports_xhigh_reasoning_effort": true, "use_openai_responses_path": true }, + "bedrock_mantle/gpt-terra-latest": { + "alias_of": "bedrock_mantle/openai.gpt-6.1-terra" + }, "clf_ai_gateway/glm-5.3": { "metadata": { "comment": "Billed price checked 2026-10-29, which the gateway shows as 60% off a list price 2.5x higher. Update the rows when the discount ends." @@ -82196,6 +82433,9 @@ "supports_vision": true, "thinking_always_on": false }, + "databricks/claude-haiku-latest": { + "alias_of": "databricks/databricks-claude-haiku-5-5" + }, "databricks/databricks-claude-sonnet-5-5": { "cache_creation_input_token_cost": 2.49998e-06, "cache_creation_input_token_cost_above_1hr": 4.00001e-06, @@ -82225,6 +82465,9 @@ "supports_vision": false, "thinking_always_on": true }, + "databricks/claude-sonnet-latest": { + "alias_of": "databricks/databricks-claude-sonnet-5-5" + }, "databricks/databricks-deepseek-v4-1-flash": { "cache_creation_input_token_cost": 3.0002e-07, "cache_read_input_token_cost": 3.003e-08, @@ -82457,6 +82700,9 @@ "cache_read_input_token_cost_above_272k_tokens_priority": 3.9998e-07, "cache_creation_input_token_cost_above_272k_tokens_priority": 9.99999e-06 }, + "databricks/gpt-sol-latest": { + "alias_of": "databricks/databricks-gpt-6-1-sol" + }, "databricks/databricks-gpt-6-1-terra": { "cache_creation_input_token_cost": 1.24999e-06, "cache_creation_input_token_cost_above_272k_tokens": 2.49998e-06, @@ -82497,6 +82743,9 @@ "cache_read_input_token_cost_above_272k_tokens_priority": 1.9999e-07, "cache_creation_input_token_cost_above_272k_tokens_priority": 5.00003e-06 }, + "databricks/gpt-terra-latest": { + "alias_of": "databricks/databricks-gpt-6-1-terra" + }, "databricks/databricks-gpt-6-astra": { "cache_creation_input_token_cost": 1.249997e-05, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, @@ -82537,6 +82786,9 @@ "cache_read_input_token_cost_above_272k_tokens_priority": 4.00001e-06, "cache_creation_input_token_cost_above_272k_tokens_priority": 5.000002e-05 }, + "databricks/gpt-astra-latest": { + "alias_of": "databricks/databricks-gpt-6-astra" + }, "databricks/databricks-gpt-6-luna": { "cache_creation_input_token_cost": 1.2502e-07, "cache_creation_input_token_cost_above_272k_tokens": 2.4997e-07, @@ -82577,6 +82829,9 @@ "cache_read_input_token_cost_above_272k_tokens_priority": 3.997e-08, "cache_creation_input_token_cost_above_272k_tokens_priority": 5.0001e-07 }, + "databricks/gpt-luna-latest": { + "alias_of": "databricks/databricks-gpt-6-luna" + }, "databricks/databricks-gpt-6-sol": { "cache_creation_input_token_cost": 2.49998e-06, "cache_creation_input_token_cost_above_272k_tokens": 5.00003e-06, @@ -82707,6 +82962,9 @@ "supports_tool_choice": true, "supports_vision": true }, + "databricks/grok-latest": { + "alias_of": "databricks/databricks-grok-4-7" + }, "databricks/databricks-kimi-k2-7": { "cache_creation_input_token_cost": 9.4997e-07, "cache_read_input_token_cost": 1.8998e-07, @@ -86639,6 +86897,9 @@ "supports_vision": true, "use_openai_responses_path": true }, + "bedrock_mantle/grok-latest": { + "alias_of": "bedrock_mantle/xai.grok-4.7" + }, "together_ai/ricursive-ai/Ripple-65B-Preview": { "cache_read_input_token_cost": 2e-08, "input_cost_per_token": 2e-07, @@ -87049,5 +87310,484 @@ "supports_parallel_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true + }, + "anthropic.claude-fable-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "eu.anthropic.claude-fable-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.1e-05, + "cache_creation_input_token_cost_above_1hr": 1.76e-05, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 8.8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output_by_default": true, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "eu.anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.1e-05, + "cache_creation_input_token_cost_above_1hr": 1.76e-05, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 8.8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "global.anthropic.claude-fable-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output_by_default": true, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "global.anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "us-gov.anthropic.claude-fable-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.2e-05, + "cache_creation_input_token_cost_above_1hr": 1.92e-05, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 9.6e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.8e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "us-gov.anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.2e-05, + "cache_creation_input_token_cost_above_1hr": 1.92e-05, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 9.6e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.8e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "us.anthropic.claude-fable-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.1e-05, + "cache_creation_input_token_cost_above_1hr": 1.76e-05, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 8.8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output_by_default": true, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "us.anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.1e-05, + "cache_creation_input_token_cost_above_1hr": 1.76e-05, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 8.8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "claude-fable-5-2": { + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_creation_input_token_cost_batches": 5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, + "input_cost_per_token": 8e-06, + "input_cost_per_token_batches": 4e-06, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "output_cost_per_token_batches": 2e-05, + "prompt_cache_min_tokens": 512.0, + "provider_specific_entry": {"us":1.1,"fast":2}, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://platform.claude.com/docs/en/about-claude/pricing", + "supports_adaptive_thinking": true, + "supports_anthropic_compaction": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_fast_mode": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_speed": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "bedrock/claude_platform/claude-fable-5-2": { + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_creation_input_token_cost_batches": 5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, + "input_cost_per_token": 8e-06, + "input_cost_per_token_batches": 4e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "output_cost_per_token_batches": 2e-05, + "prompt_cache_min_tokens": 512.0, + "provider_specific_entry": {"us":1.1}, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing", + "supports_adaptive_thinking": true, + "supports_anthropic_compaction": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "azure_ai/claude-fable-5-2": { + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://platform.claude.com/docs/en/about-claude/pricing", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": false, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "azure_ai/claude-fable-5-2-fast": { + "cache_creation_input_token_cost": 2e-05, + "cache_creation_input_token_cost_above_1hr": 3.2e-05, + "cache_read_input_token_cost": 4e-07, + "input_cost_per_token": 1.6e-05, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 8e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://platform.claude.com/docs/en/about-claude/pricing", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": false, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true } } diff --git a/litellm/proxy/_lazy_features.py b/litellm/proxy/_lazy_features.py index 680b0e6ede4..85f71defe3b 100644 --- a/litellm/proxy/_lazy_features.py +++ b/litellm/proxy/_lazy_features.py @@ -243,6 +243,7 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = ( "/mistral/", "/nvidia_nim/", "/openai/", + "/opensearch/", "/openai_passthrough/", "/tinyfish/", "/transcribe", diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index fde5be049eb..746d6a1e8d7 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -35987,6 +35987,386 @@ ] } }, + "/opensearch/{endpoint}": { + "delete": { + "description": "Call an OpenSearch index through the proxy, e.g. `POST /opensearch//_search`.\n\n`` is a LiteLLM vector store index on an OpenSearch vector store. The key or its team needs `read` or\n`write` on it in `allowed_vector_store_indexes`, and LiteLLM signs the call with the vector store's AWS identity.", + "operationId": "opensearch_proxy_route_opensearch__endpoint__delete", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "additionalProperties": true, + "type": "object" + } + }, + "application/x-ndjson": { + "schema": { + "type": "string" + } + } + }, + "required": false + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "additionalProperties": true, + "type": "object" + } + } + }, + "description": "OpenSearch's response, unmodified" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Opensearch Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "get": { + "description": "Call an OpenSearch index through the proxy, e.g. `POST /opensearch//_search`.\n\n`` is a LiteLLM vector store index on an OpenSearch vector store. The key or its team needs `read` or\n`write` on it in `allowed_vector_store_indexes`, and LiteLLM signs the call with the vector store's AWS identity.", + "operationId": "opensearch_proxy_route_opensearch__endpoint__get", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "additionalProperties": true, + "type": "object" + } + }, + "application/x-ndjson": { + "schema": { + "type": "string" + } + } + }, + "required": false + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "additionalProperties": true, + "type": "object" + } + } + }, + "description": "OpenSearch's response, unmodified" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Opensearch Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "head": { + "description": "Call an OpenSearch index through the proxy, e.g. `POST /opensearch//_search`.\n\n`` is a LiteLLM vector store index on an OpenSearch vector store. The key or its team needs `read` or\n`write` on it in `allowed_vector_store_indexes`, and LiteLLM signs the call with the vector store's AWS identity.", + "operationId": "opensearch_proxy_route_opensearch__endpoint__head", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "additionalProperties": true, + "type": "object" + } + }, + "application/x-ndjson": { + "schema": { + "type": "string" + } + } + }, + "required": false + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "additionalProperties": true, + "type": "object" + } + } + }, + "description": "OpenSearch's response, unmodified" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Opensearch Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "patch": { + "description": "Call an OpenSearch index through the proxy, e.g. `POST /opensearch//_search`.\n\n`` is a LiteLLM vector store index on an OpenSearch vector store. The key or its team needs `read` or\n`write` on it in `allowed_vector_store_indexes`, and LiteLLM signs the call with the vector store's AWS identity.", + "operationId": "opensearch_proxy_route_opensearch__endpoint__patch", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "additionalProperties": true, + "type": "object" + } + }, + "application/x-ndjson": { + "schema": { + "type": "string" + } + } + }, + "required": false + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "additionalProperties": true, + "type": "object" + } + } + }, + "description": "OpenSearch's response, unmodified" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Opensearch Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "post": { + "description": "Call an OpenSearch index through the proxy, e.g. `POST /opensearch//_search`.\n\n`` is a LiteLLM vector store index on an OpenSearch vector store. The key or its team needs `read` or\n`write` on it in `allowed_vector_store_indexes`, and LiteLLM signs the call with the vector store's AWS identity.", + "operationId": "opensearch_proxy_route_opensearch__endpoint__post", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "additionalProperties": true, + "type": "object" + } + }, + "application/x-ndjson": { + "schema": { + "type": "string" + } + } + }, + "required": false + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "additionalProperties": true, + "type": "object" + } + } + }, + "description": "OpenSearch's response, unmodified" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Opensearch Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "put": { + "description": "Call an OpenSearch index through the proxy, e.g. `POST /opensearch//_search`.\n\n`` is a LiteLLM vector store index on an OpenSearch vector store. The key or its team needs `read` or\n`write` on it in `allowed_vector_store_indexes`, and LiteLLM signs the call with the vector store's AWS identity.", + "operationId": "opensearch_proxy_route_opensearch__endpoint__put", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "additionalProperties": true, + "type": "object" + } + }, + "application/x-ndjson": { + "schema": { + "type": "string" + } + } + }, + "required": false + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "additionalProperties": true, + "type": "object" + } + } + }, + "description": "OpenSearch's response, unmodified" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Opensearch Proxy Route", + "tags": [ + "llm_passthrough" + ] + } + }, "/tinyfish/{endpoint}": { "get": { "description": "Pass-through for the TinyFish Agent API (goal-based web automation).\n\nForwarded endpoints:\n- POST /v1/automation/run \u2014 run to completion (blocking)\n- POST /v1/automation/run-async \u2014 submit a run, poll GET /v1/runs/{id} for the result\n- POST /v1/automation/run-sse \u2014 run with SSE progress events\n- GET /v1/runs/{id} \u2014 run status / result\n- POST /v1/runs/{id}/cancel \u2014 cancel a run\n\nEvery other Agent API endpoint (vault, wallet, browser profiles, and the GET /v1/runs\nlisting, which would let any caller discover other callers' run ids) returns 403: all\nproxy callers share one upstream key.\n\nCredential lookup order:\n1. passthrough_endpoint_router (config.yaml deployments with use_in_pass_through)\n2. TINYFISH_API_KEY environment variable\n\n[Docs](https://docs.litellm.ai/docs/pass_through/tinyfish)", diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 830a1063846..b7406b269e2 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -427,6 +427,7 @@ class LiteLLMRoutes(enum.Enum): "/bespoke", "/openrouter", "/milvus", + "/opensearch", "/gigachat", "/watsonx", "/nvidia_nim", diff --git a/litellm/proxy/agent_endpoints/a2a_routing.py b/litellm/proxy/agent_endpoints/a2a_routing.py index fb727367caf..934180ea964 100644 --- a/litellm/proxy/agent_endpoints/a2a_routing.py +++ b/litellm/proxy/agent_endpoints/a2a_routing.py @@ -108,7 +108,14 @@ def _bridge_request_data( if key in data and key in agent_kwargs } return _with_headers( - {**agent_kwargs, **data, **layered, **pinned, "model": bridge_model_name(litellm_params)}, + { + **agent_kwargs, + **data, + **layered, + **pinned, + "model": bridge_model_name(litellm_params), + "custom_llm_provider": agent_kwargs.get("custom_llm_provider"), + }, extra_headers, ) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index b40eb451979..4514882a45f 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -4969,11 +4969,9 @@ def _owning_team_names_for_team_deployment( if llm_router is None or team_id is None or not model.startswith(f"model_name_{team_id}_"): return () deployment: Final = llm_router.registry().deployments(model, source="first_in_group") - public_name: Final = ( - deployment.model_info.team_public_model_name - if deployment is not None and deployment.model_info.team_id == team_id - else None - ) + if deployment is None or deployment.model_info.team_id != team_id: + return () + public_name: Final = deployment.model_info.team_public_model_name legacy_aliases: Final = tuple(alias for alias, target in (team_model_aliases or {}).items() if target == model) return (*((public_name,) if public_name else ()), *legacy_aliases) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 24602dec677..d3978a58c43 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -2,7 +2,7 @@ import importlib.util import os import re import sys -from collections.abc import Collection, Iterator, Mapping +from collections.abc import Collection, Iterator, Mapping, Sequence from functools import lru_cache from itertools import accumulate, chain, repeat from logging import Logger @@ -1556,6 +1556,48 @@ def enforce_batch_limits_are_admin_only( ) +VECTOR_STORE_INDEX_GRANTS_METADATA_KEY: Final = "allowed_vector_store_indexes" + + +def _stored_vector_store_index_grants(existing_metadata: Mapping[str, object] | None) -> object: + if existing_metadata is None: + return None + return existing_metadata.get(VECTOR_STORE_INDEX_GRANTS_METADATA_KEY) or None + + +def _requested_vector_store_index_grants( + data: BatchLimitRequest, + top_level_grants: Sequence[AllowedVectorStoreIndexItem] | None, + stored_grants: object, +) -> object: + if top_level_grants is not None: + return [grant.model_dump() for grant in top_level_grants] or None + if "metadata" in data.model_fields_set: + return (data.metadata or EMPTY_MAPPING).get(VECTOR_STORE_INDEX_GRANTS_METADATA_KEY) or None + return stored_grants + + +def enforce_vector_store_index_grants_are_admin_only( + data: BatchLimitRequest, + top_level_grants: Sequence[AllowedVectorStoreIndexItem] | None, + existing_metadata: Mapping[str, object] | None, + user_api_key_dict: UserAPIKeyAuth, + entity: Literal["key", "team"], +) -> None: + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: + return + stored_grants: Final = _stored_vector_store_index_grants(existing_metadata) + if _requested_vector_store_index_grants(data, top_level_grants, stored_grants) == stored_grants: + return + raise HTTPException( + status_code=403, + detail={ + "error": f"Only proxy admins can set {VECTOR_STORE_INDEX_GRANTS_METADATA_KEY} on a {entity}. " + "It grants read and write access to vector store indexes, so the holder cannot widen it." + }, + ) + + def get_model_rate_limit_from_metadata( user_api_key_dict: UserAPIKeyAuth, metadata_accessor_key: Literal["team_metadata", "organization_metadata", "project_metadata"], @@ -2124,7 +2166,7 @@ def _extract_models_from_managed_resource_id( decoded_video: Final = decode_video_id_with_provider(resource_id) _append_model_candidates( candidates=candidates, - value=_video_access_group(decoded_video, llm_router) or decoded_video.get("model_id"), + value=video_access_group(decoded_video, llm_router) or decoded_video.get("model_id"), ) else: model_id = decode_character_id_with_provider(resource_id).get("model_id") @@ -2138,7 +2180,9 @@ def _extract_models_from_managed_resource_id( return _dedupe_model_candidates(candidates) -def _video_access_group(decoded_video: DecodedVideoId, llm_router: Router | None) -> str | None: +def video_access_group( + decoded_video: DecodedVideoId, llm_router: Router | None, *, resolved_model: str | None = None +) -> str | None: """ The group a key must be allowed to call to use a video: the creating deployment's group, or, while router fallbacks skip model access checks, the group whose router fallback created it. @@ -2146,7 +2190,9 @@ def _video_access_group(decoded_video: DecodedVideoId, llm_router: Router | None from litellm.proxy.auth.fallback_model_access import router_fallback_access_check from litellm.proxy.video_endpoints.utils import video_model_group - creating_group: Final = video_model_group(decoded_video, llm_router) + creating_group: Final = ( + resolved_model if resolved_model is not None else video_model_group(decoded_video, llm_router) + ) requested_model: Final = decoded_video.get("requested_model") if llm_router is None or creating_group is None or requested_model is None: return creating_group diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index c1db8b7b4c7..7f2aa33246d 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -7,21 +7,27 @@ import functools import itertools import re from collections.abc import Callable, Collection, Mapping, Sequence +from collections.abc import Set as AbstractSet from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, replace from types import MappingProxyType from typing import Final, TypeVar import litellm from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.credential_accessor import CredentialAccessor +from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model from litellm.litellm_core_utils.pattern_match_deployments import PatternMatchRouter from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping from litellm.proxy._types import SpecialModelNames, UserAPIKeyAuth +from litellm.proxy.common_utils.model_listing_utils import CallerAliases, alias_map, dispatched_model +from litellm.proxy.common_utils.token_count_settings import cost_map_entry from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.responses.main import strip_responses_routing_prefix from litellm.router import Router -from litellm.router_api import fallback_chain -from litellm.types.router import CredentialLiteLLMParams, LiteLLM_Params +from litellm.router_api import RouterRegistry, fallback_chain +from litellm.router_rust.surface.registry import cost_map_backend_model +from litellm.types.router import CredentialLiteLLMParams, Deployment, LiteLLM_Params from litellm.types.utils import LlmProviders from litellm.utils import ( AZURE_MODEL_PLACEHOLDER, @@ -684,6 +690,133 @@ def _known_models_without_a_deployment(wildcard_model: str) -> Sequence[str]: return get_known_models_from_wildcard(wildcard_model=wildcard_model, litellm_params=None) +def _litellm_model(deployment: Mapping[str, object]) -> object: + litellm_params: Final = deployment.get("litellm_params") + return litellm_params.get("model") if is_str_keyed_mapping(litellm_params) else None + + +def _router_deployment(registry: RouterRegistry, model_id: str) -> Mapping[str, object] | None: + deployment: Final[object] = registry.deployment_dict(model_id) + return deployment if is_str_keyed_mapping(deployment) else None + + +GLOBAL_ALIASES_ONLY: Final = CallerAliases(team=None, key=None) + + +@dataclass(frozen=True, slots=True) +class RoutedElsewhere: + pass + + +ROUTED_ELSEWHERE: Final = RoutedElsewhere() + + +@dataclass(frozen=True, slots=True) +class OtherDeploymentsBill: + model: str + + +def wildcard_row_priced_model( + row: Mapping[str, object], llm_router: Router | None, caller_aliases: CallerAliases = GLOBAL_ALIASES_ONLY +) -> str | OtherDeploymentsBill | RoutedElsewhere | None: + model_id: Final = _model_info_field(row, "id") + listed_name: Final = row.get("model_name") + if llm_router is None or not isinstance(model_id, str) or not isinstance(listed_name, str): + return None + registry: Final = llm_router.registry() + source: Final = _router_deployment(registry, model_id) + if source is None: + return None + source_pattern: Final = source.get("model_name") + source_model: Final = _litellm_model(source) + if not isinstance(source_pattern, str) or not isinstance(source_model, str) or _litellm_model(row) == source_model: + return None + return wildcard_deployment_priced_model( + llm_router, + WildcardDeployment(model_id, source_pattern, source_model, _deployment_team_id(row)), + listed_name, + caller_aliases, + ) + + +@dataclass(frozen=True, slots=True) +class WildcardDeployment: + model_id: str + model_name_pattern: str + litellm_model: str + team_id: str | None + + +def wildcard_deployment_priced_model( + llm_router: Router, deployment: WildcardDeployment, listed_name: str, caller_aliases: CallerAliases +) -> str | OtherDeploymentsBill | RoutedElsewhere | None: + served_name: Final = _served_name(llm_router, listed_name, caller_aliases) + registry: Final = llm_router.registry() + team_id: Final = deployment.team_id or caller_aliases.team_id + routed_name: Final = routed_model_name(llm_router, served_name) + candidates: Final = registry.names(served_name, kind="route_candidates", team_id=team_id) + if deployment.model_id in candidates: + return _routed_deployment_model(deployment.model_name_pattern, routed_name, deployment.litellm_model) + if not candidates: + return None + return _model_other_deployments_bill(registry, routed_name, candidates, team_id) + + +def _model_other_deployments_bill( + registry: RouterRegistry, routed_name: str, candidates: AbstractSet[str], team_id: str | None +) -> OtherDeploymentsBill | RoutedElsewhere: + rows: Final = registry.deployments(routed_name, team_id=team_id) or () + billed: Final = frozenset( + cost_map_backend_model(Deployment.model_validate(row)) + for row in rows + if _model_info_field(row, "id") in candidates + ) + if len(billed) != 1 or any(_declares_its_own_rate(registry, model_id) for model_id in candidates): + return ROUTED_ELSEWHERE + return OtherDeploymentsBill(next(iter(billed))) + + +def _declares_its_own_rate(registry: RouterRegistry, model_id: str) -> bool: + deployment: Final = _router_deployment(registry, model_id) + if deployment is None: + return False + sections: Final = (deployment.get("litellm_params"), deployment.get("model_info")) + return any(is_str_keyed_mapping(section) and use_custom_pricing_for_model(dict(section)) for section in sections) + + +def wildcard_row_filled_model( + llm_router: Router, name_pattern: str, litellm_model: str, listed_name: str, caller_aliases: CallerAliases +) -> str: + served_name: Final = _served_name(llm_router, listed_name, caller_aliases) + return _routed_deployment_model(name_pattern, routed_model_name(llm_router, served_name), litellm_model) + + +def _served_name(llm_router: Router, listed_name: str, caller_aliases: CallerAliases) -> str: + return dispatched_model(listed_name, _without_deleted_team_targets(caller_aliases, llm_router.registry())) + + +def _without_deleted_team_targets(caller_aliases: CallerAliases, registry: RouterRegistry) -> CallerAliases: + deleted_prefix: Final = f"model_name_{caller_aliases.team_id}_" + live_names: Final = registry.index("model_name") + return replace( + caller_aliases, + team={ + alias: target + for alias, target in alias_map(caller_aliases.team).items() + if not target.startswith(deleted_prefix) or target in live_names + }, + ) + + +def routed_model_name(llm_router: Router, listed_name: str) -> str: + return llm_router.registry().resolve(listed_name, kind="alias") or listed_name + + +def model_info_cost_map_entry(model: str) -> Mapping[str, object]: + provider, _, _ = model.partition("/") + return cost_map_entry(model) or cost_map_entry(model.rsplit("/", 1)[-1], custom_llm_provider=provider) + + def _get_wildcard_models( unique_models: list[str], return_wildcard_routes: bool | None = False, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index a52ad14e293..0148cafa9e9 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -14,6 +14,7 @@ import secrets from collections.abc import Mapping, Sequence from datetime import datetime, timezone from enum import Enum +from functools import partial from types import MappingProxyType from typing import Final, NamedTuple, Protocol, Union, cast @@ -144,7 +145,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg _safe_get_request_headers, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export _safe_get_request_query_params, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export _safe_set_request_parsed_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export - is_opaque_audio_pass_through_request, + is_opaque_pass_through_request, populate_request_with_path_params, read_raw_json_body, read_request_body, @@ -170,7 +171,15 @@ from litellm.proxy.db.db_lookup_gate import bounded_db_lookup from litellm.proxy.db.db_span import db_span from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup -from litellm.proxy.spend_tracking.carried_budget_state import carry_team_and_user_budget_state +from litellm.proxy.spend_tracking.carried_budget_state import ( + carry_committed_budget_spend, + carry_team_and_user_budget_state, +) +from litellm.proxy.spend_tracking.committed_spend import ( + bind_committed_admission_counters, + committed_spend_tracked, + read_admission_committed_spend, +) from litellm.proxy.spend_tracking.spend_counter_batch import ( bind_admission_counter_keys, release_spend_counter_batch, @@ -2128,7 +2137,7 @@ async def _read_request_body_deferring_parse_failure( must run (resolving identity onto the request's trace) before the 400 goes out; the caller re-raises the returned exception once identity is seeded. """ - if is_opaque_audio_pass_through_request( + if is_opaque_pass_through_request( route=get_request_route(request=request), content_type=safe_get_request_headers(request=request).get("content-type", ""), ): @@ -3650,6 +3659,7 @@ async def run_centralized_common_checks( model_max_budget_limiter, prisma_client, proxy_logging_obj, + spend_counter_store, user_api_key_cache, user_custom_auth, ) @@ -3960,6 +3970,7 @@ async def run_centralized_common_checks( ) bind_admission_counter_keys(user_api_key_auth_obj, end_user_id=end_user_id) + bind_committed_admission_counters(user_api_key_auth_obj, end_user_id=end_user_id) try: _ = await common_checks( request=request, @@ -4035,6 +4046,18 @@ async def run_centralized_common_checks( user_obj=user_object, ) + if committed_spend_tracked(): + await carry_committed_budget_spend( + valid_token=user_api_key_auth_obj, + team_object=typed_team_object, + user_object=user_object, + end_user_object=end_user_object, + read_committed=( + None + if route_skips_budget_checks(route=route) + else partial(read_admission_committed_spend, spend_counter_store) + ), + ) await _reserve_budget_after_common_checks( user_api_key_auth_obj=user_api_key_auth_obj, request=request, @@ -4348,6 +4371,7 @@ async def _prefetch_referenced_auth_objects( """Warm every object and spend counter the checks below will read, in one MGET each (one DB query when cold). Runs after the key's model access check so a denied request costs no more than it did before.""" bind_admission_counter_keys(valid_token, end_user_id=end_user_id or None) + bind_committed_admission_counters(valid_token, end_user_id=end_user_id or None) await prefetch_auth_objects( refs=AuthObjectRefs.from_token(valid_token), user_api_key_cache=user_api_key_cache, diff --git a/litellm/proxy/client/README.md b/litellm/proxy/client/README.md index ceec636567f..9a7e067c792 100644 --- a/litellm/proxy/client/README.md +++ b/litellm/proxy/client/README.md @@ -380,7 +380,7 @@ The key itself, together with the refresh token that renews a `--pkce` credentia Keychain storage needs the `keyring` package, which ships with `pip install 'litellm[cli]'`. Headless boxes and CI runners usually have no keychain either. In all of those cases the key and the refresh token stay in the same `0600` file alongside the metadata, exactly as they did before, and `lite login` names which one applies: the package is missing, the machine has no keychain, or you set `LITELLM_CLI_DISABLE_KEYRING=1` to force the file even where a keychain exists. A `token.json` written by an older `lite` keeps working and is moved into the keychain, and scrubbed from the file, the first time a keychain-capable `lite` reads it. That includes a refresh token left behind by the release that moved only the key. -`lite logout` clears both stores. If the keychain is locked at that moment it says so, and re-running it once the keychain is unlocked finishes the job. +`lite logout` clears both stores. If the keychain is locked at that moment it says so, removes `token.json`, and leaves `~/.litellm/keychain_not_cleared` holding that login's metadata without its secret. Re-running it once the keychain is unlocked revokes a `--pkce` refresh token on the proxy and finishes the job. A `lite login` made once the keychain answers revokes that login too, since it replaces the entry. A login made while the keychain is still locked keeps its secret in `token.json` and the note names the keychain's login, so the first `lite logout` or `lite login` after the unlock revokes both. Other commands leave the noted entry where it is, and a silent renewal that takes over the keychain entry revokes that login too. On a machine with no keychain at all, every later `lite logout` repeats the warning until you delete `~/.litellm/keychain_not_cleared`. If `~/.litellm` will not take the note, `lite logout` says it did not log out and leaves `token.json` where it is, so keep that file, make the directory writable, and run it again. An older `lite` does not read the note, so finish that logout with this release or a later one: an older `lite logout` erases the keychain entry without revoking its `--pkce` refresh token, which then stays valid until it expires. A machine that an older `lite` logged out while its keychain was locked still has its `token.json`, and still reads as signed in, until you run `lite logout` once more. The stored credential is a short-lived, per-session agent token, not a managed virtual key. It is scoped to the user and team you logged in as and inherits their models and budgets; spend is tracked against the shared team and user budgets rather than a separate per-session cap, so multiple logins or several concurrent agents all draw down the same allowance. It is short-lived by design (default 24h, configurable via `LITELLM_CLI_JWT_EXPIRATION_HOURS`); re-run `lite login` to refresh it and pick up your latest team and user settings. `lite auth print-token` (usable as Claude Code's `apiKeyHelper`) prints it, and once it is past 24 hours (or the client's `LITELLM_CLI_JWT_EXPIRATION_HOURS`) asks the proxy first and fails if the proxy rejects it. There is no silent renewal. It is accepted on a default deployment without `EXPERIMENTAL_UI_LOGIN`, does not appear in the Keys UI, and cannot be rotated or revoked mid-session. A credential from `lite login --pkce` is the exception: it carries a refresh token, so the CLI renews the key shortly before it expires and `lite logout` revokes the refresh token on the proxy (see [Browser sign-in with PKCE](https://docs.litellm.ai/docs/proxy/cli_sso#browser-sign-in-with-pkce)). Only the holder can end a `--pkce` session early, with `lite logout`; an admin has no button for it, but every renewal re-reads the user on the proxy, so deactivating the user or removing them from the team makes the next renewal fail and the key runs out within `LITELLM_CLI_JWT_EXPIRATION_HOURS`. On a proxy with more than one worker or replica, configure Redis (`litellm_settings.cache` with Redis `cache_params`, or `general_settings.coordination_redis`) so a refresh token stays single-use and `lite logout` holds on every worker; without Redis each worker keeps its own record. For a long-lived, rotatable, Keys-UI-visible credential, create a dedicated virtual key in the dashboard and pass it via `--api-key` or `LITELLM_PROXY_API_KEY`. diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index e38f54adc4f..7bd209140ac 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -32,11 +32,13 @@ from litellm.litellm_core_utils.cli_token_utils import ( CredentialNotCleared, CredentialNotRecorded, CredentialNotSaved, + LoginRecordNotMoved, SecretSave, clear_cli_token, get_cli_token_file_path, is_cli_token_fresh, load_cli_token, + load_cli_token_left_in_keychain, save_cli_token, ) @@ -234,22 +236,45 @@ def load_token(*, vault: SecretVault = SYSTEM_KEYRING) -> Mapping[str, object] | return None if record is None else record.model_dump(exclude_none=True) +def _login_left_in_keychain(vault: SecretVault) -> Mapping[str, object] | None: + """The login an earlier logout left in a keychain it could not reach""" + left: Final = load_cli_token_left_in_keychain(vault=vault) + return None if left is None else left.model_dump(exclude_none=True) + + +def _same_refresh_token(left: Mapping[str, object] | None, current: Mapping[str, object] | None) -> bool: + return left is not None and current is not None and left.get("refresh_token") == current.get("refresh_token") + + +def _credentials_to_revoke(vault: SecretVault) -> tuple[Mapping[str, object], ...]: + left: Final = _login_left_in_keychain(vault) + current: Final = load_token(vault=vault) + if current is None or _same_refresh_token(left, current): + return () if left is None else (left,) + return (current,) if left is None else (left, current) + + def save_token(record: CliTokenData, *, vault: SecretVault = SYSTEM_KEYRING) -> SecretSave: """Store a credential the PKCE layer produced, secret in the vault and the rest on disk""" return save_cli_token(CliTokenRecord(**record), vault=vault) -def _renewal_saver(vault: SecretVault) -> Callable[[CliTokenData], None]: +def _renewal_saver(vault: SecretVault, http: Http) -> Callable[[CliTokenData], None]: """Persist a silently renewed credential, and say on stderr when no store would keep it. A renewal rotates the refresh token, so a rotation that is never stored logs this - machine out on the next command; the user hears about it rather than guessing. + machine out on the next command; the user hears about it rather than guessing. The save + can take the keychain entry of a login an earlier logout could not end, which leaves no + other record of that login's refresh token, so it is revoked the way a login revokes it. """ def save(record: CliTokenData) -> None: + left: Final = _login_left_in_keychain(vault) outcome: Final = save_token(record, vault=vault) if isinstance(outcome, (CredentialNotSaved, CredentialNotRecorded)): _warn(storage_notice(outcome)) + if left is not None: + _revoke_the_replaced_login(left, http, _warn) return save @@ -280,10 +305,11 @@ def get_stored_api_key( return None if expected_base_url is not None and token_data.get("base_url") != expected_base_url.rstrip("/"): return None + http: Final = requests.Session() return fresh_api_key( token_data, - _renewal_saver(vault), - requests.Session(), + _renewal_saver(vault, http), + http, reload=_renewal_reader(vault), warn=_warn, ) @@ -853,17 +879,22 @@ def _finish_login(base_url: str, api_key: str, config_claude: bool, stored: Secr def _replace_stored_token(record: CliTokenData, http: Http, vault: SecretVault) -> SecretSave: - previous: Final = load_token(vault=vault) + previous: Final = _credentials_to_revoke(vault) stored: Final = save_token(record, vault=vault) - if previous is None or isinstance(stored, CredentialNotSaved): + if isinstance(stored, CredentialNotSaved): return stored - revocation: Final = revoke_stored_credential(previous, http) + for login in previous: + _revoke_the_replaced_login(login, http, click.echo) + return stored + + +def _revoke_the_replaced_login(login: Mapping[str, object], http: Http, report: Callable[[str], None]) -> None: + revocation: Final = revoke_stored_credential(login, http) if revocation is not None: - click.echo( + report( f"Could not revoke the previous login's refresh token on the proxy ({revocation.reason}); " "it expires on its own." ) - return stored def _pkce_login(base_url: str, config_claude: bool, vault: SecretVault) -> None: @@ -987,13 +1018,22 @@ def login(ctx: click.Context, config_claude: bool, pkce: bool | None) -> None: return -@click.command(name="logout") -@click.pass_context -def logout(ctx: click.Context): - """Logout and clear stored authentication""" - vault: Final = context_secret_vault(ctx) - token_data: Final = load_token(vault=vault) - revocation: Final = revoke_stored_credential(token_data, requests.Session()) if token_data is not None else None +def _revoke_before_logout(left: Mapping[str, object] | None, current: Mapping[str, object] | None, http: Http) -> None: + """Revoke every login this logout ends, holding the logout back only for the one it signs out. + + An earlier login left in the keychain is the previous login once another one is current, so a + proxy that cannot revoke it only earns the warning `lite login` gives, never a logout that the + current proxy is ready to finish. + """ + if left is None or current is None or _same_refresh_token(left, current): + _revoke_the_login_or_keep_it(current if left is None else left, http) + return + _revoke_the_replaced_login(left, http, click.echo) + _revoke_the_login_or_keep_it(current, http) + + +def _revoke_the_login_or_keep_it(token_data: Mapping[str, object] | None, http: Http) -> None: + revocation: Final = None if token_data is None else revoke_stored_credential(token_data, http) match revocation: case RevocationUnavailable(reason=reason): raise click.ClickException( @@ -1007,6 +1047,17 @@ def logout(ctx: click.Context): case _: assert_never(revocation) + +@click.command(name="logout") +@click.pass_context +def logout(ctx: click.Context): + """Logout and clear stored authentication""" + vault: Final = context_secret_vault(ctx) + http: Final = requests.Session() + left: Final = _login_left_in_keychain(vault) + current: Final = load_token(vault=vault) + _revoke_before_logout(left, current, http) + path: Final = get_cli_token_file_path() match clear_cli_token(vault=vault): case SecretErased(): @@ -1014,6 +1065,12 @@ def logout(ctx: click.Context): case CredentialNotCleared(detail=detail): click.echo(f"Your credential is still in {path}, which could not be removed: {detail}.") click.echo("Delete that file, or make the directory writable and run 'lite logout' again.") + case LoginRecordNotMoved(detail=detail): + click.echo(f"Not logged out. {path} could not be moved aside: {detail}.") + click.echo( + "Keep that file, since 'lite logout' needs it to revoke a login your OS keychain may still hold. " + "Make its directory writable and run 'lite logout' again." + ) case SecretStranded(): click.echo(STRANDED_CREDENTIAL_MESSAGE) click.echo("Unlock your keychain and run 'lite logout' again to clear it.") @@ -1067,13 +1124,14 @@ def print_token(ctx: click.Context): click.echo(keychain_unreadable_notice(vault), err=True) sys.exit(1) + http: Final = requests.Session() api_key: Final = ( ctx_obj.get("api_key") if issued_for_this_server and ctx_obj.get("api_key_from_token_file") else fresh_api_key( token_data, - _renewal_saver(vault), - requests.Session(), + _renewal_saver(vault, http), + http, reload=_renewal_reader(vault), warn=_warn, ) diff --git a/litellm/proxy/client/cli/main.py b/litellm/proxy/client/cli/main.py index a1964d8a1bb..e0bcaaf5a66 100644 --- a/litellm/proxy/client/cli/main.py +++ b/litellm/proxy/client/cli/main.py @@ -199,7 +199,8 @@ def cli(ctx: click.Context, show_version: bool, base_url: str | None, api_key: s # If no API key provided via flag or environment variable, try to load from saved token. # Pass base_url so we only use the stored key when it was issued for this server. setup_command: Final = ctx.invoked_subcommand in ("configure", "reconfigure", "unconfigure") - api_key_from_token_file: Final = api_key is None and not setup_command + reads_the_stored_login_itself: Final = ctx.invoked_subcommand in ("login", "logout") + api_key_from_token_file: Final = api_key is None and not setup_command and not reads_the_stored_login_itself resolved_api_key: Final = ( get_stored_api_key(expected_base_url=base_url, vault=context_secret_vault(ctx)) if api_key_from_token_file diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 55f2ee94d0e..8f73d47d96b 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -70,8 +70,10 @@ from litellm.litellm_core_utils.core_helpers import ( has_attributes, independent_snapshot, is_expected_client_error, + proxy_request_headers, redact_nested_match_and_regex_keys, require_attributes, + stream_cost_injection_requested, ) from litellm.litellm_core_utils.dd_tracing import NullTracer, tracer from litellm.litellm_core_utils.get_supported_openai_params import ( @@ -172,6 +174,7 @@ _LateResponseT = TypeVar("_LateResponseT", bound=Response) _LlmCallT = TypeVar("_LlmCallT") _ResponseT: Final = TypeVar("_ResponseT") _ChunkT: Final = TypeVar("_ChunkT") +_ProcessedT: Final = TypeVar("_ProcessedT") KNOWN_PROXY_ROUTES: Final = frozenset( route for member in LiteLLMRoutes for route in member.value if route.startswith("/") @@ -286,6 +289,18 @@ if TYPE_CHECKING: ProxyConfig = _ProxyConfig else: ProxyConfig = Any +from litellm.litellm_core_utils.total_deadline import ( + TOTAL_TIMEOUT_FIELD, + TOTAL_TIMEOUT_HEADER, + InvalidTotalTimeout, + TotalTimeout, + TotalTimeoutExpired, + TotalTimeoutNotApplicable, + bounded_by_total_timeout, + current_total_deadline, + total_deadline, + total_timeout_from_request, +) from litellm.proxy.anthropic_endpoints.answered_status import with_answered_status from litellm.proxy.anthropic_endpoints.streaming_model_restamp import ( AnthropicStreamModelRestamper, @@ -581,6 +596,31 @@ async def _dispatch_disconnect_billing_event(logging_obj: LiteLLMLoggingObj, par return True +_TOTAL_TIMEOUT_ROUTE_TYPES: Final = frozenset({"acompletion", "aresponses", "anthropic_messages"}) + + +def _total_timeout_refused(message: str) -> ProxyException: + return ProxyException(message=message, type="invalid_request_error", param=TOTAL_TIMEOUT_FIELD, code=400) + + +def _total_timeout_expired() -> bool: + deadline: Final = current_total_deadline() + return deadline is not None and deadline.expired() + + +async def _cost_a_call_answered_past_its_total_timeout( + llm_call: "asyncio.Task[object]", logging_obj: LiteLLMLoggingObj +) -> None: + """The provider answered but the request ran out of time after it (a slower moderation check): the + failure it is logged as still carries what the provider charged.""" + if not llm_call.done() or llm_call.cancelled() or llm_call.exception() is not None: + return + response: Final = llm_call.result() + prepare_response_cost(logging_obj, response, datetime.now(timezone.utc)) + with anyio.CancelScope(shield=True): + await finalize_proxy_response_cost(logging_obj, response) + + async def _cancel_pending_gather_tasks(tasks: Sequence["asyncio.Task[object]"]) -> None: pending_tasks: Final = [task for task in tasks if not task.done()] for task in pending_tasks: @@ -3126,33 +3166,88 @@ class ProxyBaseLLMRequestProcessing: ) return await open_sse_before_first_byte( - self._process_llm_request( + self._within_total_timeout( request=request, - fastapi_response=fastapi_response, - user_api_key_dict=user_api_key_dict, route_type=route_type, - proxy_logging_obj=proxy_logging_obj, - general_settings=general_settings, - proxy_config=proxy_config, - select_data_generator=select_data_generator, - llm_router=llm_router, - model=model, - user_model=user_model, - user_temperature=user_temperature, - user_request_timeout=user_request_timeout, - user_max_tokens=user_max_tokens, - user_api_base=user_api_base, - version=version, is_streaming_request=is_streaming_request, - contents=contents, - skip_pre_call_logic=skip_pre_call_logic, - first_chunk_framing=sse_framing.first_chunk, + processing=self._process_llm_request( + request=request, + fastapi_response=fastapi_response, + user_api_key_dict=user_api_key_dict, + route_type=route_type, + proxy_logging_obj=proxy_logging_obj, + general_settings=general_settings, + proxy_config=proxy_config, + select_data_generator=select_data_generator, + llm_router=llm_router, + model=model, + user_model=user_model, + user_temperature=user_temperature, + user_request_timeout=user_request_timeout, + user_max_tokens=user_max_tokens, + user_api_base=user_api_base, + version=version, + is_streaming_request=is_streaming_request, + contents=contents, + skip_pre_call_logic=skip_pre_call_logic, + first_chunk_framing=sse_framing.first_chunk, + ), ), ping_interval_seconds=ttft_keepalive_interval(self.data, llm_router), on_late_failure=_audit_late_failure, late_framing=ANTHROPIC_LATE_SSE_FRAMING if answers_as_anthropic else sse_framing.late, ) + def total_timeout_for( + self, request: Request, route_type: ProxyRouteType, is_streaming_request: bool | None + ) -> TotalTimeout | None: + """The request's ``total_timeout`` when it bounds this request: a body one on a request it cannot + bound is refused, a header one is ignored.""" + body: Final = _REQUEST_FIELDS.validate_python(self.data) # pyright: ignore[reportUnknownMemberType] # self.data is the untyped request dict + header: Final = request.headers.get(TOTAL_TIMEOUT_HEADER) + self.data.pop(TOTAL_TIMEOUT_FIELD, None) # pyright: ignore[reportUnknownMemberType] # self.data is the untyped request dict + bounded: Final = ( + route_type in _TOTAL_TIMEOUT_ROUTE_TYPES + and is_streaming_request is not True + and bounded_by_total_timeout(body) + ) + if header is not None and not bounded: + verbose_proxy_logger.debug("%s ignored for this %s request", TOTAL_TIMEOUT_HEADER, route_type) + return None + parsed: Final = total_timeout_from_request(header, body) + if parsed is None: + return None + if isinstance(parsed, InvalidTotalTimeout): + raise _total_timeout_refused(parsed.message) + if not bounded: + raise _total_timeout_refused( + f"{TOTAL_TIMEOUT_FIELD} applies only to non-streaming, non-background " + "/v1/chat/completions, /v1/messages and /v1/responses requests" + ) + return parsed + + async def _within_total_timeout( + self, + request: Request, + route_type: ProxyRouteType, + is_streaming_request: bool | None, + processing: Coroutine[object, object, _ProcessedT], + ) -> _ProcessedT: + try: + total_timeout: Final = self.total_timeout_for(request, route_type, is_streaming_request) + except ProxyException: + processing.close() + raise + if total_timeout is None: + return await processing + try: + async with total_deadline(total_timeout.seconds, from_body=total_timeout.from_body): + return await processing + except TotalTimeoutExpired as expired: + raise ProxyException(message=str(expired), type="timeout", param=TOTAL_TIMEOUT_FIELD, code=408) from expired + except TotalTimeoutNotApplicable as refused: + raise _total_timeout_refused(str(refused)) from refused + async def _process_llm_request( self, request: Request, @@ -3236,25 +3331,32 @@ class ProxyBaseLLMRequestProcessing: if contents: self.data["contents"] = contents - ### ROUTE THE REQUEST ### - # Do not change this - it should be a constant time fetch - ALWAYS - llm_call: Final = await route_request( - data=self.data, - route_type=route_type, - llm_router=llm_router, - user_model=user_model, - user_api_key_dict=user_api_key_dict, - ) - llm_call_task: Final = asyncio.create_task(with_proxy_compaction_executor(llm_call, request)) - tasks.append(llm_call_task) - - llm_responses: Final = asyncio.gather(*tasks) # run the moderation check in parallel to the actual llm api call - try: - if general_settings.get("cancel_on_disconnect", False): - responses = await _await_llm_call_cancelling_on_disconnect(request, llm_responses) - else: - responses = await llm_responses + ### ROUTE THE REQUEST ### + # Do not change this - it should be a constant time fetch - ALWAYS + llm_call: Final = await route_request( + data=self.data, + route_type=route_type, + llm_router=llm_router, + user_model=user_model, + user_api_key_dict=user_api_key_dict, + ) + llm_call_task: Final = asyncio.create_task(with_proxy_compaction_executor(llm_call, request)) + tasks.append(llm_call_task) + + llm_responses: Final = asyncio.gather( + *tasks + ) # run the moderation check in parallel to the actual llm api call + + try: + if general_settings.get("cancel_on_disconnect", False): + responses = await _await_llm_call_cancelling_on_disconnect(request, llm_responses) + else: + responses = await llm_responses + except asyncio.CancelledError: + if _total_timeout_expired(): + await _cost_a_call_answered_past_its_total_timeout(llm_call_task, logging_obj) + raise finally: await _cancel_pending_gather_tasks(tasks) @@ -3524,7 +3626,7 @@ class ProxyBaseLLMRequestProcessing: _guardrail_entries_for_headers(get_or_create_metadata_bucket(self.data)[1], logging_obj) ), ) - if isinstance(exc, asyncio.CancelledError): + if isinstance(exc, asyncio.CancelledError) and not _total_timeout_expired(): await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=HTTPException(status_code=499, detail=_CLIENT_DISCONNECT_DETAIL), @@ -4815,7 +4917,10 @@ class ProxyBaseLLMRequestProcessing: # await, response-string materialization, and cost-injection call are # pure overhead on the streaming hot path (the default config). caps: Final = ProxyLogging.callback_capabilities() - cost_injection_enabled: Final = bool(getattr(litellm, "include_cost_in_streaming_usage", False)) + request_headers: Final = proxy_request_headers( + request_data # pyright: ignore[reportUnknownArgumentType] # request_data is the untyped proxy request dict + ) + cost_injection_enabled: Final = stream_cost_injection_requested(request_headers) fast_path = not caps.has_streaming_chunk_override and not caps.has_guardrail and not cost_injection_enabled debug_enabled: Final = verbose_proxy_logger.isEnabledFor(logging.DEBUG) stream_completed = False @@ -4862,7 +4967,7 @@ class ProxyBaseLLMRequestProcessing: model_name = request_data.get("model", "") chunk = ProxyBaseLLMRequestProcessing.process_chunk_with_cost_injection( - chunk, model_name, request_data.get("litellm_logging_obj") + chunk, model_name, request_data.get("litellm_logging_obj"), request_headers=request_headers ) # Set before the yield: an async generator suspends at the yield, @@ -4991,18 +5096,30 @@ class ProxyBaseLLMRequestProcessing: @overload @staticmethod def process_chunk_with_cost_injection( - chunk: bytes, model_name: str, litellm_logging_obj: LiteLLMLoggingObj | None = None + chunk: bytes, + model_name: str, + litellm_logging_obj: LiteLLMLoggingObj | None = None, + *, + request_headers: Mapping[str, object] | None = None, ) -> bytes: ... @overload @staticmethod def process_chunk_with_cost_injection( - chunk: object, model_name: str, litellm_logging_obj: LiteLLMLoggingObj | None = None + chunk: object, + model_name: str, + litellm_logging_obj: LiteLLMLoggingObj | None = None, + *, + request_headers: Mapping[str, object] | None = None, ) -> object: ... @staticmethod def process_chunk_with_cost_injection( - chunk: object, model_name: str, litellm_logging_obj: LiteLLMLoggingObj | None = None + chunk: object, + model_name: str, + litellm_logging_obj: LiteLLMLoggingObj | None = None, + *, + request_headers: Mapping[str, object] | None = None, ) -> object: """ Process a streaming chunk and inject cost information if enabled. @@ -5011,11 +5128,12 @@ class ProxyBaseLLMRequestProcessing: chunk: The streaming chunk (dict, str, bytes, or bytearray) model_name: Model name for cost calculation litellm_logging_obj: The call's logging object, used for pricing + request_headers: The proxy request headers, read for the per-request opt-in header Returns: The processed chunk with cost information injected if applicable """ - if not getattr(litellm, "include_cost_in_streaming_usage", False): + if not stream_cost_injection_requested(request_headers): return chunk try: diff --git a/litellm/proxy/common_utils/encrypt_decrypt_utils.py b/litellm/proxy/common_utils/encrypt_decrypt_utils.py index e5f6bc5dda8..644a35ef4ef 100644 --- a/litellm/proxy/common_utils/encrypt_decrypt_utils.py +++ b/litellm/proxy/common_utils/encrypt_decrypt_utils.py @@ -1,11 +1,14 @@ import base64 +import ctypes import os +import re from collections.abc import Mapping from typing import Final, Literal, cast from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger +from litellm.proxy.auth.master_key_boot_check import SALT_KEY_ENV_VAR # Versioned ciphertext marker for AES-256-GCM values. # Format: "v2:gcm:" + base64url(nonce(12) || ciphertext || tag(16)). @@ -42,6 +45,72 @@ def get_salt_key() -> str | None: _get_salt_key: Final = get_salt_key +_SALT_KEY_HEX: Final = re.compile(r"(?:sk-)?((?:[0-9a-fA-F]{2}){32,})") +_OSSL_PARAM_UTF8_STRING: Final = 4 +_OSSL_PARAM_OCTET_STRING: Final = 5 +_OSSL_PARAM_UNMODIFIED: Final = ctypes.c_size_t(-1).value +_DERIVED_KEY_BYTES: Final = 32 +_KBKDF_PROPERTIES: Final = b"fips=yes" + + +class _OsslParam(ctypes.Structure): + _fields_ = ( + ("key", ctypes.c_char_p), + ("data_type", ctypes.c_uint), + ("data", ctypes.c_char_p), + ("data_size", ctypes.c_size_t), + ("return_size", ctypes.c_size_t), + ) + + +def _param(name: bytes, data_type: int, data: bytes) -> _OsslParam: + return _OsslParam(name, data_type, data, len(data), _OSSL_PARAM_UNMODIFIED) + + +def _kbkdf_hmac_sha256(key: bytes, label: bytes, context: bytes) -> bytes: + import _hashlib + + libcrypto: Final = ctypes.CDLL(_hashlib.__file__) + kdf_fetch: Final = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.c_void_p, ctypes.c_char_p, ctypes.c_char_p)( + ("EVP_KDF_fetch", libcrypto) + ) + ctx_new: Final = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.c_void_p)(("EVP_KDF_CTX_new", libcrypto)) + derive: Final = ctypes.CFUNCTYPE( + ctypes.c_int, ctypes.c_void_p, ctypes.c_char_p, ctypes.c_size_t, ctypes.POINTER(_OsslParam) + )(("EVP_KDF_derive", libcrypto)) + ctx_free: Final = ctypes.CFUNCTYPE(None, ctypes.c_void_p)(("EVP_KDF_CTX_free", libcrypto)) + kdf_free: Final = ctypes.CFUNCTYPE(None, ctypes.c_void_p)(("EVP_KDF_free", libcrypto)) + params: Final = (_OsslParam * 6)( + _param(b"mac", _OSSL_PARAM_UTF8_STRING, b"HMAC"), + _param(b"digest", _OSSL_PARAM_UTF8_STRING, b"SHA256"), + _param(b"key", _OSSL_PARAM_OCTET_STRING, key), + _param(b"salt", _OSSL_PARAM_OCTET_STRING, label), + _param(b"info", _OSSL_PARAM_OCTET_STRING, context), + ) + derived: Final = ctypes.create_string_buffer(_DERIVED_KEY_BYTES) + kdf: Final = kdf_fetch(None, b"KBKDF", _KBKDF_PROPERTIES) + ctx: Final = ctx_new(kdf) + try: + if not ctx or derive(ctx, derived, _DERIVED_KEY_BYTES, params) != 1: + raise ValueError("OpenSSL could not run KBKDF") + return derived.raw + finally: + ctx_free(ctx) + kdf_free(kdf) + + +def derive_key_from_salt_key(label: bytes, context: bytes) -> str: + """Hex of 32 bytes derived from LITELLM_SALT_KEY under label and context by SP 800-108r1 KBKDF in counter mode + with HMAC-SHA256, run by EVP_KDF from the FIPS provider of the OpenSSL that hashlib uses. Raises ValueError when + LITELLM_SALT_KEY is not at least 32 bytes of hex, optionally prefixed with sk-, or when no FIPS provider runs it.""" + salt_key: Final = _SALT_KEY_HEX.fullmatch(os.getenv(SALT_KEY_ENV_VAR, "")) + if salt_key is None: + raise ValueError( + f"With general_settings.fips_mode on, {SALT_KEY_ENV_VAR} must be at least 32 bytes of hex, such as " + "sk-$(openssl rand -hex 32)" + ) + return _kbkdf_hmac_sha256(bytes.fromhex(salt_key.group(1)), label, context).hex() + def _get_encryption_algorithm() -> str: """ diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index 08461aa3a44..80ba2f49745 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -19,6 +19,7 @@ from litellm.constants import ( LITELLM_CLIENT_DISCONNECT_DETAIL, LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, MAX_REQUEST_BODY_SIZE_TO_REPAIR_MB, + OPENSEARCH_PASS_THROUGH_ROUTE_PREFIX, ) from litellm.integrations.otel.runtime import phase_event from litellm.litellm_core_utils.core_helpers import as_object_list, has_attributes, require_attributes @@ -404,8 +405,10 @@ async def read_request_body(request: Request | None) -> dict: _read_request_body: Final = read_request_body -def is_opaque_audio_pass_through_request(route: str, content_type: str) -> bool: - """Azure Speech bodies (raw audio, multipart uploads) are forwarded byte for byte, so auth must not consume them.""" +def is_opaque_pass_through_request(route: str, content_type: str) -> bool: + """Azure Speech audio and every OpenSearch body are forwarded byte for byte, so auth must not parse them.""" + if route.startswith(f"{OPENSEARCH_PASS_THROUGH_ROUTE_PREFIX}/"): + return True media_type: Final = _normalize_media_type(content_type) return route.startswith(f"{AZURE_SPEECH_PASS_THROUGH_ROUTE_PREFIX}/") and ( media_type.startswith("audio/") or media_type == "multipart/form-data" diff --git a/litellm/proxy/common_utils/load_config_utils.py b/litellm/proxy/common_utils/load_config_utils.py index 640a947f9c5..d6c639354bd 100644 --- a/litellm/proxy/common_utils/load_config_utils.py +++ b/litellm/proxy/common_utils/load_config_utils.py @@ -3,12 +3,14 @@ import os import posixpath from collections.abc import Awaitable, Mapping from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Protocol +from typing import TYPE_CHECKING, ClassVar, Final, Protocol import yaml from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.web_identity_session_policy import S3_OBJECT_READ_SESSION_POLICY, WebIdentitySessionPolicy from litellm.proxy.common_utils.config_includes import resolve_includes if TYPE_CHECKING: @@ -18,6 +20,10 @@ _BUCKET_CONFIG_ADAPTER: Final = TypeAdapter(dict[str, object]) _DECODED_YAML: Final = TypeAdapter(object) +class _ConfigBucketAWS(BaseAWSLLM): + _web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = S3_OBJECT_READ_SESSION_POLICY + + class BucketObjectFetcher(Protocol): def __call__(self, object_key: str, /) -> Awaitable[Mapping[str, object] | None]: ... @@ -48,9 +54,7 @@ def s3_object_reader(bucket_name: str) -> SyncBucketObjectReader: import boto3 from botocore.credentials import Credentials - from litellm.main import bedrock_converse_chat_completion - - credentials: Final[Credentials] = bedrock_converse_chat_completion.get_credentials() + credentials: Final[Credentials] = _ConfigBucketAWS().get_credentials() s3_client: Final = boto3.client( "s3", aws_access_key_id=credentials.access_key, @@ -219,11 +223,7 @@ def download_python_file_from_s3( import boto3 from botocore.credentials import Credentials - from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM - - base_aws_llm: Final = BaseAWSLLM() - - credentials: Final[Credentials] = base_aws_llm.get_credentials() + credentials: Final[Credentials] = _ConfigBucketAWS().get_credentials() s3_client: Final = boto3.client( "s3", aws_access_key_id=credentials.access_key, diff --git a/litellm/proxy/common_utils/model_listing_utils.py b/litellm/proxy/common_utils/model_listing_utils.py index 9db67aa3808..86b444663f0 100644 --- a/litellm/proxy/common_utils/model_listing_utils.py +++ b/litellm/proxy/common_utils/model_listing_utils.py @@ -160,8 +160,8 @@ class ClaudeCodeRoutingNames: @dataclass(frozen=True, slots=True) class CallerAliases: - """The caller's `team` and `key` alias maps, the names `/v1/models` lists for it, and `router`, - its key's or team's `router_settings.model_group_alias`. + """The caller's `team` and `key` alias maps, the names `/v1/models` lists for it, `router`, + its key's or team's `router_settings.model_group_alias`, and `team_id`, the key's team. `rewrite` are the maps `/chat/completions` rewrites its model through after auth's `router` rewrite, in the order it applies them: the team's, the key's in `add_litellm_data_to_request`, then the global `model_alias_map` and the key's again in `common_processing_pre_call_logic`, @@ -170,6 +170,7 @@ class CallerAliases: team: object key: object router: object = None + team_id: str | None = None @property def own(self) -> tuple[object, ...]: @@ -214,7 +215,7 @@ def caller_alias_maps( """Team aliases count only when listing the team the key authenticated as.""" if listed_team_id is not None and listed_team_id != key_team_id: return CallerAliases(None, key_aliases, router_aliases) - return CallerAliases(team_aliases, key_aliases, router_aliases) + return CallerAliases(team_aliases, key_aliases, router_aliases, key_team_id) def alias_map(aliases: object) -> Mapping[str, str]: diff --git a/litellm/proxy/common_utils/registry_read_through.py b/litellm/proxy/common_utils/registry_read_through.py index 487c6feb68d..d58c58fa1a5 100644 --- a/litellm/proxy/common_utils/registry_read_through.py +++ b/litellm/proxy/common_utils/registry_read_through.py @@ -133,7 +133,7 @@ async def _resync_model_deployments(model_name: str) -> bool: rows: Final = await _find_rows_by_name_or_id(model_name, ModelRepository(prisma_client).table) if not rows: return False - async with proxy_server.MODEL_RECONCILE_LOCK: + async with proxy_server.DB_RECONCILE_LOCK: outcome: Final = await _install_model_rows(rows, _always) if outcome in ("no_router", "auto_router"): await _reload_every_model() @@ -150,7 +150,7 @@ async def resync_model_deployments_by_ids(model_ids: tuple[str, ...]) -> None: prisma_client: Final = proxy_server.prisma_client assert prisma_client is not None writer_table: Final = ModelRepository(WriterPinnedClient(prisma_client.db)).table - async with proxy_server.MODEL_RECONCILE_LOCK: + async with proxy_server.DB_RECONCILE_LOCK: rows: Final = await _find_rows_by_ids(model_ids, writer_table) _ = await _install_model_rows(_without_auto_routers(rows), _names_an_unloaded_credential) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 131f57568e8..9c06e3afc18 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -48,6 +48,7 @@ from litellm.proxy.db.budget_window_spend_writer import roll_window_spend_row from litellm.proxy.db.db_span import db_span, db_spanned from litellm.proxy.db.db_transaction_queue.pod_lock_manager import POD_LOCK_TARGET, PodLockManager from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry +from litellm.proxy.spend_tracking.committed_spend import committed_spend_keys from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.redis.stores.spend_counters import ( SPEND_COUNTERS_TARGET, @@ -621,7 +622,7 @@ class ResetBudgetJob: from litellm.proxy.proxy_server import spend_counter_store try: - await spend_counter_store.delete((counter_key,)) + await spend_counter_store.delete((counter_key, *committed_spend_keys((counter_key,)))) except Exception as e: verbose_proxy_logger.warning( "Failed to reset spend counter %s in Redis: %s. Budget may be over-enforced until counter expires.", @@ -680,7 +681,7 @@ class ResetBudgetJob: match cache: case "spend counter": - await spend_counter_store.delete(keys) + await spend_counter_store.delete((*keys, *committed_spend_keys(keys))) case "user_api_key_cache": await user_api_key_cache.async_delete_cache_keys(keys) case _: diff --git a/litellm/proxy/credential_endpoints/endpoints.py b/litellm/proxy/credential_endpoints/endpoints.py index 86e716f19e3..3026f2706ac 100644 --- a/litellm/proxy/credential_endpoints/endpoints.py +++ b/litellm/proxy/credential_endpoints/endpoints.py @@ -324,7 +324,7 @@ async def create_credential( Stores credential in DB. Reloads credentials in memory. """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import DB_RECONCILE_LOCK, prisma_client try: if prisma_client is None: @@ -365,21 +365,22 @@ async def create_credential( credentials_dict_jsonified: Final = cast( # cast-ok: deep-copies a model_dump, so keys are str "dict[str, object]", jsonify_object(credentials_dict) ) - try: - await CredentialsRepository(prisma_client).create( - data={ - **credentials_dict_jsonified, - "created_by": user_api_key_dict.user_id, - "updated_by": user_api_key_dict.user_id, - } - ) - except Exception as e: - if not is_unique_violation(e): - raise - raise HTTPException(status_code=409, detail=_credential_exists_detail(credential.credential_name)) + async with DB_RECONCILE_LOCK: + try: + await CredentialsRepository(prisma_client).create( + data={ + **credentials_dict_jsonified, + "created_by": user_api_key_dict.user_id, + "updated_by": user_api_key_dict.user_id, + } + ) + except Exception as e: + if not is_unique_violation(e): + raise + raise HTTPException(status_code=409, detail=_credential_exists_detail(credential.credential_name)) - ## ADD TO LITELLM ## - CredentialAccessor.upsert_credentials([processed_credential]) + ## ADD TO LITELLM ## + CredentialAccessor.upsert_credentials([processed_credential]) return {"success": True, "message": "Credential created successfully"} except Exception as e: @@ -865,7 +866,7 @@ async def delete_credential( """ [BETA] endpoint. This might change unexpectedly. """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import DB_RECONCILE_LOCK, prisma_client try: if prisma_client is None: @@ -876,14 +877,17 @@ async def delete_credential( _reject_non_admin_wif_fields( await named_credential_wif_fields(credential_name, prisma_client), user_api_key_dict ) - deleted: Final = await CredentialsRepository(prisma_client).delete_by_name(credential_name) - if deleted is None: - raise _not_found_unless_config_defined( - credential_name, "Credential not found. Got credential name: " + credential_name - ) + async with DB_RECONCILE_LOCK: + deleted: Final = await CredentialsRepository(prisma_client).delete_by_name(credential_name) + if deleted is None: + raise _not_found_unless_config_defined( + credential_name, "Credential not found. Got credential name: " + credential_name + ) - ## DELETE FROM LITELLM ## - litellm.credential_list = [cred for cred in litellm.credential_list if cred.credential_name != credential_name] + ## DELETE FROM LITELLM ## + litellm.credential_list = [ + cred for cred in litellm.credential_list if cred.credential_name != credential_name + ] await _purge_user_connections_for_credential(credential_name) return {"success": True, "message": "Credential deleted successfully"} except Exception as e: @@ -983,7 +987,7 @@ async def update_credential( """ [BETA] endpoint. This might change unexpectedly. """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import DB_RECONCILE_LOCK, prisma_client try: if credential.credential_name and credential.credential_name != credential_name: @@ -1007,53 +1011,55 @@ async def update_credential( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) credentials_repository: Final = CredentialsRepository(WriterPinnedClient(prisma_client.db)) - db_credential: Final = await credentials_repository.find_by_name(credential_name) - if db_credential is None: - raise _not_found_unless_config_defined(credential_name, "Credential not found in DB.") - _reject_non_admin_wif_fields(_stored_wif_fields(db_credential), user_api_key_dict) - decrypted_stored_credential: Final = decrypted_credential(db_credential) - decrypted_values: Final = _CREDENTIAL_DICT_ADAPTER.validate_python( - decrypted_stored_credential.credential_values - ) - incoming_values: Final = _without_masked_read_backs(requested_values, _stored_objects(decrypted_values)) - patch: Final = CredentialItem( - credential_name=credential_name, - display_name=( - requested_display_name if "display_name" in credential.model_fields_set else db_credential.display_name - ), - credential_info=_CREDENTIAL_DICT_ADAPTER.validate_python(credential.credential_info), - credential_values=_CREDENTIAL_DICT_ADAPTER.validate_python(incoming_values), - credential_values_to_delete=credential.credential_values_to_delete, - ) - stored_info: Final = _CREDENTIAL_DICT_ADAPTER.validate_python( - db_credential.credential_info # pyright: ignore[reportUnknownMemberType] # untyped credential field - ) - patch_info: Final = _CREDENTIAL_DICT_ADAPTER.validate_python( - patch.credential_info # pyright: ignore[reportUnknownMemberType] # untyped credential field - ) - effective_info: Final = {**stored_info, **patch_info} - normalized_patch: Final = ( - patch.model_copy( - update={"credential_values": _normalized_github_copilot_values(incoming_values, effective_info)} + async with DB_RECONCILE_LOCK: + db_credential: Final = await credentials_repository.find_by_name(credential_name) + if db_credential is None: + raise _not_found_unless_config_defined(credential_name, "Credential not found in DB.") + _reject_non_admin_wif_fields(_stored_wif_fields(db_credential), user_api_key_dict) + decrypted_stored_credential: Final = decrypted_credential(db_credential) + decrypted_values: Final = _CREDENTIAL_DICT_ADAPTER.validate_python( + decrypted_stored_credential.credential_values ) - if incoming_values.get(GITHUB_COPILOT_HOST_KEY) is not None - else patch - ) - merged_credential: Final = update_db_credential(db_credential, normalized_patch) - credential_object_jsonified: Final = cast( # cast-ok: deep-copies a model_dump, so keys are str - "dict[str, object]", jsonify_object(merged_credential.model_dump(exclude_none=True)) - ) - await credentials_repository.update_by_name( - credential_name, - data={ - **credential_object_jsonified, - "display_name": merged_credential.display_name, - "updated_by": user_api_key_dict.user_id, - }, - ) - - # Sync in-memory credential_list (skip if not in memory - e.g., proxy restarted) - _sync_in_memory_credential(normalized_patch, credential_name) + incoming_values: Final = _without_masked_read_backs(requested_values, _stored_objects(decrypted_values)) + patch: Final = CredentialItem( + credential_name=credential_name, + display_name=( + requested_display_name + if "display_name" in credential.model_fields_set + else db_credential.display_name + ), + credential_info=_CREDENTIAL_DICT_ADAPTER.validate_python(credential.credential_info), + credential_values=_CREDENTIAL_DICT_ADAPTER.validate_python(incoming_values), + credential_values_to_delete=credential.credential_values_to_delete, + ) + stored_info: Final = _CREDENTIAL_DICT_ADAPTER.validate_python( + db_credential.credential_info # pyright: ignore[reportUnknownMemberType] # untyped credential field + ) + patch_info: Final = _CREDENTIAL_DICT_ADAPTER.validate_python( + patch.credential_info # pyright: ignore[reportUnknownMemberType] # untyped credential field + ) + effective_info: Final = {**stored_info, **patch_info} + normalized_patch: Final = ( + patch.model_copy( + update={"credential_values": _normalized_github_copilot_values(incoming_values, effective_info)} + ) + if incoming_values.get(GITHUB_COPILOT_HOST_KEY) is not None + else patch + ) + merged_credential: Final = update_db_credential(db_credential, normalized_patch) + credential_object_jsonified: Final = cast( # cast-ok: deep-copies a model_dump, so keys are str + "dict[str, object]", jsonify_object(merged_credential.model_dump(exclude_none=True)) + ) + await credentials_repository.update_by_name( + credential_name, + data={ + **credential_object_jsonified, + "display_name": merged_credential.display_name, + "updated_by": user_api_key_dict.user_id, + }, + ) + # Sync in-memory credential_list (skip if not in memory - e.g., proxy restarted) + _sync_in_memory_credential(normalized_patch, credential_name) from litellm.llms.github_copilot.github_host import normalized_github_host diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index ff57596232f..4a1871ff1f5 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -28,6 +28,7 @@ from typing_extensions import LiteralString, ReadOnly, TypedDict import litellm from litellm._internal_context import service_target, with_service_target from litellm._logging import verbose_proxy_logger +from litellm.accounting.metrics import Cadence from litellm.caching import RedisCache from litellm.constants import ( BILLED_RETRY_ATTEMPT_METADATA_KEY, @@ -682,8 +683,11 @@ class DBSpendUpdateWriter: def __init__( self, redis_cache: RedisCache | None = None, + accounting_clock: Callable[[], float] = time.monotonic, ): self.redis_cache = redis_cache + self.accounting_storage_checks = Cadence(accounting_clock) + self.accounting_delay_alerts = Cadence(accounting_clock) self.redis_update_buffer = RedisUpdateBuffer(redis_cache=self.redis_cache) self.pod_lock_manager = PodLockManager() self.spend_update_queue = SpendUpdateQueue() diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 4ed3ac30be7..ae45b43f15f 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -94,6 +94,10 @@ class SpendCounterReseed: _locks: ClassVar["OrderedDict[str, asyncio.Lock]"] = OrderedDict() _registry_lock: ClassVar[asyncio.Lock | None] = None + @staticmethod + async def counter_lock(counter_key: str) -> asyncio.Lock: + return await SpendCounterReseed._get_lock(counter_key) + @staticmethod async def _get_lock(counter_key: str) -> asyncio.Lock: if SpendCounterReseed._registry_lock is None: diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py index b1e3d51bbc1..6f077650a92 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py @@ -6,15 +6,17 @@ to detect and block/mask sensitive content. """ import asyncio +import bisect import contextvars import itertools import json import os import re import time -from collections.abc import AsyncGenerator, AsyncIterable, Callable, Coroutine, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterable, Callable, Coroutine, Iterable, Mapping, Sequence from dataclasses import dataclass, replace from datetime import datetime +from functools import partial from re import Pattern from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict, TypeVar, cast @@ -79,7 +81,7 @@ from .competitor_intent import ( AirlineCompetitorIntentChecker, BaseCompetitorIntentChecker, ) -from .patterns import PATTERN_EXTRA_CONFIG, get_compiled_pattern +from .patterns import PATTERN_EXTRA_CONFIG, PATTERN_STREAMING_HOLDBACK_CHARS, get_compiled_pattern MAX_KEYWORD_VALUE_GAP_WORDS: Final = 1 _DECODED_ARGUMENTS: Final = TypeAdapter(object) @@ -223,10 +225,70 @@ async def _stream_then_end_of_stream_chunk( @dataclass(frozen=True, slots=True) class _StreamedScanPlan: context_chars: int + holdback_chars: int exception_phrases: tuple[str, ...] conditional_words: tuple[str, ...] +@dataclass(frozen=True, slots=True) +class _RedactionSpan: + start: int + end: int + pattern_index: int + pattern_name: str + + +def _merge_redaction_spans(spans: Iterable[_RedactionSpan]) -> tuple[_RedactionSpan, ...]: + ordered: Final = sorted(spans, key=lambda span: (span.start, -span.end)) + reach: Final = tuple(itertools.accumulate((span.end for span in ordered), max)) + first_indices: Final = tuple( + index for index, span in enumerate(ordered) if index == 0 or span.start >= reach[index - 1] + ) + last_indices: Final = (*(index - 1 for index in first_indices[1:]), len(ordered) - 1) + return tuple(replace(ordered[first], end=reach[last]) for first, last in zip(first_indices, last_indices)) + + +@dataclass(frozen=True, slots=True) +class _RedactedText: + spans: tuple[_RedactionSpan, ...] + tags: tuple[str, ...] + tag_starts: tuple[int, ...] + + def render(self, text: str) -> str: + previous_ends: Final = (0, *(span.end for span in self.spans)) + kept_parts: Final = ( + text[previous_end : span.start] + tag + for previous_end, span, tag in zip(previous_ends, self.spans, self.tags) + ) + return "".join((*kept_parts, text[previous_ends[-1] :])) + + def original_start(self, position: int) -> int: + index: Final = bisect.bisect_right(self.tag_starts, position) - 1 + if index < 0: + return position + tag_end: Final = self.tag_starts[index] + len(self.tags[index]) + if position < tag_end: + return self.spans[index].start + return position - tag_end + self.spans[index].end + + def original_end(self, position: int) -> int: + index: Final = bisect.bisect_left(self.tag_starts, position) - 1 + if index < 0: + return position + tag_end: Final = self.tag_starts[index] + len(self.tags[index]) + if position <= tag_end: + return self.spans[index].end + return position - tag_end + self.spans[index].end + + +def _redacted_text(spans: tuple[_RedactionSpan, ...], redaction_format: str) -> _RedactedText: + tags: Final = tuple(redaction_format.format(pattern_name=span.pattern_name.upper()) for span in spans) + shifts: Final = itertools.accumulate( + (len(tag) - (span.end - span.start) for span, tag in zip(spans, tags)), initial=0 + ) + return _RedactedText(spans, tags, tuple(span.start + shift for span, shift in zip(spans, shifts))) + + class CategoryFileData(TypedDict, total=False): category_name: str description: str @@ -987,21 +1049,6 @@ class ContentFilterGuardrail(CustomGuardrail): merged.append((start, end)) return merged - def _mask_spans(self, text: str, spans: list[tuple[int, int]], redaction: str) -> str: - """Apply masking for the provided spans using the given redaction tag.""" - - if not spans: - return text - - result_parts: Final[list[str]] = [] - previous_end = 0 - for start, end in spans: - result_parts.append(text[previous_end:start]) - result_parts.append(redaction) - previous_end = end - result_parts.append(text[previous_end:]) - return "".join(result_parts) - def _convert_word_number_sequence(self, sequence: str) -> str | None: """Convert a spelled-out digit sequence (e.g., 'One-Two') into digits.""" @@ -1392,10 +1439,8 @@ class ContentFilterGuardrail(CustomGuardrail): self, pattern_name: str, action: ContentFilterAction, - text: str, - spans: list[tuple[int, int]], detections: list[ContentFilterDetection] | None, - ) -> str: + ) -> None: """Handle regex pattern match detection and action.""" if detections is not None: pattern_detection: Final[PatternDetection] = { @@ -1413,11 +1458,49 @@ class ContentFilterGuardrail(CustomGuardrail): detail={"error": error_msg, "pattern": pattern_name}, ) elif action == ContentFilterAction.MASK: - redaction_tag: Final = self.pattern_redaction_format.format(pattern_name=pattern_name.upper()) - text = self._mask_spans(text, spans, redaction_tag) verbose_proxy_logger.info("Masked all %s matches in content", pattern_name) - return text + def _mask_step( + self, + text: str, + spans: tuple[_RedactionSpan, ...], + indexed_entry: tuple[int, CompiledPatternEntry], + ) -> tuple[_RedactionSpan, ...]: + index, pattern_entry = indexed_entry + if pattern_entry["action"] != ContentFilterAction.MASK: + return spans + original_matches: Final = self._find_pattern_spans(text, pattern_entry) + redacted_matches: Final = self._redacted_text_matches(text, spans, pattern_entry) if spans else () + return _merge_redaction_spans( + ( + *spans, + *( + _RedactionSpan(start, end, index, pattern_entry["pattern_name"]) + for start, end in (*original_matches, *redacted_matches) + ), + ) + ) + + def _redacted_text_matches( + self, text: str, spans: tuple[_RedactionSpan, ...], pattern_entry: CompiledPatternEntry + ) -> tuple[tuple[int, int], ...]: + redacted: Final = _redacted_text(spans, self.pattern_redaction_format) + return tuple( + (redacted.original_start(start), redacted.original_end(end)) + for start, end in self._find_pattern_spans(redacted.render(text), pattern_entry) + ) + + def _pattern_detected( + self, + text: str, + index: int, + pattern_entry: CompiledPatternEntry, + span_states: Sequence[tuple[_RedactionSpan, ...]], + ) -> bool: + if pattern_entry["action"] == ContentFilterAction.MASK: + return any(span.pattern_index == index for span in span_states[-1]) + text_after_earlier_masks: Final = _redacted_text(span_states[index], self.pattern_redaction_format).render(text) + return bool(self._find_pattern_spans(text_after_earlier_masks, pattern_entry)) def _handle_blocked_word_match( self, @@ -1506,13 +1589,14 @@ class ContentFilterGuardrail(CustomGuardrail): keyword, category_name, severity, action = category_keyword_match text = self._handle_category_keyword_match(keyword, category_name, severity, action, text, detections) - # Check regex patterns - process ALL patterns, not just first match - for pattern_entry in self.compiled_patterns: - spans = self._find_pattern_spans(text, pattern_entry) - if spans: - pattern_name = pattern_entry["pattern_name"] - action = pattern_entry["action"] - text = self._handle_pattern_match(pattern_name, action, text, spans, detections) + no_spans: Final[tuple[_RedactionSpan, ...]] = () + span_states: Final = tuple( + itertools.accumulate(enumerate(self.compiled_patterns), partial(self._mask_step, text), initial=no_spans) + ) + for index, pattern_entry in enumerate(self.compiled_patterns): + if self._pattern_detected(text, index, pattern_entry, span_states): + self._handle_pattern_match(pattern_entry["pattern_name"], pattern_entry["action"], detections) + text = _redacted_text(span_states[-1], self.pattern_redaction_format).render(text) # Check blocked words - iterate through ALL blocked words text_lower = text.lower() @@ -2071,7 +2155,9 @@ class ContentFilterGuardrail(CustomGuardrail): def _streamed_scan_plan(self) -> _StreamedScanPlan: """ Per-stream inputs for buffer trimming: the retained tail length (the default - context, widened to the longest configured keyword), the category exception + context, widened to the longest configured keyword), the unreleased tail length + (widened to the longest holdback a configured prebuilt pattern declares, since + provider keys run longer than the default), the category exception phrases, which suppress matches anywhere in the scanned text, and the conditional category words, which only match when paired inside one sentence. """ @@ -2081,6 +2167,13 @@ class ContentFilterGuardrail(CustomGuardrail): ) return _StreamedScanPlan( context_chars=max(CONTENT_FILTER_STREAMING_SCAN_CONTEXT_CHARS, longest_keyword), + holdback_chars=max( + ( + PATTERN_STREAMING_HOLDBACK_CHARS.get(entry["pattern_name"], CONTENT_FILTER_STREAMING_HOLDBACK_CHARS) + for entry in self.compiled_patterns + ), + default=CONTENT_FILTER_STREAMING_HOLDBACK_CHARS, + ), exception_phrases=tuple( phrase for category in self.loaded_categories.values() for phrase in category.exceptions ), @@ -2221,7 +2314,7 @@ class ContentFilterGuardrail(CustomGuardrail): safe_to_yield_len = max( previous_state.yielded_masked_text_len, - len(masked_text) - (0 if is_final else CONTENT_FILTER_STREAMING_HOLDBACK_CHARS), + len(masked_text) - (0 if is_final else plan.holdback_chars), ) released_text = masked_text[previous_state.yielded_masked_text_len : safe_to_yield_len] if isinstance(choice, TextChoices): diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.json b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.json index db6dd36726c..dcb2dc5ce77 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.json +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.json @@ -98,6 +98,28 @@ "category": "Credential Patterns", "description": "Detects generic API key patterns" }, + { + "name": "openai_api_key", + "display_name": "OpenAI API Key", + "pattern": "sk-(?:(?:proj|svcacct|admin|None)-[A-Za-z0-9_\\-]{20,200}T3BlbkFJ(?:(?!sk-(?:proj|svcacct|admin|None)-)[A-Za-z0-9_\\-]){20,200}|[A-Za-z0-9]{20}T3BlbkFJ[A-Za-z0-9]{20}(?![A-Za-z0-9_]))", + "streaming_holdback_chars": 419, + "category": "Credential Patterns", + "description": "Detects OpenAI API keys (sk-proj-..., sk-svcacct-..., sk-admin-..., sk-None-... and legacy sk-... keys, all carrying the T3BlbkFJ marker)" + }, + { + "name": "anthropic_api_key", + "display_name": "Anthropic API Key", + "pattern": "(?:(? DBSpendUpdateWriter: return proxy_logging_obj.db_spend_update_writer +def _committed_spend_scope(user_api_key_dict: UserAPIKeyAuth) -> EntitySpendCounterKeys: + committed: Final = user_api_key_dict.committed_budget_spend + return committed.counter_keys if committed is not None else EntitySpendCounterKeys() + + def _proxy_spend_counters() -> "_IncrementSpendCounters": from litellm.proxy.proxy_server import increment_spend_counters @@ -639,6 +647,7 @@ class _ProxyDBLogger(CustomLogger): team_id=user_api_key_dict.team_id, org_id=user_api_key_dict.org_id, project_id=user_api_key_dict.project_id, + committed_spend_scope=_committed_spend_scope(user_api_key_dict), kwargs=dict(row_kwargs), completion_response=original_exception, start_time=start_time, @@ -757,6 +766,7 @@ class _ProxyDBLogger(CustomLogger): router=get_llm_router(), ), project_id=user_api_key_dict.project_id, + committed_spend_scope=_committed_spend_scope(user_api_key_dict), **_ledger_counter_options(row_kwargs.get(ACCOUNTING_LEDGER_KWARG)), ) except Exception: # noqa: BLE001 # one row's counter failure must not keep the other rows' counters from landing @@ -1269,10 +1279,38 @@ class _IncrementSpendCounters(Protocol): request_started_at: datetime | None = None, model_access_groups: Sequence[str] | None = None, project_id: str | None = None, + committed_spend_scope: EntitySpendCounterKeys | None = None, windows_from_ledger: bool = False, ) -> None: ... +async def _committed_spend_scope_of_row( + kwargs: Mapping[str, object], + response_cost: float, + token: str | None, + team_id: str | None, + user_id: str | None, + org_id: str | None, + end_user_id: str | None, +) -> EntitySpendCounterKeys: + metadata: Final = _OBJECT_MAPPING.validate_python(get_litellm_metadata_from_kwargs(kwargs=dict(kwargs))) + carried: Final = committed_budget_spend_of_row( + metadata, _OBJECT_MAPPING.validate_python(kwargs.get("litellm_params") or {}) + ) + if carried is not None: + return carried.counter_keys + if response_cost <= 0 or not committed_spend_tracked(): + return EntitySpendCounterKeys() + from litellm.proxy.proxy_server import spend_counter_store + + return await warm_committed_counter_keys( + spend_counter_store, + EntitySpendCounterKeys.of( + token=token, team_id=team_id, user_id=user_id, org_id=org_id, end_user_id=end_user_id + ), + ) + + async def _update_database_and_spend_counters( proxy_logging_obj: "ProxyLogging", increment_spend_counters: _IncrementSpendCounters, @@ -1291,6 +1329,7 @@ async def _update_database_and_spend_counters( model_access_groups: Sequence[str] | None = None, project_id: str | None = None, spend_writer: DBSpendUpdateWriter | None = None, + committed_spend_scope: EntitySpendCounterKeys | None = None, ) -> bool: """The reservation is reconciled before the spend is persisted, from its own read. One spend counter batch then spans the database write and the counter update, so the post-call counters are read with a single MGET after the @@ -1333,6 +1372,7 @@ async def _update_database_and_spend_counters( request_tags=request_tags, model_access_groups=model_access_groups, project_id=project_id, + committed_spend_scope=committed_spend_scope, ) @@ -1354,6 +1394,7 @@ async def _update_database_and_spend_counters_in_batch( model_access_groups: Sequence[str] | None, project_id: str | None, spend_writer: DBSpendUpdateWriter | None = None, + committed_spend_scope: EntitySpendCounterKeys | None = None, ) -> bool: try: writer: Final = spend_writer if spend_writer is not None else proxy_logging_obj.db_spend_update_writer @@ -1407,6 +1448,19 @@ async def _update_database_and_spend_counters_in_batch( request_started_at=start_time, model_access_groups=model_access_groups, project_id=project_id, + committed_spend_scope=( + committed_spend_scope + if committed_spend_scope is not None + else await _committed_spend_scope_of_row( + kwargs=kwargs, + response_cost=response_cost, + token=user_api_key, + team_id=team_id, + user_id=user_id, + org_id=org_id, + end_user_id=end_user_id, + ) + ), **_ledger_counter_options(_BATCH_CALLBACK_FIELDS.validate_python(kwargs).get(ACCOUNTING_LEDGER_KWARG)), ) except Exception: diff --git a/litellm/proxy/hooks/responses_id_security.py b/litellm/proxy/hooks/responses_id_security.py index 0ab3d1ca7af..f99dd4c699b 100644 --- a/litellm/proxy/hooks/responses_id_security.py +++ b/litellm/proxy/hooks/responses_id_security.py @@ -16,6 +16,7 @@ from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.base_llm.managed_resources.isolation import resolve_resource_owner_id from litellm.proxy._types import LitellmUserRoles from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, @@ -276,7 +277,12 @@ class ResponsesIDSecurity(CustomLogger): ): return True - if response_id_user_id and response_id_user_id != user_api_key_dict.user_id: + if not response_id_user_id and not response_id_team_id: + if general_settings.get("disable_responses_id_security", False): + return True + raise HTTPException(status_code=403, detail="Forbidden. This response id has no associated owner") + + if response_id_user_id and response_id_user_id != resolve_resource_owner_id(user_api_key_dict): if general_settings.get("disable_responses_id_security", False): verbose_proxy_logger.debug( "Responses ID Security is disabled. User %s is accessing response id %s which is not associated with them.", @@ -394,7 +400,7 @@ class ResponsesIDSecurity(CustomLogger): managed_id: Final = SpecialEnums.LITELLM_MANAGED_RESPONSE_API_RESPONSE_ID_COMPLETE_STR.value.format( original_id, - user_api_key_dict.user_id or "", + resolve_resource_owner_id(user_api_key_dict) or "", user_api_key_dict.team_id or "", ) encrypted_id: Final = f"resp_{encrypt_value_helper(value=managed_id)}" diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 85a730a4848..7f3dce4b3ad 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -106,7 +106,7 @@ _REDACTED_HEADER_VALUE: Final = "***REDACTED***" _DECODED_JSON: Final = TypeAdapter(object) _OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) _CREDENTIAL_HEADER_NAMES: Final = SpecialHeaders.litellm_credential_header_names() | frozenset( - {"cookie", "proxy-authorization"} + {"cookie", "proxy-authorization", "x-amz-security-token"} ) _TRANSPORT_ONLY_CREDENTIAL_KEYS: Final = frozenset({"provider_specific_header", "headers", "api_key"}) _REQUEST_FIELDS: Final = TypeAdapter(Mapping[str, object]) @@ -2155,7 +2155,14 @@ class LiteLLMProxyRequestSetup: user_api_key_dict.matched_model_access_groups ) # UserAPIKeyAuth object for MCP server access control - data[_metadata_variable_name]["user_api_key_auth"] = user_api_key_dict.model_copy( + data[_metadata_variable_name]["user_api_key_auth"] = LiteLLMProxyRequestSetup.auth_for_request_metadata( + user_api_key_dict + ) + return data + + @staticmethod + def auth_for_request_metadata(user_api_key_dict: UserAPIKeyAuth) -> UserAPIKeyAuth: + return user_api_key_dict.model_copy( update={ "metadata": strip_callback_config(user_api_key_dict.metadata), "team_metadata": strip_callback_config(user_api_key_dict.team_metadata), @@ -2163,7 +2170,6 @@ class LiteLLMProxyRequestSetup: "organization_metadata": strip_callback_config(user_api_key_dict.organization_metadata), } ) - return data @staticmethod def add_management_endpoint_metadata_to_request_metadata( diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index 0640bf0fa4b..672bbdb164a 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -42,7 +42,9 @@ from litellm.proxy.management_helpers.object_permission_utils import ( # noqa: handle_update_object_permission_common, set_object_permission, ) +from litellm.proxy.spend_tracking.committed_spend import drop_committed_spend from litellm.proxy.utils import handle_exception_on_proxy +from litellm.redis.stores.spend_counters import end_user_counter from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.chunked_in import find_many_in from litellm.repositories.table_repositories import EndUserRepository @@ -812,7 +814,7 @@ async def delete_end_user( See below for all params ``` """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import prisma_client, spend_counter_store try: if prisma_client is None: @@ -842,6 +844,9 @@ async def delete_end_user( verbose_proxy_logger.debug("received response from updating prisma client. response=%s", response) await _evict_end_user_cache_keys(_end_user_cache_keys(data.user_ids)) + await drop_committed_spend( + spend_counter_store, tuple(end_user_counter(user_id) for user_id in data.user_ids) + ) return DeleteCustomersResponse( deleted_customers=response, diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index ac7463104e2..f88ac445a0b 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -39,7 +39,10 @@ from litellm.proxy.auth.auth_checks import ( get_team_object, get_user_object, ) -from litellm.proxy.auth.auth_utils import enforce_batch_limits_are_admin_only +from litellm.proxy.auth.auth_utils import ( + enforce_batch_limits_are_admin_only, + enforce_vector_store_index_grants_are_admin_only, +) from litellm.proxy.auth.password_policy import ( validate_password_not_breached, validate_password_policy, @@ -595,6 +598,7 @@ async def new_user( if data.auto_create_key and isinstance(user_api_key_dict, UserAPIKeyAuth): enforce_batch_limits_are_admin_only(data, None, user_api_key_dict, "key") + enforce_vector_store_index_grants_are_admin_only(data, None, None, user_api_key_dict, "key") raise_on_logging_credential_change(data.metadata, None, user_api_key_dict) check_permissions_caller_permission( @@ -1481,8 +1485,11 @@ async def _invalidate_user_spend_counter_if_changed( """ if non_default_values.get("spend") is not None: from litellm.proxy.proxy_server import invalidate_spend_counter + from litellm.proxy.spend_tracking.committed_spend import committed_spend_key - await invalidate_spend_counter(counter_key=user_counter(str(non_default_values["user_id"]))) + counter_key: Final = user_counter(str(non_default_values["user_id"])) + await invalidate_spend_counter(counter_key=counter_key) + await invalidate_spend_counter(counter_key=committed_spend_key(counter_key)) def _clears_object_permission(user_request: UpdateUserRequest) -> bool: diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 2b6cbd6529d..038c1bdf5a6 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -62,6 +62,7 @@ from litellm.proxy.auth.auth_utils import ( abbreviate_api_key, enforce_batch_limits_are_admin_only, enforce_output_token_estimates_are_admin_only, + enforce_vector_store_index_grants_are_admin_only, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( @@ -162,6 +163,7 @@ from litellm.proxy.mcp_server.outbound_credentials.sso_assertion_store import ( ) from litellm.proxy.search_endpoints.search_tool_registry import rotate_search_tools_master_key from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start +from litellm.proxy.spend_tracking.committed_spend import committed_spend_key from litellm.proxy.spend_tracking.spend_tracking_utils import ( # noqa: F401 # legacy module exports _is_master_key, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export is_master_key, @@ -1538,6 +1540,13 @@ async def _common_key_generation_helper( user_api_key_dict=user_api_key_dict, entity="key", ) + enforce_vector_store_index_grants_are_admin_only( + data=data, + top_level_grants=data.allowed_vector_store_indexes, + existing_metadata=None, + user_api_key_dict=user_api_key_dict, + entity="key", + ) if data.metadata is not None and data.metadata.get("service_account_id") is not None and data.team_id is None: await validate_team_id_used_in_service_account_request( @@ -2447,7 +2456,7 @@ async def generate_key_fn( - prompts: Optional[List[str]] - List of allowed prompts for the key. If specified, the key will only be able to use these specific prompts. - auto_rotate: Optional[bool] - Whether this key should be automatically rotated (regenerated) - rotation_interval: Optional[str] - How often to auto-rotate this key (e.g., '30s', '30m', '30h', '30d'). Required if auto_rotate=True. - - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint. + - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Only a proxy admin can set or change it. Create index, using `/v1/indexes` endpoint. - router_settings: Optional[UpdateRouterConfig] - key-specific router settings. Example - {"model_group_retry_policy": {"gpt-4": {"RateLimitErrorRetries": 5}}}. IF null or {} then no router settings. - access_group_ids: Optional[List[str]] - List of access group IDs to associate with the key. Access groups define which models a key can access. Example - ["access_group_1", "access_group_2"]. - budget_limits: Optional[list] - List of concurrent budget windows for the key. Each window specifies a budget_limit, time_period, and optional budget_duration. Example - [{"budget_limit": 10.0, "time_period": "1d"}, {"budget_limit": 50.0, "time_period": "7d"}]. @@ -2649,7 +2658,7 @@ async def generate_service_account_key_fn( - allowed_routes: Optional[list] - List of allowed routes for the key. Store the actual route or store a wildcard pattern for a set of routes. Example - ["/chat/completions", "/embeddings", "/keys/*"] - object_permission: Optional[LiteLLM_ObjectPermissionBase] - key-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission. Examples: - - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint. + - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Only a proxy admin can set or change it. Create index, using `/v1/indexes` endpoint. 1. Allow users to turn on/off pii masking @@ -3368,6 +3377,13 @@ async def _process_single_key_update( user_api_key_dict=user_api_key_dict, entity="key", ) + enforce_vector_store_index_grants_are_admin_only( + data=update_key_request, + top_level_grants=update_key_request.allowed_vector_store_indexes, + existing_metadata=existing_key_row.metadata, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # LiteLLM_VerificationToken.metadata is a bare dict + user_api_key_dict=user_api_key_dict, + entity="key", + ) await _check_single_key_update_team_permissions( user_api_key_dict=user_api_key_dict, @@ -3877,6 +3893,13 @@ async def _validate_update_key_data( user_api_key_dict=user_api_key_dict, entity="key", ) + enforce_vector_store_index_grants_are_admin_only( + data=data, + top_level_grants=data.allowed_vector_store_indexes, + existing_metadata=_existing_metadata if isinstance(_existing_metadata, dict) else None, # pyright: ignore[reportUnknownArgumentType] # stored metadata is a bare dict + user_api_key_dict=user_api_key_dict, + entity="key", + ) # Personal-key bypass: the caller both created the key AND still owns it # (user_id == caller). Checking only created_by would let a demoted admin @@ -4120,7 +4143,7 @@ async def update_key_fn( - object_permission: Optional[LiteLLM_ObjectPermissionBase] - key-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission. - auto_rotate: Optional[bool] - Whether this key should be automatically rotated - rotation_interval: Optional[str] - How often to rotate this key (e.g., '30d', '90d'). Required if auto_rotate=True - - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint. + - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Only a proxy admin can set or change it. Create index, using `/v1/indexes` endpoint. - router_settings: Optional[UpdateRouterConfig] - key-specific router settings. Example - {"model_group_retry_policy": {"gpt-4": {"RateLimitErrorRetries": 5}}}. IF null or {} then no router settings. - access_group_ids: Optional[List[str]] - List of access group IDs to associate with the key. Access groups define which models a key can access. Example - ["access_group_1", "access_group_2"]. - budget_limits: Optional[list] - List of concurrent budget windows for the key. Each window specifies a budget_limit, time_period, and optional budget_duration. Example - [{"budget_limit": 10.0, "time_period": "1d"}, {"budget_limit": 50.0, "time_period": "7d"}]. @@ -4297,7 +4320,7 @@ async def update_key_fn( ) if data.spend is not None: - from litellm.proxy.proxy_server import spend_counter_store + from litellm.proxy.proxy_server import invalidate_spend_counter, spend_counter_store counter_key: Final = key_counter(hash_token_if_needed(key)) try: @@ -4309,6 +4332,7 @@ async def update_key_fn( counter_key, redis_err, ) + await invalidate_spend_counter(counter_key=committed_spend_key(counter_key)) asyncio.create_task( KeyManagementEventHooks.async_key_updated_hook( @@ -5135,9 +5159,14 @@ async def info_key_fn( # default to using Auth token if no key is passed in key = key or user_api_key_dict.api_key - hashed_key: str | None = key - if key is not None: - hashed_key = hash_token_if_needed(token=key) # rebind-ok: pre-existing rebinding on a rename-only line + if not key: + raise ProxyException( + message="No key passed in. Pass ?key= or authenticate with a virtual key.", + type=ProxyErrorTypes.bad_request_error, + param="key", + code=status.HTTP_400_BAD_REQUEST, + ) + hashed_key: str | None = hash_token_if_needed(token=key) live_key_info: Final = await _prisma_table(VerificationTokenRepository(prisma_client)).find_unique( where={"token": hashed_key}, include={"litellm_budget_table": True}, @@ -6338,6 +6367,13 @@ async def _execute_virtual_key_regeneration( user_api_key_dict=user_api_key_dict, entity="key", ) + enforce_vector_store_index_grants_are_admin_only( + data=data, + top_level_grants=data.allowed_vector_store_indexes, + existing_metadata=_existing_key_metadata if isinstance(_existing_key_metadata, dict) else None, # pyright: ignore[reportUnknownArgumentType] # stored metadata is a bare dict + user_api_key_dict=user_api_key_dict, + entity="key", + ) await _validate_end_user_budget_id_change( requested_budget_id=_requested_end_user_budget_id(data), existing_budget_id=get_key_end_user_budget_id( @@ -6965,6 +7001,7 @@ async def reset_key_spend_fn( try: from litellm.proxy.proxy_server import ( hash_token, + invalidate_spend_counter, prisma_client, proxy_logging_obj, set_spend_counter_with_floor_and_broadcast, @@ -7020,6 +7057,7 @@ async def reset_key_spend_fn( # pre-reset value. _counter_key: Final = key_counter(hashed_api_key) await set_spend_counter_with_floor_and_broadcast(counter_key=_counter_key, value=reset_to) + await invalidate_spend_counter(counter_key=committed_spend_key(_counter_key)) await _reset_key_budget_windows( prisma_client=prisma_client, hashed_api_key=hashed_api_key, diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index b3c1f48bf99..38b814225ce 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -15,7 +15,7 @@ import datetime import itertools import json from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence -from contextlib import AbstractAsyncContextManager, asynccontextmanager, suppress +from contextlib import AbstractAsyncContextManager, asynccontextmanager from dataclasses import dataclass from fnmatch import fnmatchcase from functools import partial @@ -62,6 +62,14 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.litellm_license import AUTO_ROUTER_LICENSE_REMEDY +from litellm.proxy.auth.model_checks import ( + GLOBAL_ALIASES_ONLY, + OtherDeploymentsBill, + WildcardDeployment, + model_info_cost_map_entry, + wildcard_deployment_priced_model, + wildcard_row_filled_model, +) from litellm.proxy.auth.team_grants import team_model_aliases from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.config_sync_pubsub import ( @@ -78,6 +86,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, ) +from litellm.proxy.common_utils.model_listing_utils import CallerAliases from litellm.proxy.common_utils.openai_error_payload import proxy_exception_for from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient @@ -1314,14 +1323,47 @@ def _ptu_priced_deployment(model_params: Deployment) -> Deployment: ) -def _cost_map_entry(db_model: Deployment, incoming_model_info: Mapping[str, object]) -> Mapping[str, object]: +@dataclass(frozen=True, slots=True) +class _ListedRow: + name: str | None = None + llm_router: Router | None = None + caller_aliases: CallerAliases = GLOBAL_ALIASES_ONLY + + +_NO_LISTED_ROW: Final = _ListedRow() + + +def _live_cost_map_model(db_model: Deployment, listed_row: _ListedRow) -> str | None: + stored_model: Final = _decrypted_model(db_model.litellm_params.model) + if stored_model is None or "*" not in stored_model or listed_row.name is None or listed_row.llm_router is None: + return stored_model + return wildcard_row_filled_model( + listed_row.llm_router, db_model.model_name, stored_model, listed_row.name, listed_row.caller_aliases + ) + + +def _cost_map_entry( + db_model: Deployment, incoming_model_info: Mapping[str, object], listed_row: _ListedRow = _NO_LISTED_ROW +) -> Mapping[str, object]: base_model: Final = incoming_model_info.get("base_model") - lookup: Final = base_model if isinstance(base_model, str) else _decrypted_model(db_model.litellm_params.model) - if lookup is None: + lookup: Final = base_model if isinstance(base_model, str) else _live_cost_map_model(db_model, listed_row) + return MappingProxyType({}) if lookup is None else model_info_cost_map_entry(lookup) + + +def _routed_cost_map_entry(db_model: Deployment, listed_row: _ListedRow) -> Mapping[str, object]: + stored_model: Final = _decrypted_model(db_model.litellm_params.model) + if stored_model is None or "*" not in stored_model or listed_row.name is None or listed_row.llm_router is None: return MappingProxyType({}) - with suppress(Exception): - return MappingProxyType(dict(litellm.get_model_info(model=lookup))) - return MappingProxyType({}) + priced_as: Final = wildcard_deployment_priced_model( + listed_row.llm_router, + WildcardDeployment( + db_model.model_info.id or "", db_model.model_name, stored_model, db_model.model_info.team_id + ), + listed_row.name, + listed_row.caller_aliases, + ) + billed_model: Final = priced_as.model if isinstance(priced_as, OtherDeploymentsBill) else priced_as + return model_info_cost_map_entry(billed_model) if isinstance(billed_model, str) else MappingProxyType({}) LoadedCatalog: TypeAlias = Callable[[], Mapping[str, Mapping[str, object]]] @@ -1361,6 +1403,7 @@ def legacy_update_model_info_changes( incoming_model_info: ModelInfo, discovered_model_info: Mapping[str, object] = MappingProxyType({}), loaded_catalog: LoadedCatalog = GetModelCostMap.loaded_model_cost_map, + listed_row: _ListedRow = _NO_LISTED_ROW, ) -> Mapping[str, object]: """The ``model_info`` fields ``POST /model/update`` persists from its request. @@ -1380,7 +1423,8 @@ def legacy_update_model_info_changes( echoed: Final = frozenset( echoed_cost_map_fields( requested, - _cost_map_entry(db_model, {**stored, **sent}), + _cost_map_entry(db_model, {**stored, **sent}, listed_row), + _routed_cost_map_entry(db_model, listed_row), _loaded_catalog_entry(requested, loaded_catalog), discovered_model_info, ) @@ -1399,9 +1443,10 @@ def _legacy_update_model_info_writes( incoming_model_info: ModelInfo, discovered_model_info: Mapping[str, object], member_marker: bool | None, + listed_row: _ListedRow, ) -> Mapping[str, object]: changes: Final = { - **legacy_update_model_info_changes(current, incoming_model_info, discovered_model_info), + **legacy_update_model_info_changes(current, incoming_model_info, discovered_model_info, listed_row=listed_row), **({"member_auto_router": member_marker} if member_marker is not None else {}), } return MappingProxyType({k: v for k, v in changes.items() if getattr(current.model_info, k, None) != v}) @@ -1413,13 +1458,14 @@ async def _locked_legacy_update_model_info_json( incoming_model_info: ModelInfo, discovered_model_info: Mapping[str, object], member_marker: bool | None, + listed_row: _ListedRow, ) -> str | None: row: Final = await table.find_unique(where={"model_id": model_id}) if row is None: return None payload: Final = row.model_dump() writes: Final = _legacy_update_model_info_writes( - Deployment(**payload), incoming_model_info, discovered_model_info, member_marker + Deployment(**payload), incoming_model_info, discovered_model_info, member_marker, listed_row ) if not writes: return None @@ -1445,6 +1491,9 @@ def update_db_model( db_model: Deployment, updated_patch: updateDeployment, loaded_catalog: LoadedCatalog = GetModelCostMap.loaded_model_cost_map, + llm_router: Router | None = None, + caller_aliases: CallerAliases = GLOBAL_ALIASES_ONLY, + listed_name: str | None = None, ) -> PrismaCompatibleUpdateDBModel: if updated_patch.model_info is not None: _raise_if_ptu_cost_attribution_disabled(updated_patch.model_info.model_dump(exclude_none=True)) @@ -1474,9 +1523,11 @@ def update_db_model( # update model info if updated_patch.model_info: incoming_model_info: Final = updated_patch.model_info.model_dump(exclude_none=True) + listed_row: Final = _ListedRow(listed_name or updated_patch.model_name, llm_router, caller_aliases) echoed_fields: Final = echoed_cost_map_fields( incoming_model_info, - _cost_map_entry(db_model, incoming_model_info), + _cost_map_entry(db_model, incoming_model_info, listed_row), + _routed_cost_map_entry(db_model, listed_row), _loaded_catalog_entry(incoming_model_info, loaded_catalog), ) stored_model_info_fields: Final[Mapping[str, object]] = without_server_derived_pricing(incoming_model_info) @@ -1656,6 +1707,7 @@ async def patch_model( llm_router, premium_user, prisma_client, + request_caller_aliases, store_model_in_db, user_api_key_cache, ) @@ -1826,6 +1878,8 @@ async def patch_model( user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, write_row=write_row, + llm_router=llm_router, + caller_aliases=await request_caller_aliases(user_api_key_dict), ) if updated_model is None: @@ -2181,6 +2235,8 @@ async def _update_team_model_in_db( user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, write_row: Callable[[PrismaCompatibleUpdateDBModel], Awaitable[_RowT]], + llm_router: Router | None = None, + caller_aliases: CallerAliases = GLOBAL_ALIASES_ONLY, ) -> _RowT: """ Handle team model updates with proper alias management. @@ -2225,7 +2281,11 @@ async def _update_team_model_in_db( # No team_id in patch, proceed with standard update if patch_team_id is None: - return await write_row(update_db_model(db_model=db_model, updated_patch=patch_data)) + return await write_row( + update_db_model( + db_model=db_model, updated_patch=patch_data, llm_router=llm_router, caller_aliases=caller_aliases + ) + ) # Determine public model name public_model_name: Final = _get_public_model_name( @@ -2246,7 +2306,15 @@ async def _update_team_model_in_db( ), } ) - row: Final = await write_row(update_db_model(db_model=db_model, updated_patch=team_patch)) + row: Final = await write_row( + update_db_model( + db_model=db_model, + updated_patch=team_patch, + llm_router=llm_router, + caller_aliases=caller_aliases, + listed_name=patch_data.model_name, + ) + ) if is_new_team_assignment: await _setup_new_team_model_assignment( @@ -2392,12 +2460,12 @@ async def delete_team_models( if deleted_model_ids: await publish_config_change(redis_cache=coordination_redis_cache(), object_type="litellm_proxymodeltable") - # Under MODEL_RECONCILE_LOCK, for the same reason as delete_model: the rows are + # Under DB_RECONCILE_LOCK, for the same reason as delete_model: the rows are # gone, but a reconcile holding a pre-delete snapshot would upsert these ids back # onto this pod. The lock orders the eviction after any in-flight reconcile. - from litellm.proxy.proxy_server import MODEL_RECONCILE_LOCK, proxy_config + from litellm.proxy.proxy_server import DB_RECONCILE_LOCK, proxy_config - async with MODEL_RECONCILE_LOCK: + async with DB_RECONCILE_LOCK: proxy_config.remove_auto_router_catalog_entries(frozenset(deleted_model_ids)) if llm_router is not None: for model_id in deleted_model_ids: @@ -2884,7 +2952,7 @@ async def delete_model( """ from litellm.proxy.proxy_server import ( - MODEL_RECONCILE_LOCK, + DB_RECONCILE_LOCK, general_settings_view, llm_router, premium_user, @@ -2945,13 +3013,13 @@ async def delete_model( ) ## DELETE FROM ROUTER ## - # Under MODEL_RECONCILE_LOCK. The db row is already gone, but a reconcile + # Under DB_RECONCILE_LOCK. The db row is already gone, but a reconcile # that snapshotted the db BEFORE that delete still lists this id as desired, # and its _add_deployment upserts the deployment straight back -- leaving # this pod serving a model the database no longer has, until the next # reconcile. Taking the lock orders this eviction after any such in-flight # reconcile's re-add, so the eviction is the last word. - async with MODEL_RECONCILE_LOCK: + async with DB_RECONCILE_LOCK: proxy_config.remove_auto_router_catalog_entries(frozenset({model_info.id})) if llm_router is not None: llm_router.apply(ChangeSet(delete=[model_info.id])) @@ -3292,6 +3360,7 @@ async def update_model( llm_router, premium_user, prisma_client, + request_caller_aliases, store_model_in_db, user_api_key_cache, ) @@ -3429,8 +3498,11 @@ async def update_model( if llm_router is not None else MappingProxyType({}) ) + listed_row: Final = _ListedRow( + model_params.model_name, llm_router, await request_caller_aliases(user_api_key_dict) + ) requested_model_info: Final = _legacy_update_model_info_writes( - deployment, _model_info, discovered_model_info, member_marker + deployment, _model_info, discovered_model_info, member_marker, listed_row ) _raise_if_access_groups_malformed(requested_model_info) async with _auto_router_capability_slot( @@ -3442,7 +3514,7 @@ async def update_model( ) as table: model_info_json: Final = ( await _locked_legacy_update_model_info_json( - table, _model_id, _model_info, discovered_model_info, member_marker + table, _model_id, _model_info, discovered_model_info, member_marker, listed_row ) if requested_model_info else None @@ -3892,7 +3964,7 @@ def reload_serving_verdict( Without it (no reconcile ran) every drop is reported, which is the safe direction. ``live_after`` is the router's serving state captured by the reload itself, while it - still held MODEL_RECONCILE_LOCK. Pass it whenever the caller has it: re-reading the + still held DB_RECONCILE_LOCK. Pass it whenever the caller has it: re-reading the router here instead means sampling it after the lock was released, where the NEXT reconcile's leading wipe (clear_cache un-serves every db model before reloading them) shows up as this reload having dropped them. Falling back to a fresh read is @@ -3970,13 +4042,13 @@ async def clear_cache() -> ReconcileOutcome: Returns what the reload saw (see ReconcileOutcome) so callers can pass it to raise_if_reload_degraded_serving. - Runs under MODEL_RECONCILE_LOCK for its whole extent, not just the reload at the + Runs under DB_RECONCILE_LOCK for its whole extent, not just the reload at the end, so the auto-router reset and the reload that rebuilds those routers are atomic to any other reconcile. The inner call is _add_deployment_locked because add_deployment would re-acquire the same non-reentrant lock and deadlock. """ from litellm.proxy.proxy_server import ( - MODEL_RECONCILE_LOCK, + DB_RECONCILE_LOCK, llm_router, prisma_client, proxy_config, @@ -3988,7 +4060,7 @@ async def clear_cache() -> ReconcileOutcome: verbose_proxy_logger.debug("llm_router or prisma_client is None, skipping cache clear") return ReconcileOutcome(still_desired=None, live_after=None) - async with MODEL_RECONCILE_LOCK: + async with DB_RECONCILE_LOCK: try: # Only clear DB models, preserve config models verbose_proxy_logger.debug("Clearing only DB models, preserving config models") @@ -4060,7 +4132,7 @@ async def clear_cache() -> ReconcileOutcome: llm_router.registry().plugins("switchyard").pop(model_name, None) # Reload only DB models. _add_deployment_locked, not add_deployment: this - # coroutine already holds MODEL_RECONCILE_LOCK and asyncio.Lock is not + # coroutine already holds DB_RECONCILE_LOCK and asyncio.Lock is not # reentrant, so the public wrapper would deadlock against itself. outcome: Final = await proxy_config.add_deployment_locked( prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 15d05633679..bc1963af0bb 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -66,7 +66,9 @@ from litellm.proxy.management_helpers.utils import ( get_new_internal_user_defaults, management_endpoint_wrapper, ) +from litellm.proxy.spend_tracking.committed_spend import drop_committed_spend from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.redis.stores.spend_counters import org_counter, team_counter from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.chunked_in import find_many_in from litellm.repositories.object_permission_repository import ObjectPermissionRepository @@ -93,6 +95,7 @@ if TYPE_CHECKING: LiteLLM_OrganizationMembership as PrismaOrganizationMembership, ) from prisma.models import LiteLLM_OrganizationTable as PrismaOrganizationTable + from prisma.models import LiteLLM_TeamTable as PrismaTeamTable from prisma.models import LiteLLM_UserTable as PrismaUserTable from prisma.models import LiteLLM_VerificationToken as PrismaVerificationToken @@ -180,6 +183,8 @@ class _OrganizationMembershipTableClient(Protocol): class _TeamTableClient(Protocol): + async def find_many(self, *, where: Mapping[str, object]) -> "Sequence[PrismaTeamTable]": ... + async def delete_many(self, where: Mapping[str, object]) -> int: ... @@ -819,6 +824,11 @@ async def update_organization( include={"members": True, "teams": True, "litellm_budget_table": True}, ) + if data.spend is not None: + from litellm.proxy.proxy_server import spend_counter_store + + await drop_committed_spend(spend_counter_store, (org_counter(data.organization_id),)) + return response @@ -1045,7 +1055,8 @@ async def delete_organization( teams_in_deleted_orgs: Final = await find_many_in( _table(TeamRepository(prisma_client)), "organization_id", data.organization_ids ) - raise_if_teams_hold_ptu_shares(tuple(team.team_id for team in teams_in_deleted_orgs), llm_router) + team_ids_in_deleted_orgs: Final = tuple(team.team_id for team in teams_in_deleted_orgs) + raise_if_teams_hold_ptu_shares(team_ids_in_deleted_orgs, llm_router) deleted_orgs: Final = [] for organization_id in data.organization_ids: @@ -1073,6 +1084,15 @@ async def delete_organization( ) deleted_orgs.append(deleted_org) + from litellm.proxy.proxy_server import spend_counter_store + + await drop_committed_spend( + spend_counter_store, + ( + *(org_counter(organization_id) for organization_id in data.organization_ids), + *(team_counter(team_id) for team_id in team_ids_in_deleted_orgs), + ), + ) return deleted_orgs diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 59fa8a18ac7..210aa1aaa6f 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -122,6 +122,7 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo from litellm.proxy.auth.auth_utils import ( enforce_batch_limits_are_admin_only, enforce_output_token_estimates_are_admin_only, + enforce_vector_store_index_grants_are_admin_only, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( @@ -225,12 +226,14 @@ from litellm.proxy.management_helpers.utils import ( add_new_member, management_endpoint_wrapper, ) +from litellm.proxy.spend_tracking.committed_spend import drop_committed_spend from litellm.proxy.utils import ( PrismaClient, ProxyLogging, handle_exception_on_proxy, query_raw_within_statement_timeout, ) +from litellm.redis.stores.spend_counters import team_counter from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.chunked_in import delete_many_in, find_many_in from litellm.repositories.daily_activity_buckets import api_key_from_bucket @@ -1674,7 +1677,7 @@ async def new_team( - team_member_key_duration: Optional[str] - The duration for a team member's key. e.g. "1d", "1w", "1mo" - allowed_passthrough_routes: Optional[List[str]] - List of allowed pass through routes for the team. - denied_passthrough_routes: Optional[List[str]] - List of pass through routes the team's keys may not call, even if allowed. Applies together with each key's `denied_passthrough_routes`. - - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint. + - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Only a proxy admin can set or change it. Create index, using `/v1/indexes` endpoint. - secret_manager_settings: Optional[dict] - Secret manager settings for the team. [Docs](https://docs.litellm.ai/docs/secret_managers/overview) - router_settings: Optional[UpdateRouterConfig] - team-specific router settings. Example - {"model_group_retry_policy": {"gpt-4": {"RateLimitErrorRetries": 5}}}. IF null or {} then no router settings. - access_group_ids: Optional[List[str]] - List of access group IDs to associate with the team. Access groups define which models the team can access. Example - ["access_group_1", "access_group_2"]. @@ -1785,6 +1788,13 @@ async def new_team( user_api_key_dict=user_api_key_dict, entity="team", ) + enforce_vector_store_index_grants_are_admin_only( + data=data, + top_level_grants=data.allowed_vector_store_indexes, + existing_metadata=None, + user_api_key_dict=user_api_key_dict, + entity="team", + ) # Check if license is over limit total_teams: Final = await _team_db(prisma_client).count() @@ -2459,7 +2469,7 @@ async def update_team( - default_estimated_output_tokens_per_model: Optional[Dict[str, int]] - Per-model override of the above. Example: {"gpt-4": 4096, "gpt-3.5-turbo": 1024} - mcp_rpm_limit: Optional[Dict[str, int]] - Per-MCP-server RPM limit for this team, keyed by MCP server name (alias if set, else the configured name). Example: {"github": 100, "slack": 200}. Applied across all keys for this team. Example - update team TPM Limit - - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint. + - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Only a proxy admin can set or change it. Create index, using `/v1/indexes` endpoint. - secret_manager_settings: Optional[dict] - Secret manager settings for the team. [Docs](https://docs.litellm.ai/docs/secret_managers/overview) - router_settings: Optional[UpdateRouterConfig] - team-specific router settings. Example - {"model_group_retry_policy": {"gpt-4": {"RateLimitErrorRetries": 5}}}. IF null or {} then no router settings. - access_group_ids: Optional[List[str]] - List of access group IDs to associate with the team. Access groups define which models the team can access. Example - ["access_group_1", "access_group_2"]. @@ -2597,6 +2607,13 @@ async def update_team( user_api_key_dict=user_api_key_dict, entity="team", ) + enforce_vector_store_index_grants_are_admin_only( + data=data, + top_level_grants=data.allowed_vector_store_indexes, + existing_metadata=_existing_team_metadata if isinstance(_existing_team_metadata, dict) else None, # pyright: ignore[reportUnknownArgumentType] # stored metadata is a bare dict + user_api_key_dict=user_api_key_dict, + entity="team", + ) check_passthrough_routes_caller_permission( data, @@ -4848,6 +4865,7 @@ async def delete_validated_team_rows( """ from litellm.proxy.proxy_server import ( proxy_logging_obj, + spend_counter_store, user_api_key_cache, ) @@ -4919,7 +4937,9 @@ async def delete_validated_team_rows( _emit_team_members_metric(deleted_team.model_copy(update={"members_with_roles": ()})) await sync_team_access_group_membership(prisma_client=prisma_client, team_id=deleted_team.team_id) - return DeletedTeams(team_ids=tuple(team.team_id for team in locked)) + deleted_team_ids: Final = tuple(team.team_id for team in locked) + await drop_committed_spend(spend_counter_store, tuple(team_counter(team_id) for team_id in deleted_team_ids)) + return DeletedTeams(team_ids=deleted_team_ids) def _locked_roster_user_ids(locked: Sequence[LiteLLM_TeamTable]) -> tuple[str, ...]: diff --git a/litellm/proxy/management_helpers/bulk_user_creation.py b/litellm/proxy/management_helpers/bulk_user_creation.py index 62228f6faa8..77e37cb11cb 100644 --- a/litellm/proxy/management_helpers/bulk_user_creation.py +++ b/litellm/proxy/management_helpers/bulk_user_creation.py @@ -29,7 +29,10 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.auth_checks import invalidate_team_member_spend_state -from litellm.proxy.auth.auth_utils import enforce_batch_limits_are_admin_only +from litellm.proxy.auth.auth_utils import ( + enforce_batch_limits_are_admin_only, + enforce_vector_store_index_grants_are_admin_only, +) from litellm.proxy.auth.litellm_license import LicenseCheck from litellm.proxy.common_utils.callback_config_validation import credential_reference_rejection from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time @@ -220,6 +223,7 @@ def _row_error(item: BulkNewUserItem, user_api_key_dict: UserAPIKeyAuth) -> str check_permissions_caller_permission(data=item, user_api_key_dict=user_api_key_dict) if item.auto_create_key: enforce_batch_limits_are_admin_only(item, None, user_api_key_dict, "key") + enforce_vector_store_index_grants_are_admin_only(item, None, None, user_api_key_dict, "key") except Exception as exc: # noqa: BLE001 # any validation failure is reported on this row only return _error_message(exc) rejection: Final = ( diff --git a/litellm/proxy/management_helpers/bulk_user_deletion.py b/litellm/proxy/management_helpers/bulk_user_deletion.py index 04dd0b3f57f..cc60f58efb0 100644 --- a/litellm/proxy/management_helpers/bulk_user_deletion.py +++ b/litellm/proxy/management_helpers/bulk_user_deletion.py @@ -46,7 +46,9 @@ from litellm.proxy.management_helpers.access_group_team_sync import ( TEAM_ADVISORY_LOCK_SQL, USER_ROWS_LOCK_SQL, ) +from litellm.proxy.spend_tracking.committed_spend import drop_committed_spend from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.redis.stores.spend_counters import user_counter from litellm.repositories.table_repositories import ( OrganizationMembershipRepository, TeamMembershipRepository, @@ -585,6 +587,9 @@ async def _after_users_deleted( ) await evict_and_broadcast(cache_keys=deletion.jwt_mapping_cache_keys, user_api_key_cache=user_api_key_cache) await evict_and_broadcast(cache_keys=sorted(user_ids), user_api_key_cache=user_api_key_cache) + from litellm.proxy.proxy_server import spend_counter_store + + await drop_committed_spend(spend_counter_store, tuple(user_counter(user_id) for user_id in sorted(user_ids))) for removal in deletion.removals.values(): _emit_team_members_metric(removal.team) diff --git a/litellm/proxy/management_helpers/team_member_permission_checks.py b/litellm/proxy/management_helpers/team_member_permission_checks.py index 834b032b04f..f13da0bbd1a 100644 --- a/litellm/proxy/management_helpers/team_member_permission_checks.py +++ b/litellm/proxy/management_helpers/team_member_permission_checks.py @@ -180,6 +180,9 @@ class TeamMemberPermissionChecks: caller_team_role: Final = get_caller_team_role(team_table=team_table, user_api_key_dict=user_api_key_dict) + if not set(access_group_ids).issubset(team_table.access_group_ids or []): + raise HTTPException(status_code=403, detail="Access groups must be assigned to the key's team first") + # Team admins always bypass (consistent with other member-permission checks). if caller_team_role == "admin": return diff --git a/litellm/proxy/mcp_server/outbound_credentials/bridge_credentials.py b/litellm/proxy/mcp_server/outbound_credentials/bridge_credentials.py index 2326bcebdc3..cc3e99263c9 100644 --- a/litellm/proxy/mcp_server/outbound_credentials/bridge_credentials.py +++ b/litellm/proxy/mcp_server/outbound_credentials/bridge_credentials.py @@ -16,6 +16,8 @@ from typing import Final, Literal, TypeAlias from pydantic import ConfigDict, SecretStr +from litellm.proxy.common_utils.encrypt_decrypt_utils import derive_key_from_salt_key +from litellm.proxy.common_utils.fips import fips_mode_active from litellm.proxy.mcp_server.outbound_credentials.envelope import ( EnvelopeIdentity, EnvelopeKeys, @@ -61,7 +63,15 @@ def envelope_keys_from_master_key(master_key: str) -> EnvelopeKeys: fixed for a process), so the KDF runs once per key and adds nothing to the per-request admission path. The derivation is deterministic; rotating ``master_key`` invalidates every outstanding envelope, which is the intended behavior for a signing-key change. + + When general_settings.fips_mode is on, both keys come from :func:`derive_key_from_salt_key` with the master key as + context instead. """ + if fips_mode_active(): + return EnvelopeKeys( + signing_key=SecretStr(derive_key_from_salt_key(_SIGNING_KEY_DOMAIN, master_key.encode())), + encryption_key=SecretStr(derive_key_from_salt_key(_ENCRYPTION_KEY_DOMAIN, master_key.encode())), + ) signing: Final = hashlib.scrypt( master_key.encode(), salt=_SIGNING_KEY_DOMAIN, diff --git a/litellm/proxy/mcp_server/outbound_credentials/session_credentials.py b/litellm/proxy/mcp_server/outbound_credentials/session_credentials.py index 6e3d09f4dd2..7cb89e53cdb 100644 --- a/litellm/proxy/mcp_server/outbound_credentials/session_credentials.py +++ b/litellm/proxy/mcp_server/outbound_credentials/session_credentials.py @@ -22,6 +22,8 @@ from typing import Final, Literal, TypeAlias from pydantic import ConfigDict, Field, SecretStr, ValidationError +from litellm.proxy.common_utils.encrypt_decrypt_utils import derive_key_from_salt_key +from litellm.proxy.common_utils.fips import fips_mode_active from litellm.proxy.mcp_server.outbound_credentials.session_token import ( AsymmetricSessionKeys, OpenedSessionToken, @@ -59,7 +61,14 @@ def session_keys_from_master_key(master_key: str) -> SessionKeys: into the other. The result is cached (the master key is fixed for a process); rotating ``master_key`` invalidates every outstanding session, which is the intended behavior for a signing-key change. + + When general_settings.fips_mode is on, the key comes from :func:`derive_key_from_salt_key` with the master key as + context instead. """ + if fips_mode_active(): + return SessionKeys( + signing_key=SecretStr(derive_key_from_salt_key(_SESSION_SIGNING_KEY_DOMAIN, master_key.encode())) + ) signing: Final = hashlib.scrypt( master_key.encode(), salt=_SESSION_SIGNING_KEY_DOMAIN, diff --git a/litellm/proxy/middleware/billable_request_metrics_middleware.py b/litellm/proxy/middleware/billable_request_metrics_middleware.py index fd78b6b58da..95e5c9d1293 100644 --- a/litellm/proxy/middleware/billable_request_metrics_middleware.py +++ b/litellm/proxy/middleware/billable_request_metrics_middleware.py @@ -24,6 +24,7 @@ from typing import Final, Protocol, runtime_checkable from starlette.types import ASGIApp, Message, Receive, Scope, Send from litellm._logging import verbose_proxy_logger +from litellm.constants import OPENSEARCH_PASS_THROUGH_ROUTE_PREFIX from litellm.proxy._types import LiteLLMRoutes from litellm.proxy.auth.route_checks import RouteChecks @@ -101,16 +102,21 @@ _LLM_ROUTE_EXACT: Final[tuple[str, ...]] = ( # inference calls that write SpendLogs rows, so they bill. Anchored to the # routes enum so new providers are picked up without touching this module. # /langfuse forwards observability traffic, not inference: it writes no -# SpendLogs row and must not bill. -_NON_BILLABLE_PASSTHROUGH_PREFIXES: Final = frozenset({"/langfuse"}) +# SpendLogs row and must not bill. /opensearch bills only index searches, like +# /v1/vector_stores/{id}/search, and never its document writes or lookups. +_NON_BILLABLE_PASSTHROUGH_PREFIXES: Final = frozenset({"/langfuse", OPENSEARCH_PASS_THROUGH_ROUTE_PREFIX}) _PASSTHROUGH_PREFIXES: Final[tuple[str, ...]] = tuple( prefix for prefix in LiteLLMRoutes.mapped_pass_through_routes.value if prefix not in _NON_BILLABLE_PASSTHROUGH_PREFIXES ) +_OPENSEARCH_INDEX_SEARCH: Final = re.compile(rf"{re.escape(OPENSEARCH_PASS_THROUGH_ROUTE_PREFIX)}/[^/]+/_search") + def _classify_llm_route(path: str) -> str | None: + if path.startswith(f"{OPENSEARCH_PASS_THROUGH_ROUTE_PREFIX}/"): + return OPENSEARCH_PASS_THROUGH_ROUTE_PREFIX if _OPENSEARCH_INDEX_SEARCH.fullmatch(path) else None exact_match: Final = next((route for route in _LLM_ROUTE_EXACT if path == route), None) if exact_match is not None: return exact_match diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 3653fcee112..5bdb2219194 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -159,11 +159,15 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( ) from litellm.proxy.utils import ProxyLogging as ProxyLoggingType from litellm.proxy.utils import is_known_model, normalize_route_for_root_path +from litellm.proxy.vector_store_endpoints.endpoints import build_request_data_from_managed_vector_store from litellm.proxy.vector_store_endpoints.utils import ( assert_proxy_admin_for_vector_store_index_management, assert_user_can_access_vector_store, + assert_vector_store_index_access, get_litellm_managed_vector_store, + has_vector_store_index_access, is_allowed_to_call_vector_store_endpoint, + normalize_litellm_params, ) from litellm.secret_managers.main import get_secret_str, normalize_nonempty_secret_str, str_to_bool from litellm.types.passthrough_endpoints.eu_jev import EuJevDecisionRequest, EuJevDecisionResponse @@ -182,8 +186,8 @@ from litellm.types.passthrough_endpoints.tinyfish import ( from litellm.types.passthrough_endpoints.vertex_ai import VertexPassThroughCredentials from litellm.types.router import LiteLLMParamsTypedDict from litellm.types.utils import CallTypes, LlmProviders -from litellm.types.vector_stores import LiteLLM_ManagedVectorStore -from litellm.utils import ProviderConfigManager +from litellm.types.vector_stores import LiteLLM_ManagedVectorStore, LiteLLM_ManagedVectorStoreIndex +from litellm.utils import ProviderConfigManager, with_named_credential from .passthrough_endpoint_router import PassthroughEndpointRouter @@ -1056,6 +1060,131 @@ async def milvus_proxy_route( return received_value +def _registered_vector_store_index(index_name: str) -> LiteLLM_ManagedVectorStoreIndex: + registry: Final = litellm.vector_store_index_registry + index_object: Final = registry.get_vector_store_index_by_name(index_name) if registry is not None else None + if index_object is None: + raise HTTPException( + status_code=400, + detail=f"Index {index_name} is not a litellm managed vector store index. " + "Only litellm managed vector store indexes are supported.", + ) + return index_object + + +async def _opensearch_store_params( + index_object: LiteLLM_ManagedVectorStoreIndex, user_api_key_dict: UserAPIKeyAuth +) -> Mapping[str, object]: + registry: Final = litellm.vector_store_registry + store_name: Final = index_object.litellm_params.vector_store_name + named_stores: Final = tuple( + store + for store in (registry.vector_stores if registry is not None else ()) + if store.get("vector_store_name") == store_name + ) + if len(named_stores) > 1: + raise HTTPException( + status_code=400, + detail=f"{len(named_stores)} vector stores are named {store_name}, so LiteLLM cannot tell which one " + f"index {index_object.index_name} is on. Give each vector store a unique name.", + ) + if not named_stores: + raise HTTPException( + status_code=400, + detail=f"Vector store {store_name} of index {index_object.index_name} is not registered.", + ) + vector_store: Final = named_stores[0] + params: Final = with_named_credential( + build_request_data_from_managed_vector_store(normalize_litellm_params(vector_store)) + ) + provider: Final = params.get("custom_llm_provider") + if not isinstance(provider, str) or provider.split("/", 1)[0] != LlmProviders.OPENSEARCH.value: + raise HTTPException( + status_code=400, + detail=f"Index {index_object.index_name} is on vector store {store_name}, " + "which is not an OpenSearch vector store.", + ) + await assert_user_can_access_vector_store(vector_store=vector_store, user_api_key_dict=user_api_key_dict) + return params + + +@router.api_route( + "/opensearch/{endpoint:path}", + methods=["GET", "HEAD", "POST", "PUT", "DELETE", "PATCH"], + tags=["OpenSearch Pass-through", "pass-through"], + openapi_extra={ + "requestBody": { + "required": False, + "content": { + "application/json": {"schema": {"type": "object", "additionalProperties": True}}, + "application/x-ndjson": {"schema": {"type": "string"}}, + }, + } + }, + responses={ + 200: { + "description": "OpenSearch's response, unmodified", + "content": {"application/json": {"schema": {"type": "object", "additionalProperties": True}}}, + } + }, +) +async def opensearch_proxy_route( # noqa: ANN201 # returns create_pass_through_route's untyped handler result + endpoint: str, + request: Request, + fastapi_response: Response, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +): + """ + Call an OpenSearch index through the proxy, e.g. `POST /opensearch//_search`. + + `` is a LiteLLM vector store index on an OpenSearch vector store. The key or its team needs `read` or + `write` on it in `allowed_vector_store_indexes`, and LiteLLM signs the call with the vector store's AWS identity. + """ + from litellm.llms.opensearch.vector_stores.index_passthrough import ( + IndexRefusal, + index_operation, + parse_index_path, + prepare_index_request, + ) + from litellm.llms.opensearch.vector_stores.transformation import OpenSearchIndexAWS, OpenSearchVectorStoreConfig + + path: Final = parse_index_path(endpoint) + if isinstance(path, IndexRefusal): + raise HTTPException(status_code=path.status_code, detail=path.message) + operation: Final = index_operation(request.method, path.api) + assert_vector_store_index_access( + user_api_key_dict, path.index_name, operation.permission if operation is not None else None + ) + index_object: Final = _registered_vector_store_index(path.index_name) + store_params: Final = await _opensearch_store_params(index_object, user_api_key_dict) + prepared: Final = await prepare_index_request( + config=OpenSearchVectorStoreConfig(aws=OpenSearchIndexAWS()), + method=request.method, + path=path, + operation=operation, + provider_index=index_object.litellm_params.vector_store_index, + content_type=request.headers.get("content-type"), + query=tuple(request.query_params.multi_items()), + body=await request.body(), + store_params=store_params, + can_read=has_vector_store_index_access(user_api_key_dict, path.index_name, "read"), + ) + if isinstance(prepared, IndexRefusal): + raise HTTPException(status_code=prepared.status_code, detail=prepared.message) + endpoint_func: Final = create_pass_through_route( + endpoint=endpoint, + target=prepared.url, + custom_headers=prepared.headers, + custom_llm_provider=LlmProviders.OPENSEARCH.value, + keep_target_query=True, + upstream_headers_filter=lambda _forwarded: prepared.headers, + ) + setattr(request.state, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, prepared.body) + setattr(request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, {}) + safe_set_request_parsed_body(request, {}) + return await endpoint_func(request, fastapi_response, user_api_key_dict) + + async def is_streaming_request_fn(request: Request) -> bool: if request.method == "POST": content_type: Final = request.headers.get("content-type", None) @@ -1620,7 +1749,7 @@ async def bedrock_proxy_route( ) from litellm.llms.bedrock.base_aws_llm import run_aws_signing, sign_aws_json_post - from litellm.llms.bedrock.chat import BedrockConverseLLM + from litellm.llms.bedrock.passthrough.transformation import BedrockAgentRuntimeAWS aws_region_name: Final = await run_aws_signing(get_secret_str, secret_name="AWS_REGION_NAME") base_target_url: Final = f"https://bedrock-agent-runtime.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}" @@ -1637,7 +1766,7 @@ async def bedrock_proxy_route( path=HttpPassThroughEndpointHelpers.join_base_and_endpoint_path(base_url, encoded_endpoint) ) - bedrock_llm: Final = BedrockConverseLLM() + agent_runtime_aws: Final = BedrockAgentRuntimeAWS() # Assuming the body contains JSON data, parse it try: data: Final = await _json_request_body(request) @@ -1645,7 +1774,7 @@ async def bedrock_proxy_route( raise HTTPException(status_code=400, detail={"error": e}) prepped: Final = await run_aws_signing( sign_aws_json_post, - get_credentials=bedrock_llm.get_credentials, + get_credentials=agent_runtime_aws.get_credentials, service_name="bedrock", aws_region_name=aws_region_name, url=str(updated_url), @@ -1716,6 +1845,7 @@ async def comprehend_medical_proxy_route( """ from .llm_provider_handlers.comprehend_medical_passthrough_logging_handler import ( COMPREHEND_MEDICAL_SUPPORTED_OPERATIONS, + ComprehendMedicalAWS, ) if operation not in COMPREHEND_MEDICAL_SUPPORTED_OPERATIONS: @@ -1744,12 +1874,12 @@ async def comprehend_medical_proxy_route( if "stream" in data: raise HTTPException(status_code=400, detail="'stream' is not a Comprehend Medical request member") - from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing, sign_aws_json_post + from litellm.llms.bedrock.base_aws_llm import run_aws_signing, sign_aws_json_post target_url: Final = f"https://comprehendmedical.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}/" prepped: Final = await run_aws_signing( sign_aws_json_post, - get_credentials=partial(BaseAWSLLM().get_credentials, aws_region_name=aws_region_name), + get_credentials=partial(ComprehendMedicalAWS().get_credentials, aws_region_name=aws_region_name), service_name="comprehendmedical", aws_region_name=aws_region_name, url=target_url, @@ -1973,6 +2103,7 @@ async def transcribe_proxy_route( TRANSCRIBE_OWNED_JOB_OPERATIONS, TRANSCRIBE_PRICED_OPERATION, TRANSCRIBE_TARGET_PREFIX, + TranscribeAWS, TranscribeRefusal, transcribe_admin_only_refusal, transcribe_cost_per_second, @@ -2038,12 +2169,12 @@ async def transcribe_proxy_route( if access_refusal is not None: raise HTTPException(status_code=access_refusal.status_code, detail=access_refusal.detail) - from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing, sign_aws_json_post + from litellm.llms.bedrock.base_aws_llm import run_aws_signing, sign_aws_json_post target_url: Final = f"https://transcribe.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}/" prepped: Final = await run_aws_signing( sign_aws_json_post, - get_credentials=partial(BaseAWSLLM().get_credentials, aws_region_name=aws_region_name), + get_credentials=partial(TranscribeAWS().get_credentials, aws_region_name=aws_region_name), service_name="transcribe", aws_region_name=aws_region_name, url=target_url, diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/comprehend_medical_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/comprehend_medical_passthrough_logging_handler.py index 3fb4fbd555d..741a8fb8838 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/comprehend_medical_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/comprehend_medical_passthrough_logging_handler.py @@ -2,7 +2,7 @@ import math from collections.abc import Mapping from datetime import datetime from types import MappingProxyType -from typing import Final +from typing import ClassVar, Final import httpx @@ -11,6 +11,8 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.litellm_logging import ( get_standard_logging_object_payload, ) +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.web_identity_session_policy import COMPREHEND_MEDICAL_SESSION_POLICY, WebIdentitySessionPolicy from litellm.proxy._types import PassThroughEndpointLoggingTypedDict from litellm.types.utils import StandardPassThroughResponseObject @@ -27,6 +29,10 @@ COMPREHEND_MEDICAL_COST_PER_UNIT_USD: Final[Mapping[str, float]] = MappingProxyT COMPREHEND_MEDICAL_SUPPORTED_OPERATIONS: Final = frozenset(COMPREHEND_MEDICAL_COST_PER_UNIT_USD) +class ComprehendMedicalAWS(BaseAWSLLM): + _web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = COMPREHEND_MEDICAL_SESSION_POLICY + + def _operation_from_headers(headers: Mapping[str, str]) -> str: target: Final = headers.get("x-amz-target", "") return target.split(".")[-1] diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/transcribe_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/transcribe_passthrough_logging_handler.py index 5b45316a44c..cd90d5c1a13 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/transcribe_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/transcribe_passthrough_logging_handler.py @@ -9,7 +9,7 @@ from email.utils import parsedate_to_datetime from functools import lru_cache, partial from pathlib import Path from types import MappingProxyType -from typing import IO, Final, Protocol, TypeAlias +from typing import IO, ClassVar, Final, Protocol, TypeAlias from urllib.parse import quote import httpx @@ -34,6 +34,8 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.litellm_logging import ( get_standard_logging_object_payload, ) +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.web_identity_session_policy import TRANSCRIBE_SESSION_POLICY, WebIdentitySessionPolicy from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.proxy._types import ( PassThroughEndpointLoggingResultValues, @@ -69,6 +71,10 @@ JobLookup: TypeAlias = Callable[[str], Awaitable[Mapping[str, object]]] MediaDurationProbe: TypeAlias = Callable[[str, float], Awaitable[float | None]] +class TranscribeAWS(BaseAWSLLM): + _web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = TRANSCRIBE_SESSION_POLICY + + class GetTranscriptionJobRequest(TypedDict): TranscriptionJobName: ReadOnly[str] @@ -446,7 +452,7 @@ def _as_json_object(response: httpx.Response) -> Mapping[str, object]: def transcribe_job_lookup(aws_region_name: str) -> JobLookup: - from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing, sign_aws_json_post + from litellm.llms.bedrock.base_aws_llm import run_aws_signing, sign_aws_json_post url: Final = f"https://transcribe.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}/" headers: Final = MappingProxyType( @@ -461,7 +467,7 @@ def transcribe_job_lookup(aws_region_name: str) -> JobLookup: payload: Final = json.dumps(body) prepped: Final = await run_aws_signing( sign_aws_json_post, - get_credentials=partial(BaseAWSLLM().get_credentials, aws_region_name=aws_region_name), + get_credentials=partial(TranscribeAWS().get_credentials, aws_region_name=aws_region_name), service_name="transcribe", aws_region_name=aws_region_name, url=url, @@ -527,11 +533,11 @@ def transcribe_media_duration_probe(aws_region_name: str, download_slots: asynci from botocore.auth import S3SigV4Auth from botocore.awsrequest import AWSRequest - from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing + from litellm.llms.bedrock.base_aws_llm import run_aws_signing def sign_s3_get(url: str) -> dict[str, str]: # mutable-ok: httpx request headers take a dict aws_request: Final = AWSRequest(method="GET", url=url) - credentials: Final = BaseAWSLLM().get_credentials(aws_region_name=aws_region_name) + credentials: Final = TranscribeAWS().get_credentials(aws_region_name=aws_region_name) S3SigV4Auth(credentials, "s3", aws_region_name).add_auth(aws_request) return dict(aws_request.prepare().headers.items()) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 53317010af4..3e25f7ce9af 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -119,6 +119,7 @@ from litellm.proxy.common_utils.error_body_call_id import JSON_OBJECT, error_bod from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports _read_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export _safe_get_request_headers, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export + is_opaque_pass_through_request, read_request_body, safe_get_request_headers, ) @@ -164,6 +165,8 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( WebSocketSessionUsage, WebsocketUpstreamFrameKind, ) +from litellm.types.proxy.auth.user_api_key_auth import COMMITTED_BUDGET_SPEND_PARAM +from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend from litellm.types.utils import TRUSTED_CALLBACK_VARS_FIELD, CallTypes, Usage from .llm_provider_handlers.tinyfish_passthrough_logging_handler import ( @@ -423,6 +426,10 @@ async def chat_completion_pass_through_endpoint( ) +def _is_provider_payload(request: Request) -> bool: + return is_opaque_pass_through_request(get_request_route(request), request.headers.get("content-type", "")) + + class HttpPassThroughEndpointHelpers(BasePassthroughUtils): @staticmethod def get_response_headers( @@ -774,6 +781,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): "litellm_params": { **litellm_params_in_body, "metadata": _metadata, + COMMITTED_BUDGET_SPEND_PARAM: user_api_key_dict.committed_budget_spend or CommittedBudgetSpend(), "proxy_server_request": { "url": str(request.url), "method": request.method, @@ -1209,6 +1217,14 @@ def _unfiltered_headers(headers: Mapping[str, str]) -> Mapping[str, str]: return headers +def _caller_query_source( + keep_target_query: bool, query_params: Mapping[str, object] | None, request: Request +) -> Mapping[str, object]: + if keep_target_query: + return MappingProxyType({}) + return query_params or request.query_params + + def _json_object(value: object) -> Mapping[str, object]: return cast("Mapping[str, object]", value) if isinstance(value, Mapping) else MappingProxyType({}) # cast-ok: json @@ -1266,6 +1282,7 @@ async def pass_through_request( timeout: float | None = None, upstream_headers_filter: Callable[[Mapping[str, str]], Mapping[str, str]] = _unfiltered_headers, operation: GenAIOperation | None = None, + keep_target_query: bool = False, ): """ Pass through endpoint handler, makes the httpx request for pass-through endpoints and ensures logging hooks are called @@ -1288,6 +1305,7 @@ async def pass_through_request( general_settings.pass_through_request_timeout, then 600s. upstream_headers_filter: Optional last rewrite of the headers sent upstream, after x-pass- expansion operation: Optional GenAI operation the route reports, unless the upstream names its own + keep_target_query: Send the target URL's query exactly as given, adding none of the caller's query params """ from litellm.exceptions import ModifyResponseException from litellm.litellm_core_utils.litellm_logging import Logging @@ -1326,7 +1344,10 @@ async def pass_through_request( parent_span=user_api_key_dict.parent_otel_span, ) - requested_query_params: dict | None = query_params or dict(request.query_params) or None + requested_query_params: dict | None = ( + dict(_caller_query_source(keep_target_query, query_params, request)) # pyright: ignore[reportUnknownArgumentType] # query_params keeps its legacy untyped dict annotation + or None + ) endpoint_type: Final[EndpointType] = HttpPassThroughEndpointHelpers.get_endpoint_type( str(url), custom_llm_provider @@ -2263,6 +2284,8 @@ async def _parse_request_data_by_content_type( Returns: Tuple of (query_params_data, custom_body_data, file_data, stream) """ + if _is_provider_payload(request): + return None, None, None, None content_type: Final = request.headers.get("content-type", "") query_params_data = None @@ -2333,6 +2356,7 @@ def create_pass_through_route( timeout: float | None = None, upstream_headers_filter: Callable[[Mapping[str, str]], Mapping[str, str]] = _unfiltered_headers, operation: GenAIOperation | None = None, + keep_target_query: bool = False, ): # check if target is an adapter.py or a url from litellm._uuid import uuid @@ -2482,6 +2506,7 @@ def create_pass_through_route( timeout=cast(float | None, param_timeout), upstream_headers_filter=upstream_headers_filter, operation=operation, + keep_target_query=keep_target_query, ) finally: if hasattr(request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY): diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index c94dc16a8f5..ef8ffdf640e 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -6,10 +6,14 @@ from typing import Final, Protocol import httpx -import litellm from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.asyncify import asyncify -from litellm.litellm_core_utils.core_helpers import bind_budget_reservation_to_callbacks +from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, + bind_budget_reservation_to_callbacks, + proxy_request_headers, + stream_cost_injection_requested, +) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.anthropic.common_utils import anthropic_error_frame_exception @@ -194,12 +198,11 @@ class PassThroughStreamingHandler: litellm_logging_obj=litellm_logging_obj, ) - # Resolve once per stream rather than re-reading the global + - # re-branching on every chunk. ``include_cost_in_streaming_usage`` is - # set at config load and stable for the process, matching how the - # proxy-level streaming fast path resolves it. + request_headers: Final = proxy_request_headers( + as_str_mapping(litellm_logging_obj.model_call_details.get("litellm_params")) + ) cost_injection_active: Final = ( - bool(getattr(litellm, "include_cost_in_streaming_usage", False)) + stream_cost_injection_requested(request_headers) and bool(model_name) and ( endpoint_type in (EndpointType.ANTHROPIC, EndpointType.OPENAI) @@ -242,7 +245,7 @@ class PassThroughStreamingHandler: complete_frames, pending = split_complete_sse_frames(pending + chunk) if complete_frames: yield ProxyBaseLLMRequestProcessing.process_chunk_with_cost_injection( - complete_frames, resolved_model_name, litellm_logging_obj + complete_frames, resolved_model_name, litellm_logging_obj, request_headers=request_headers ) if pending: yield pending diff --git a/litellm/proxy/plugin_routes.py b/litellm/proxy/plugin_routes.py index cfbcfefe9e8..966136f6fa7 100644 --- a/litellm/proxy/plugin_routes.py +++ b/litellm/proxy/plugin_routes.py @@ -18,6 +18,10 @@ Plugin iframe auth: litellm bearer token, so a compromised plugin learns only the caller's identity, never their credential. LITELLM_SALT_KEY itself is never shared with plugins — each plugin holds only its own derived key. + + With general_settings.fips_mode on, the per-plugin key is instead + derive_key_from_salt_key(b"litellm-plugin:session-claim:", plugin_name), and + the Fernet token is built from the FIPS provider's DRBG, AES-CBC and HMAC. """ import base64 @@ -25,11 +29,14 @@ import hashlib import hmac as _hmac import json import os +import ssl import time from collections.abc import Mapping from typing import Final from cryptography.fernet import Fernet, InvalidToken +from cryptography.hazmat.primitives import padding +from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes from fastapi import APIRouter, Depends, HTTPException, Request, Response from pydantic import ConfigDict, TypeAdapter, with_config from typing_extensions import NotRequired, ReadOnly, TypedDict @@ -37,6 +44,8 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.proxy._types import PluginConfig, SpecialHeaders, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.encrypt_decrypt_utils import derive_key_from_salt_key +from litellm.proxy.common_utils.fips import fips_mode_active from litellm.types.llms.custom_http import httpxSpecialProvider router: Final = APIRouter() @@ -145,11 +154,30 @@ def _plugin_fernet(plugin_name: str) -> Fernet: A plugin possessing its own key cannot derive the master salt or forge claims intended for a different plugin. """ + if fips_mode_active(): + return Fernet(base64.urlsafe_b64encode(_fips_plugin_key(plugin_name))) salt: Final = os.getenv("LITELLM_SALT_KEY", "").encode() derived: Final = _hmac.new(salt, plugin_name.encode(), hashlib.sha256).digest() return Fernet(base64.urlsafe_b64encode(derived)) +_FIPS_PLUGIN_KEY_LABEL: Final = b"litellm-plugin:session-claim:" + + +def _fips_plugin_key(plugin_name: str) -> bytes: + return bytes.fromhex(derive_key_from_salt_key(_FIPS_PLUGIN_KEY_LABEL, plugin_name.encode())) + + +def _fips_fernet_token(key: bytes, data: bytes) -> str: + """Encrypt data into a standard Fernet token using only the FIPS provider's DRBG, AES-128-CBC and HMAC-SHA256.""" + iv: Final = ssl.RAND_bytes(16) + padder: Final = padding.PKCS7(128).padder() + encryptor: Final = Cipher(algorithms.AES(key[16:]), modes.CBC(iv)).encryptor() + ciphertext: Final = encryptor.update(padder.update(data) + padder.finalize()) + encryptor.finalize() + body: Final = b"\x80" + int(time.time()).to_bytes(8, "big") + iv + ciphertext + return base64.urlsafe_b64encode(body + _hmac.digest(key[:16], body, "sha256")).decode() + + _CLAIM_TTL_SECONDS: Final = 30 # identity claims expire after 30 s @@ -166,6 +194,8 @@ def issue_plugin_session_claim(plugin_name: str, user_id: str | None, user_role: "user_role": user_role or "", "exp": int(time.time()) + _CLAIM_TTL_SECONDS, } + if fips_mode_active(): + return _fips_fernet_token(_fips_plugin_key(plugin_name), json.dumps(claim).encode()) return _plugin_fernet(plugin_name).encrypt(json.dumps(claim).encode()).decode() diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index fe633238e20..61f81907ff4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -93,6 +93,7 @@ from litellm.constants import ( LITELLM_UI_SESSION_DURATION, OPENAI_LIVE_SESSION_START_TIMEOUT_SECONDS, RUNTIME_UPDATABLE_ROUTER_SETTINGS, + SPEND_LOG_QUEUE_MAX_BACKOFF_SECONDS, VLLM_CHOICE_FIELDS, ) from litellm.litellm_core_utils.access_windows import access_windows_config_error @@ -338,7 +339,7 @@ from litellm._logging import ( verbose_router_logger, ) from litellm.caching.caching import DualCache, RedisCache -from litellm.caching.redis_batch import drain_post_call_redis_batches +from litellm.caching.redis_batch import active_request_redis_batches, drain_post_call_redis_batches from litellm.caching.redis_cache import RedisCircuitBreakerOpenError, is_redis_timeout_failure from litellm.caching.redis_cluster_cache import RedisClusterCache from litellm.constants import ( @@ -471,13 +472,18 @@ from litellm.proxy.auth.master_key_boot_check import ( with_stored_secrets_counted, ) from litellm.proxy.auth.model_checks import ( + GLOBAL_ALIASES_ONLY, + OtherDeploymentsBill, + RoutedElsewhere, expand_wildcard_deployments_for_model_info, get_all_fallbacks, get_complete_model_list, get_key_models, get_mcp_server_ids, get_team_models, + model_info_cost_map_entry, run_model_listing, + wildcard_row_priced_model, ) from litellm.proxy.auth.password_policy import validate_password_not_breached, validate_password_policy from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401, RUF100 # legacy module exports @@ -990,12 +996,20 @@ from litellm.proxy.spend_tracking.budget_reservation import ( release_unbound_budget_reservation, stamp_budget_reservation_actual_cost, ) +from litellm.proxy.spend_tracking.committed_spend import ( + committed_spend_key, + committed_spend_keys, + has_committed_twin, + is_committed_spend_key, + seed_committed_spend, +) from litellm.proxy.spend_tracking.spend_capture_rate import ( run_scheduled_spend_capture_rate_check, ) from litellm.proxy.spend_tracking.spend_counter_batch import ( PendingSpendIncrement, active_spend_counter_batch, + bind_spend_counter_keys, forget_spend_counter, post_call_counter_keys, read_batched_spend_counter, @@ -1100,6 +1114,7 @@ from litellm.types.llms.openai import ( ResponseInProgressEvent, ResponsesAPIResponse, ) +from litellm.types.proxy.carried_budget_state import EntitySpendCounterKeys from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry from litellm.types.proxy.management_endpoints.model_management_endpoints import ( ModelGroupInfoProxy, @@ -3140,15 +3155,16 @@ experimental = False llm_router: Router | None = None llm_model_list: list | None = None heuristic_v1_tuning_baselines: Mapping[str, str] | None = None -# Serializes every model reconcile (ProxyConfig.add_deployment and clear_cache) so the -# read-modify-write of llm_router above is atomic. Without it, two concurrent model -# writes each reconcile the router against their OWN db snapshot, and the one holding -# the older snapshot evicts the deployment the newer one just added -- the db keeps the -# row, this pod stops serving it. Control-plane only (model create/update/delete and -# the config-sync tick), never on a completion path, so the serialization is free. -# Module-level rather than per-ProxyConfig because llm_router is a module global and a -# second ProxyConfig instance must not get its own independent lock over it. -MODEL_RECONCILE_LOCK: Final = asyncio.Lock() +# Serializes every db-to-memory reconcile (ProxyConfig.add_deployment and clear_cache) +# with every control-plane write to the state those reconciles rebuild: llm_router above +# and litellm.credential_list. Without it, two concurrent writes each reconcile against +# their OWN db snapshot, and the one holding the older snapshot evicts the deployment or +# credential the newer one just added -- the db keeps the row, this pod stops serving it. +# Control-plane only (model and credential create/update/delete and the config-sync +# tick), never on a completion path, so the serialization is free. Module-level rather +# than per-ProxyConfig because llm_router is a module global and a second ProxyConfig +# instance must not get its own independent lock over it. +DB_RECONCILE_LOCK: Final = asyncio.Lock() general_settings: dict = {} _GENERAL_SETTINGS_VIEW: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) @@ -3234,6 +3250,7 @@ proxy_budget_rescheduler_min_time = PROXY_BUDGET_RESCHEDULER_MIN_TIME proxy_budget_rescheduler_max_time = PROXY_BUDGET_RESCHEDULER_MAX_TIME proxy_batch_polling_interval = PROXY_BATCH_POLLING_INTERVAL proxy_batch_write_at = PROXY_BATCH_WRITE_AT +_BATCH_WRITE_JITTER_MAX_SECONDS: Final = 5 proxy_config_reload_interval_seconds = PROXY_CONFIG_RELOAD_INTERVAL_SECONDS litellm_master_key_hash = None disable_spend_logs = False @@ -3632,6 +3649,7 @@ async def increment_spend_counters( request_started_at: datetime | None = None, model_access_groups: Sequence[str] | None = None, project_id: str | None = None, + committed_spend_scope: EntitySpendCounterKeys | None = None, windows_from_ledger: bool = False, ): """ @@ -3669,6 +3687,7 @@ async def increment_spend_counters( request_started_at=request_started_at, model_access_groups=model_access_groups, project_id=project_id, + committed_spend_scope=committed_spend_scope, windows_from_ledger=windows_from_ledger, ) @@ -3784,10 +3803,17 @@ async def _increment_spend_counters_batched( request_started_at: datetime | None, model_access_groups: Sequence[str] | None, project_id: str | None = None, + committed_spend_scope: EntitySpendCounterKeys | None = None, windows_from_ledger: bool = False, ): """Runs inside one spend counter batch: the reservation reconcile and the warm checks share a single MGET, and the reconcile adjustments go out in the same INCRBYFLOAT pipeline as the counter increments.""" + committed_counter_keys: Final = ( + tuple(counter_key for counter_key in committed_spend_scope.tracked() if has_committed_twin(counter_key)) + if committed_spend_scope is not None and response_cost is not None and response_cost > 0 + else () + ) + bind_spend_counter_keys(frozenset(committed_spend_keys(committed_counter_keys))) reservation_update: Final = await _reconcile_budget_reservation_for_counter_update( budget_reservation=budget_reservation, response_cost=response_cost, @@ -3974,6 +4000,9 @@ async def _increment_spend_counters_batched( ) if project_id is not None else None, + _prepare_committed_spend_increments(counter_keys=committed_counter_keys, increment=cost) + if committed_counter_keys + else None, ) if coro is not None ) @@ -4076,6 +4105,23 @@ async def _reconcile_budget_reservation_for_counter_update( return _ReservationCounterUpdate(reserved_counter_keys=frozenset(reserved_counter_keys), pending=pending) +async def _committed_spend_seeded(counter_key: str) -> bool: + return await _is_spend_counter_cache_warm(committed_spend_key(counter_key)) or await seed_committed_spend( + spend_counter_store=spend_counter_store, prisma_client=prisma_client, counter_key=counter_key + ) + + +async def _prepare_committed_spend_increments( + counter_keys: tuple[str, ...], increment: float +) -> tuple[PendingSpendIncrement, ...]: + seeded: Final = await asyncio.gather(*(_committed_spend_seeded(counter_key) for counter_key in counter_keys)) + return tuple( + PendingSpendIncrement(counter_key=committed_spend_key(counter_key), increment=increment) + for counter_key, is_seeded in zip(counter_keys, seeded) + if is_seeded + ) + + async def _prepare_end_user_and_tag_spend_increments( end_user_id: str | None, tags: list[str] | None, @@ -4481,7 +4527,15 @@ def _defer_spend_counter_increments(pending: Sequence[PendingSpendIncrement]) -> """Post-call increments ride the request's native post-call pipeline with the other counters. Each counter's new value lands in memory when the pipeline settles; a failed one is dropped everywhere so no reader trusts a counter whose increment may not have applied, as ``increment_spend_counters_pipeline`` does.""" - return spend_counter_store.increment_many_post_call(tuple((item.counter_key, item.increment) for item in pending)) + settlement: Final = spend_counter_store.increment_many_post_call( + tuple((item.counter_key, item.increment) for item in pending) + ) + if settlement is None: + return False + batches: Final = active_request_redis_batches() + if batches is not None and any(is_committed_spend_key(item.counter_key) for item in pending): + batches.track_post_call_settlement(settlement) + return True async def increment_spend_counters_pipeline(pending: Sequence[PendingSpendIncrement]) -> tuple[float | None, ...]: @@ -8922,7 +8976,7 @@ class ProxyConfig: - Check if model id's in router already - If not, add to router - Serialized against every other model reconcile by MODEL_RECONCILE_LOCK, because + Serialized against every other model reconcile by DB_RECONCILE_LOCK, because the work below is a read-modify-write of the shared ``llm_router`` global: it reads the db into a snapshot and then makes the router match that snapshot. Two of those interleaving is not a lost update but an eviction -- the request whose @@ -8941,7 +8995,7 @@ class ProxyConfig: if not ui_settings_already_synced: await sync_ui_settings_to_general_settings(prisma_client) - async with MODEL_RECONCILE_LOCK: + async with DB_RECONCILE_LOCK: return await self.add_deployment_locked(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj) async def add_deployment_locked( @@ -8949,7 +9003,7 @@ class ProxyConfig: prisma_client: PrismaClient, proxy_logging_obj: ProxyLogging, ) -> ReconcileOutcome: - """add_deployment's body, minus the locking. MODEL_RECONCILE_LOCK MUST already + """add_deployment's body, minus the locking. DB_RECONCILE_LOCK MUST already be held. Split out for the one caller that has to hold the lock across more than this reconcile -- clear_cache, which un-serves every db model before calling it and would deadlock on a re-acquire.""" @@ -9898,7 +9952,7 @@ class ProxyConfig: immediately, and by the periodic job in store_model_in_db-off deployments. Gated the same way as startup, so an admin who excluded search_tools from supported_db_objects opts out. - Serialized by MODEL_RECONCILE_LOCK for the reason add_deployment documents: the body is a + Serialized by DB_RECONCILE_LOCK for the reason add_deployment documents: the body is a read-modify-write of the shared ``llm_router`` global, so two of them interleaving lets the older snapshot's wholesale assignment land last and restore a tool the newer one deleted. The lock belongs here rather than in _init_search_tools_in_db, which _init_non_llm_objects_in_db @@ -9908,7 +9962,7 @@ class ProxyConfig: return if prisma_client is None: return - async with MODEL_RECONCILE_LOCK: + async with DB_RECONCILE_LOCK: await self._init_search_tools_in_db(prisma_client=prisma_client) @staticmethod @@ -11107,7 +11161,7 @@ def _pricing_override_stamps( return MappingProxyType({**own_pricing, PRICING_OVERRIDES_KEY: pricing_override_fields(model_info, own_pricing)}) -def get_litellm_model_info(model: dict = {}): +def get_litellm_model_info(model: dict = {}) -> Mapping[str, object]: model_info: Final = model.get("model_info", {}) model_to_lookup = model.get("litellm_params", {}).get("model", None) try: @@ -11645,7 +11699,10 @@ class ProxyStartupEvent: ) # Ensure minimum interval of 30 seconds for batch writing to prevent memory issues - batch_writing_interval: Final = proxy_batch_write_at + random.randint(0, 5) + batch_writing_interval: Final = proxy_batch_write_at + random.randint(0, _BATCH_WRITE_JITTER_MAX_SECONDS) + spend_counter_store.set_committed_ttl( + 2 * (SPEND_LOG_QUEUE_MAX_BACKOFF_SECONDS + proxy_batch_write_at + _BATCH_WRITE_JITTER_MAX_SECONDS) + ) ### PROXY WORKER HEARTBEAT ### worker_heartbeat: Final = ProxyWorkerHeartbeat(prisma_client=prisma_client) @@ -12655,16 +12712,16 @@ async def _deployment_hidden_by_listing_callbacks(deployment: Deployment, user_a return listed_name in await _names_hidden_by_listing_callbacks(user_api_key_dict, (listed_name,)) -async def _caller_aliases(user_api_key_dict: UserAPIKeyAuth, team_id: str | None) -> CallerAliases: - caller_router_settings: Final = await proxy_config.get_hierarchical_router_settings( +async def request_caller_aliases(user_api_key_dict: UserAPIKeyAuth, listed_team_id: str | None = None) -> CallerAliases: + router_settings: Final = await proxy_config.get_hierarchical_router_settings( user_api_key_dict, prisma_client, proxy_logging_obj ) return caller_alias_maps( user_api_key_dict.aliases, user_api_key_dict.team_model_aliases, user_api_key_dict.team_id, - team_id, - caller_router_settings.get("model_group_alias") if isinstance(caller_router_settings, Mapping) else None, + listed_team_id, + router_settings.get("model_group_alias") if isinstance(router_settings, Mapping) else None, ) @@ -12757,7 +12814,7 @@ async def _key_model_listing( async def codex_catalog_rows_for_key(user_api_key_dict: UserAPIKeyAuth) -> tuple[CodexCatalogRow, ...]: listing: Final = await _key_model_listing( user_api_key_dict, - await _caller_aliases(user_api_key_dict, None), + await request_caller_aliases(user_api_key_dict), await _listing_hidden_names(None), ) return await asyncio.to_thread( @@ -12840,7 +12897,7 @@ async def model_list( http_request is not None and http_request.headers.get("anthropic-version") is not None ) client_headers: Final[Mapping[str, str]] = http_request.headers if http_request is not None else _EMPTY_HEADERS - caller_aliases: Final = await _caller_aliases(user_api_key_dict, team_id) + caller_aliases: Final = await request_caller_aliases(user_api_key_dict, team_id) routing_names: Final = ClaudeCodeRoutingNames( llm_router, team_id or user_api_key_dict.team_id, @@ -13084,16 +13141,7 @@ async def model_info( ) hidden_names: Final = blocked_names | unhealthy_names internal_to_public: Final = TeamModelNameTranslator.build_internal_to_public_map(llm_router, settings) - caller_router_settings: Final = await proxy_config.get_hierarchical_router_settings( - user_api_key_dict, prisma_client, proxy_logging_obj - ) - caller_aliases: Final = caller_alias_maps( - user_api_key_dict.aliases, - user_api_key_dict.team_model_aliases, - user_api_key_dict.team_id, - team_id, - caller_router_settings.get("model_group_alias") if isinstance(caller_router_settings, Mapping) else None, - ) + caller_aliases: Final = await request_caller_aliases(user_api_key_dict, team_id) callback_hidden_names: Final = await _names_hidden_by_listing_callbacks( user_api_key_dict, tuple( @@ -16807,7 +16855,10 @@ async def get_all_team_and_direct_access_models( def _enrich_model_info_with_litellm_data( - model: dict[str, Any], debug: bool = False, llm_router: Router | None = None + model: dict[str, Any], + debug: bool = False, + llm_router: Router | None = None, + caller_aliases: CallerAliases = GLOBAL_ALIASES_ONLY, ) -> dict[str, Any]: """ Enrich a model dictionary with litellm model info (pricing, context window, etc.) @@ -16817,6 +16868,7 @@ def _enrich_model_info_with_litellm_data( model: Model dictionary to enrich debug: Whether to include debug information like openai_client llm_router: Optional router instance for debug info + caller_aliases: The caller's alias maps, applied before deciding which model a wildcard row prices as Returns: Enriched model dictionary with sensitive info removed @@ -16834,30 +16886,29 @@ def _enrich_model_info_with_litellm_data( # read litellm model_prices_and_context_window.json to get the following: # input_cost_per_token, output_cost_per_token, max_tokens - litellm_model_info = get_litellm_model_info(model=model) - - # 2nd pass on the model, try seeing if we can find model in litellm model_cost map - if litellm_model_info == {}: - # use litellm_param model_name to get model_info - litellm_params = model.get("litellm_params", {}) - litellm_model = litellm_params.get("model", None) - try: - litellm_model_info = litellm.get_model_info(model=litellm_model) - except Exception: - litellm_model_info = {} - # 3rd pass on the model, try seeing if we can find model but without the "/" in model cost map - if litellm_model_info == {}: - # use litellm_param model_name to get model_info - litellm_params = model.get("litellm_params", {}) - litellm_model = litellm_params.get("model", None) - if litellm_model: - split_model: Final = litellm_model.split("/") - if len(split_model) > 0: - litellm_model = split_model[-1] - try: - litellm_model_info = litellm.get_model_info(model=litellm_model, custom_llm_provider=split_model[0]) - except Exception: - litellm_model_info = {} + priced_as: Final = wildcard_row_priced_model(model, llm_router, caller_aliases) + billed_model: Final = priced_as.model if isinstance(priced_as, OtherDeploymentsBill) else priced_as + lookup_model_info: Final = ( + {k: v for k, v in model_info.items() if k != "base_model"} + if isinstance(priced_as, OtherDeploymentsBill) + else model_info + ) + priced_model: Final = ( + {**model, "litellm_params": {**model["litellm_params"], "model": billed_model}, "model_info": lookup_model_info} + if isinstance(billed_model, str) + else model + ) + priced_litellm_model: Final = priced_model.get("litellm_params", {}).get("model") + litellm_model_info: Final[Mapping[str, object]] = ( + MappingProxyType({}) + if isinstance(priced_as, RoutedElsewhere) + else get_litellm_model_info(model=priced_model) + or ( + model_info_cost_map_entry(priced_litellm_model) + if isinstance(priced_litellm_model, str) + else MappingProxyType({}) + ) + ) discovered_model_info: Final = ( llm_router.registry().discovered_model_info(model_info.get("id")) if llm_router is not None @@ -16907,7 +16958,7 @@ async def _get_caller_byok_team_scope( must have `model_info.team_id` ∈ that set, otherwise they belong to a team the caller is not a member of and must be dropped. """ - if user_api_key_dict is None or prisma_client is None: + if user_api_key_dict is None: return None if user_api_key_dict.user_role in ( LitellmUserRoles.PROXY_ADMIN, @@ -16916,7 +16967,7 @@ async def _get_caller_byok_team_scope( return None key_team_scope: Final[set[str]] = {user_api_key_dict.team_id} if user_api_key_dict.team_id else set() user_id: Final = user_api_key_dict.user_id - if user_id is None: + if user_id is None or prisma_client is None: return key_team_scope try: user_row: Final = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id}) @@ -17598,7 +17649,7 @@ async def _find_model_by_id( llm_router: Router | None, prisma_client: PrismaClient | None, proxy_config: "ProxyConfig", -) -> tuple[list, int | None]: +) -> tuple[Sequence[dict[str, Any]], int | None]: """Find a model by its ID and optionally filter by search term.""" found_model = None @@ -17795,18 +17846,32 @@ async def model_info_v2( # Load existing config await proxy_config.get_config() + allowed_team_ids: Final = await _get_caller_byok_team_scope(user_api_key_dict, prisma_client) + # If modelId is provided, search for the specific model if modelId is not None: - all_models, search_total_count = await _find_model_by_id( + matched_models: Final = await _find_model_by_id( model_id=modelId, search=search, llm_router=llm_router, prisma_client=prisma_client, proxy_config=proxy_config, ) + all_models, search_total_count = ( + [ + m + for m in matched_models[0] + if not _byok_row_outside_caller_teams(m.get("model_info", {}), allowed_team_ids) + ], + matched_models[1], + ) else: # Normal flow when modelId is not provided - all_models = copy.deepcopy(llm_router.registry().all()) + all_models = [ + m + for m in copy.deepcopy(llm_router.registry().all()) + if not _byok_row_outside_caller_teams(m.get("model_info", {}), allowed_team_ids) + ] if user_model is not None: # if user does not use a config.yaml, https://github.com/BerriAI/litellm/issues/2061 @@ -18460,8 +18525,10 @@ def _translate_model_name_for_response(model: dict) -> dict[str, object]: return {**model, "model_name": team_public} -def _get_proxy_model_info(model: dict) -> dict: - return _translate_model_name_for_response(_enrich_model_info_with_litellm_data(model=model, llm_router=llm_router)) +def _get_proxy_model_info(model: dict, caller_aliases: CallerAliases = GLOBAL_ALIASES_ONLY) -> dict: + return _translate_model_name_for_response( + _enrich_model_info_with_litellm_data(model=model, llm_router=llm_router, caller_aliases=caller_aliases) + ) def _model_info_json_response(data: Sequence[Mapping[str, object]] | Mapping[str, object]) -> Response: @@ -18598,7 +18665,20 @@ async def model_info_v1( status_code=400, detail={"error": f"Model id = {litellm_model_id} not found on litellm proxy"}, ) - _deployment_info_dict = _get_proxy_model_info(model=deployment_info.model_dump(exclude_none=True)) + if _byok_row_outside_caller_teams( + deployment_info.model_info.model_dump(), + await _get_caller_byok_team_scope(user_api_key_dict, prisma_client), + ): + if include_team_models: + return _model_info_json_response([]) + raise HTTPException( + status_code=400, + detail={"error": f"Model id = {litellm_model_id} not found on litellm proxy"}, + ) + _deployment_info_dict = _get_proxy_model_info( + model=deployment_info.model_dump(exclude_none=True), + caller_aliases=await request_caller_aliases(user_api_key_dict), + ) single_model_list: list[dict] = [_deployment_info_dict] if prisma_client is not None: single_model_list = await _populate_team_access_on_models( @@ -18674,8 +18754,11 @@ async def model_info_v1( if include_team_models: all_models = _filter_models_to_user_accessible(all_models) + caller_aliases: Final = await request_caller_aliases(user_api_key_dict) all_models = [ - _translate_model_name_for_response(_enrich_model_info_with_litellm_data(model=model, llm_router=llm_router)) + _translate_model_name_for_response( + _enrich_model_info_with_litellm_data(model=model, llm_router=llm_router, caller_aliases=caller_aliases) + ) for model in all_models ] diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 50d1c7c5b59..574a677061d 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -398,6 +398,9 @@ async def responses_api( # instead of a polling ID that immediately fails in the background task. processor = ProxyBaseLLMRequestProcessing(data=data) try: + # A background request runs after its polling id is returned, so a total_timeout it + # cannot apply is refused (or a header one dropped) before the id is handed out. + _ = processor.total_timeout_for(request, "aresponses", is_streaming_request=None) data, _logging_obj = await processor.common_processing_pre_call_logic( request=request, general_settings=general_settings, diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 08942851af9..ffc0e31408d 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -17,6 +17,7 @@ import litellm from litellm.litellm_core_utils.core_helpers import normalize_drop_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.routing_common_utils import is_proxy_admin_request +from litellm.litellm_core_utils.total_deadline import enforce_total_timeout_scope from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.router import router_arguments from litellm.router_api import router_init_args @@ -761,6 +762,7 @@ async def route_request( """ Common helper to route the request """ + enforce_total_timeout_scope(data) # pyright: ignore[reportUnknownArgumentType] # route_request takes an untyped dict body router_only_settings_from_override: Final = _merge_router_settings_override(data) # pyright: ignore[reportUnknownArgumentType] # route_request takes an untyped dict body try: return await _route_request_single_attempt( diff --git a/litellm/proxy/spend_tracking/carried_budget_state.py b/litellm/proxy/spend_tracking/carried_budget_state.py index 0dfb38272b0..150ddd1efe9 100644 --- a/litellm/proxy/spend_tracking/carried_budget_state.py +++ b/litellm/proxy/spend_tracking/carried_budget_state.py @@ -1,14 +1,17 @@ """Pins the budget state auth resolved onto ``UserAPIKeyAuth`` and emits it as request metadata.""" -from collections.abc import Mapping +from collections.abc import Awaitable, Callable, Mapping from types import MappingProxyType from typing import Final +from litellm.models.end_user import LiteLLM_EndUserTable from litellm.models.organization import LiteLLM_OrganizationTable from litellm.models.team import LiteLLM_TeamTable from litellm.models.user import LiteLLM_UserTable from litellm.proxy._types import UserAPIKeyAuth from litellm.types.proxy.carried_budget_state import ( + CommittedBudgetSpend, + EntitySpendCounterKeys, OrgBudgetSnapshot, TeamBudgetSnapshot, UserBudgetSnapshot, @@ -34,6 +37,41 @@ def carry_team_and_user_budget_state( ) +def _budgeted_org_id(valid_token: UserAPIKeyAuth, team_object: LiteLLM_TeamTable | None) -> str | None: + org_snapshot: Final = valid_token.org_budget_snapshot + if org_snapshot is None or org_snapshot.max_budget is None: + return None + return valid_token.org_id or (team_object.organization_id if team_object is not None else None) + + +def _budgeted_end_user_id(end_user_object: LiteLLM_EndUserTable | None) -> str | None: + budget_table: Final = end_user_object.litellm_budget_table if end_user_object is not None else None + if end_user_object is None or budget_table is None or budget_table.max_budget is None: + return None + return end_user_object.user_id + + +async def carry_committed_budget_spend( + valid_token: UserAPIKeyAuth, + team_object: LiteLLM_TeamTable | None, + user_object: LiteLLM_UserTable | None, + end_user_object: LiteLLM_EndUserTable | None, + read_committed: Callable[[EntitySpendCounterKeys], Awaitable[CommittedBudgetSpend]] | None, +) -> None: + counter_keys: Final = EntitySpendCounterKeys.of( + token=valid_token.token if valid_token.max_budget is not None else None, + team_id=team_object.team_id if team_object is not None and team_object.max_budget is not None else None, + user_id=user_object.user_id if user_object is not None and user_object.max_budget is not None else None, + org_id=_budgeted_org_id(valid_token=valid_token, team_object=team_object), + end_user_id=_budgeted_end_user_id(end_user_object=end_user_object), + ) + valid_token.committed_budget_spend = ( # rebind-ok: request-scoped credential + CommittedBudgetSpend(counter_keys=counter_keys) + if read_committed is None + else await read_committed(counter_keys) + ) + + def carry_organization_budget_state(valid_token: UserAPIKeyAuth, org_table: LiteLLM_OrganizationTable) -> None: budget_table: Final = org_table.litellm_budget_table valid_token.organization_alias = org_table.organization_alias diff --git a/litellm/proxy/spend_tracking/committed_spend.py b/litellm/proxy/spend_tracking/committed_spend.py new file mode 100644 index 00000000000..76066339cf1 --- /dev/null +++ b/litellm/proxy/spend_tracking/committed_spend.py @@ -0,0 +1,181 @@ +"""Committed spend counters: twins of the entity spend counters that only ever take charged cost, never budget holds. + +The remaining budget gauges read them, so a request's worst-case reservation never shows as spent. They are kept +only for entities with a budget, while a ``PrometheusLogger`` is registered, in the same ``SpendCounterStore`` as the +counters they twin. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import Iterable, Mapping, Sequence +from typing import TYPE_CHECKING, Final + +import litellm +from litellm._internal_context import with_service_target +from litellm._logging import verbose_proxy_logger +from litellm.caching.redis_batch import Batch +from litellm.integrations.prometheus import PrometheusLogger +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed +from litellm.proxy.spend_tracking.spend_counter_batch import ( + admission_counter_keys, + bind_spend_counter_keys, + read_batched_spend_counter, +) +from litellm.redis.stores.spend_counters import ( + COMMITTED_COUNTER_PREFIX, + END_USER_COUNTER_PREFIX, + SPEND_COUNTERS_TARGET, + SpendCounterStore, + committed_counter, +) +from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + +_BASE_COUNTER_PREFIXES: Final = ("spend:key:", "spend:team:", "spend:user:", "spend:org:", END_USER_COUNTER_PREFIX) + + +def committed_spend_key(counter_key: str) -> str: + return committed_counter(counter_key) + + +def has_committed_twin(counter_key: str) -> bool: + return counter_key.startswith(_BASE_COUNTER_PREFIXES) and ":window:" not in counter_key + + +def committed_spend_keys(counter_keys: Iterable[str]) -> tuple[str, ...]: + return tuple(committed_spend_key(counter_key) for counter_key in counter_keys if has_committed_twin(counter_key)) + + +def is_committed_spend_key(counter_key: str) -> bool: + return counter_key.startswith(COMMITTED_COUNTER_PREFIX) + + +def committed_spend_tracked() -> bool: + return litellm.logging_callback_manager.callback_is_active(callback_type=PrometheusLogger) + + +def bind_committed_admission_counters(token: UserAPIKeyAuth, end_user_id: str | None) -> None: + if committed_spend_tracked(): + bind_spend_counter_keys(frozenset(committed_spend_keys(admission_counter_keys(token, end_user_id)))) + + +async def _read_committed_values( + spend_counter_store: SpendCounterStore, twin_keys: tuple[str, ...] +) -> Mapping[str, float]: + batched: Final = await spend_counter_store.read_many(twin_keys, batch=Batch.NONE) + if batched is not None: + return {twin_key: value for twin_key, value in batched.items() if value is not None} + reads: Final = await asyncio.gather(*(spend_counter_store.read(twin_key) for twin_key in twin_keys)) + return {twin_key: value for twin_key, (value, _) in zip(twin_keys, reads, strict=True) if value is not None} + + +def _committed_budget_spend( + counter_keys: EntitySpendCounterKeys, committed: Mapping[str, float] +) -> CommittedBudgetSpend: + def spend_of(counter_key: str | None) -> float | None: + return committed.get(committed_spend_key(counter_key)) if counter_key is not None else None + + return CommittedBudgetSpend( + counter_keys=counter_keys, + key=spend_of(counter_keys.key), + team=spend_of(counter_keys.team), + user=spend_of(counter_keys.user), + org=spend_of(counter_keys.org), + end_user=spend_of(counter_keys.end_user), + ) + + +@with_service_target(SPEND_COUNTERS_TARGET) +async def read_committed_budget_spend( + spend_counter_store: SpendCounterStore, counter_keys: EntitySpendCounterKeys +) -> CommittedBudgetSpend: + twin_keys: Final = committed_spend_keys(counter_keys.tracked()) + if not twin_keys: + return CommittedBudgetSpend(counter_keys=counter_keys) + try: + committed: Final = await _read_committed_values(spend_counter_store, twin_keys) + except Exception as e: # noqa: BLE001 # a gauge read never fails the request; the gauges keep the stored spend + verbose_proxy_logger.debug("committed spend read failed for %s: %s", twin_keys, e) + return CommittedBudgetSpend(counter_keys=counter_keys) + return _committed_budget_spend(counter_keys=counter_keys, committed=committed) + + +async def warm_committed_counter_keys( + spend_counter_store: SpendCounterStore, counter_keys: EntitySpendCounterKeys +) -> EntitySpendCounterKeys: + """For a charge with no auth to name the budgeted entities: only a budgeted entity's twin exists.""" + committed: Final = await read_committed_budget_spend(spend_counter_store, counter_keys) + return EntitySpendCounterKeys( + key=counter_keys.key if committed.key is not None else None, + team=counter_keys.team if committed.team is not None else None, + user=counter_keys.user if committed.user is not None else None, + org=counter_keys.org if committed.org is not None else None, + end_user=counter_keys.end_user if committed.end_user is not None else None, + ) + + +async def read_admission_committed_spend( + spend_counter_store: SpendCounterStore, counter_keys: EntitySpendCounterKeys +) -> CommittedBudgetSpend: + """Rides the admission MGET: a counter the batch cannot answer stays unknown here and the emit-time read has it.""" + if not spend_counter_store.has_redis: + return await read_committed_budget_spend(spend_counter_store, counter_keys) + twin_keys: Final = committed_spend_keys(counter_keys.tracked()) + batched: Final = await asyncio.gather(*(read_batched_spend_counter(twin_key) for twin_key in twin_keys)) + return _committed_budget_spend( + counter_keys=counter_keys, + committed={ + twin_key: read[0] for twin_key, read in zip(twin_keys, batched) if read is not None and read[0] is not None + }, + ) + + +async def drop_committed_spend(spend_counter_store: SpendCounterStore, counter_keys: Sequence[str]) -> None: + """After a direct write to the stored spend the twin holds the old total, so the next charge seeds it again.""" + twin_keys: Final = committed_spend_keys(counter_keys) + if not twin_keys: + return + try: + await spend_counter_store.delete(twin_keys) + except Exception as e: # noqa: BLE001 # the write already committed; a kept twin lasts until charges stop for its TTL + verbose_proxy_logger.warning("committed spend drop failed: %s", e) + + +@with_service_target(SPEND_COUNTERS_TARGET) +async def seed_committed_spend( + spend_counter_store: SpendCounterStore, prisma_client: PrismaClient | None, counter_key: str +) -> bool: + """One DB read per worker for a cold twin: charges waiting on the lock find the twin the first one wrote.""" + twin_key: Final = committed_spend_key(counter_key) + async with await SpendCounterReseed.counter_lock(twin_key): + try: + stored, _ = await spend_counter_store.read(twin_key) + except Exception as e: # noqa: BLE001 # an unseeded twin only skips this charge; the next one seeds it again + verbose_proxy_logger.debug("committed spend seed read failed for %s: %s", twin_key, e) + return False + return stored is not None or await _seed_committed_spend_from_db( + spend_counter_store, prisma_client, counter_key + ) + + +async def _seed_committed_spend_from_db( + spend_counter_store: SpendCounterStore, prisma_client: PrismaClient | None, counter_key: str +) -> bool: + db_spend: Final = await ( + SpendCounterReseed.end_user_from_db(prisma_client=prisma_client, counter_key=counter_key) + if counter_key.startswith(END_USER_COUNTER_PREFIX) + else SpendCounterReseed.from_db(prisma_client=prisma_client, counter_key=counter_key) + ) + if db_spend is None: + return False + twin_key: Final = committed_spend_key(counter_key) + try: + await spend_counter_store.seed(twin_key, db_spend, increment_if_lost=False) + except Exception as e: # noqa: BLE001 # an unseeded twin only skips this charge; the next one seeds it again + verbose_proxy_logger.debug("committed spend seed failed for %s: %s", twin_key, e) + return False + return True diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 0348d8e1f82..f3dc672a072 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -2,7 +2,7 @@ import collections import json import os -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass from datetime import date, datetime, timedelta, timezone from functools import partial @@ -53,6 +53,7 @@ from litellm.proxy.auth.authorization_dependencies import ( from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.openai_error_payload import error_status_code, openai_error_param, openai_error_type from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup +from litellm.proxy.spend_tracking.committed_spend import committed_spend_tracked, drop_committed_spend from litellm.proxy.spend_tracking.session_page import ( AGENT_CALL_TYPE_SQL, MCP_CALL_TYPES_SQL, @@ -74,6 +75,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( get_spend_by_team_and_customer, ) from litellm.proxy.utils import handle_exception_on_proxy, query_raw_within_statement_timeout +from litellm.redis.stores.spend_counters import key_counter, team_counter from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import SpendLogsRepository from litellm.repositories.team_repository import TeamRepository @@ -3595,6 +3597,50 @@ async def view_spend_logs( ) +class _EntityIdRow(TypedDict): + id: ReadOnly[str] + + +_COMMITTED_DROP_PAGE_SIZE: Final = 1000 +_CommittedDropEntity: TypeAlias = Literal["key", "team"] +_COMMITTED_DROP_PAGE_QUERIES: Final[Mapping[_CommittedDropEntity, str]] = MappingProxyType( + { + "key": 'SELECT token AS id FROM "LiteLLM_VerificationToken" WHERE token > $1 ORDER BY token LIMIT $2', + "team": 'SELECT team_id AS id FROM "LiteLLM_TeamTable" WHERE team_id > $1 ORDER BY team_id LIMIT $2', + } +) +_COMMITTED_DROP_COUNTERS: Final[Mapping[_CommittedDropEntity, Callable[[str], str]]] = MappingProxyType( + {"key": key_counter, "team": team_counter} +) + + +async def _drop_committed_spend_page( + prisma_client: PrismaClient, entity: _CommittedDropEntity, after_id: str +) -> str | None: + from litellm.proxy.proxy_server import spend_counter_store + + rows: Final[Sequence[_EntityIdRow]] = await _query_raw( + prisma_client, _COMMITTED_DROP_PAGE_QUERIES[entity], after_id, _COMMITTED_DROP_PAGE_SIZE + ) + counter_of: Final = _COMMITTED_DROP_COUNTERS[entity] + await drop_committed_spend(spend_counter_store, tuple(counter_of(row["id"]) for row in rows)) + return rows[-1]["id"] if len(rows) == _COMMITTED_DROP_PAGE_SIZE else None + + +async def _drop_committed_spend_pages(prisma_client: PrismaClient, entity: _CommittedDropEntity) -> None: + after_id: str | None = "" # rebind-ok: the keyset cursor moves one page at a time + while after_id is not None: + after_id = await _drop_committed_spend_page(prisma_client, entity, after_id) + + +async def _drop_all_key_and_team_committed_spend(prisma_client: PrismaClient) -> None: + try: + await _drop_committed_spend_pages(prisma_client, "key") + await _drop_committed_spend_pages(prisma_client, "team") + except Exception as e: # noqa: BLE001 # the reset already committed; a kept twin lasts until charges stop for its TTL + verbose_proxy_logger.warning("committed spend drop after a global reset failed: %s", e) + + @router.post( "/global/spend/reset", tags=["Budget & Spend Tracking"], @@ -3626,6 +3672,8 @@ async def global_spend_reset(): _ = await transaction.litellm_teamtable.update_many(data={"spend": 0.0}, where={}) _ = await transaction.query_raw(_LOCK_ALL_KEYS_SQL) _ = await transaction.litellm_verificationtoken.update_many(data={"spend": 0.0}, where={}) + if committed_spend_tracked(): + await _drop_all_key_and_team_committed_spend(prisma_client) return { "message": "Spend for all API Keys and Teams reset successfully", diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 503dfb69793..97c327cbe4e 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -188,6 +188,7 @@ from litellm.litellm_core_utils.served_output_texts import ( served_stream_output_texts, ) from litellm.litellm_core_utils.token_counter import offload_token_count, usage_tokenizer_model +from litellm.litellm_core_utils.total_deadline import TotalTimeoutExpired from litellm.llms import load_guardrail_translation_mappings from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field, stream_item_items from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @@ -1737,8 +1738,17 @@ def _client_disconnected(exception: Exception) -> bool: return isinstance(exception, HTTPException) and exception.status_code == LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED +def _ran_out_of_total_timeout(exception: Exception) -> bool: + """The request's ``total_timeout`` ran out while the provider was working on it, like a disconnect.""" + return any(isinstance(cause, TotalTimeoutExpired) for cause in _exception_chain(exception)) + + def _provider_billed_the_failure(exception: Exception) -> bool: - return _timed_out_reading_from_provider(exception) or _client_disconnected(exception) + return ( + _timed_out_reading_from_provider(exception) + or _client_disconnected(exception) + or _ran_out_of_total_timeout(exception) + ) def _as_str_mapping(value: object) -> Mapping[str, object]: @@ -9253,13 +9263,14 @@ async def monitor_spend_logs_queue( proxy_logging_obj: Proxy logging object """ from litellm.constants import ( + SPEND_LOG_QUEUE_MAX_BACKOFF_SECONDS, SPEND_LOG_QUEUE_POLL_INTERVAL, SPEND_LOG_QUEUE_SIZE_THRESHOLD, ) threshold: Final = SPEND_LOG_QUEUE_SIZE_THRESHOLD base_interval: Final = SPEND_LOG_QUEUE_POLL_INTERVAL - max_backoff: Final = 30.0 # Maximum backoff interval in seconds + max_backoff: Final = SPEND_LOG_QUEUE_MAX_BACKOFF_SECONDS backoff_multiplier: Final = 1.5 # Exponential backoff multiplier current_interval = base_interval flush_requested: Final = asyncio.Event() diff --git a/litellm/proxy/vector_store_endpoints/utils.py b/litellm/proxy/vector_store_endpoints/utils.py index 8ce48fc3516..4dec4c7b38b 100644 --- a/litellm/proxy/vector_store_endpoints/utils.py +++ b/litellm/proxy/vector_store_endpoints/utils.py @@ -25,7 +25,7 @@ from litellm.utils import ProviderConfigManager _JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True)) -def _normalize_litellm_params( +def normalize_litellm_params( vector_store: LiteLLM_ManagedVectorStore, ) -> LiteLLM_ManagedVectorStore: litellm_params: Final = vector_store.get("litellm_params") @@ -46,6 +46,10 @@ def _is_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> bool: ) +def _key_metadata(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, object]: + return user_api_key_dict.metadata + + def assert_proxy_admin_for_vector_store_index_management( user_api_key_dict: UserAPIKeyAuth, *, @@ -278,7 +282,7 @@ async def get_litellm_managed_vector_store( vector_store_id=vector_store_id ) if vector_store is not None: - return _normalize_litellm_params(vector_store) + return normalize_litellm_params(vector_store) except Exception as e: verbose_proxy_logger.warning( "Failed to resolve vector store id=%s from registry: %s", @@ -310,7 +314,7 @@ async def get_litellm_managed_vector_store( ) if not rows: return None - return _normalize_litellm_params(LiteLLM_ManagedVectorStore(**rows[0].model_dump())) + return normalize_litellm_params(LiteLLM_ManagedVectorStore(**rows[0].model_dump())) except Exception as e: verbose_proxy_logger.warning( "Failed to resolve vector store id=%s from shared cache: %s", @@ -413,6 +417,34 @@ def check_vector_store_permission( return False +def vector_store_endpoint_forbidden(index_name: str) -> HTTPException: + return HTTPException( + status_code=403, + detail=( + f"User does not have permission to call vector store endpoint {index_name}. " + "Ask your administrator to add the necessary permissions to your API key/Team." + ), + ) + + +def has_vector_store_index_access( + user_api_key_dict: UserAPIKeyAuth, index_name: str, permission: Literal["read", "write"] | None +) -> bool: + return _is_proxy_admin(user_api_key_dict) or ( + permission is not None + and check_vector_store_permission( + index_name, permission, _key_metadata(user_api_key_dict), user_api_key_dict.team_metadata + ) + ) + + +def assert_vector_store_index_access( + user_api_key_dict: UserAPIKeyAuth, index_name: str, permission: Literal["read", "write"] | None +) -> None: + if not has_vector_store_index_access(user_api_key_dict, index_name, permission): + raise vector_store_endpoint_forbidden(index_name) + + def is_allowed_to_call_vector_store_endpoint( provider: LlmProviders, index_name: str, @@ -432,7 +464,7 @@ def is_allowed_to_call_vector_store_endpoint( ): return True # check what allowed permissions are for the key - key_metadata: Final = user_api_key_dict.metadata + key_metadata: Final = _key_metadata(user_api_key_dict) team_metadata: Final = user_api_key_dict.team_metadata provider_config: Final = ProviderConfigManager.get_provider_vector_stores_config(provider=provider) @@ -478,14 +510,7 @@ def is_allowed_to_call_vector_store_endpoint( break if permission_type is None: - raise HTTPException( - status_code=403, - detail=( - f"User does not have permission to call vector store endpoint " - f"{index_name}. Ask your administrator to add the necessary " - "permissions to your API key/Team." - ), - ) + raise vector_store_endpoint_forbidden(index_name) # Check if key has specific permission for allowed_vector_store_indexes has_permission: Final = check_vector_store_permission( @@ -496,10 +521,7 @@ def is_allowed_to_call_vector_store_endpoint( ) if not has_permission: - raise HTTPException( - status_code=403, - detail=f"User does not have permission to call vector store endpoint {index_name}. Ask your administrator to add the necessary permissions to your API key/Team.", - ) + raise vector_store_endpoint_forbidden(index_name) return has_permission @@ -516,7 +538,7 @@ def is_allowed_to_call_vector_store_files_endpoint( ): return True - key_metadata: Final = user_api_key_dict.metadata + key_metadata: Final = _key_metadata(user_api_key_dict) team_metadata: Final = user_api_key_dict.team_metadata provider_config: Final = ProviderConfigManager.get_provider_vector_store_files_config(provider=provider) diff --git a/litellm/proxy/video_endpoints/endpoints.py b/litellm/proxy/video_endpoints/endpoints.py index 31d6af71594..4e14eb31762 100644 --- a/litellm/proxy/video_endpoints/endpoints.py +++ b/litellm/proxy/video_endpoints/endpoints.py @@ -22,13 +22,13 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( from litellm.proxy.image_endpoints.endpoints import batch_to_bytesio from litellm.proxy.video_endpoints.utils import ( VideoRequestRoute, + authorized_video_model, encode_character_id_in_response, encode_video_id_in_response, extract_model_from_target_model_names, fallback_requested_model, get_custom_provider_from_data, routable_model_name, - video_model_group, video_reference_to_id, ) from litellm.types.videos.utils import ( @@ -283,7 +283,7 @@ async def video_status( if custom_llm_provider: data["custom_llm_provider"] = custom_llm_provider - resolved_model: Final = video_model_group(decoded, llm_router) + resolved_model: Final = await authorized_video_model(decoded, llm_router, user_api_key_dict) if resolved_model: data["model"] = resolved_model @@ -384,7 +384,7 @@ async def video_content( if custom_llm_provider: data["custom_llm_provider"] = custom_llm_provider - resolved_model: Final = video_model_group(decoded, llm_router) + resolved_model: Final = await authorized_video_model(decoded, llm_router, user_api_key_dict) if resolved_model: data["model"] = resolved_model # Process request using ProxyBaseLLMRequestProcessing @@ -490,7 +490,7 @@ async def video_remix( if custom_llm_provider: data["custom_llm_provider"] = custom_llm_provider - resolved_model: Final = video_model_group(decoded, llm_router) + resolved_model: Final = await authorized_video_model(decoded, llm_router, user_api_key_dict) if resolved_model: data["model"] = resolved_model @@ -814,7 +814,7 @@ async def video_edit( ) data["custom_llm_provider"] = custom_llm_provider - resolved_model: Final = video_model_group(decoded, llm_router) + resolved_model: Final = await authorized_video_model(decoded, llm_router, user_api_key_dict) if resolved_model: data["model"] = resolved_model @@ -917,7 +917,7 @@ async def video_extension( ) data["custom_llm_provider"] = custom_llm_provider - resolved_model: Final = video_model_group(decoded, llm_router) + resolved_model: Final = await authorized_video_model(decoded, llm_router, user_api_key_dict) if resolved_model: data["model"] = resolved_model diff --git a/litellm/proxy/video_endpoints/utils.py b/litellm/proxy/video_endpoints/utils.py index 696ad4ab170..24fd443b68a 100644 --- a/litellm/proxy/video_endpoints/utils.py +++ b/litellm/proxy/video_endpoints/utils.py @@ -3,6 +3,9 @@ from typing import TYPE_CHECKING, Any, Final import orjson from pydantic import BaseModel, ConfigDict, TypeAdapter +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.auth_utils import video_access_group +from litellm.proxy.openai_files_endpoints.common_utils import authorize_model_for_key from litellm.types.videos.main import DecodedVideoId, VideoObject from litellm.types.videos.utils import ( decode_video_id_with_provider, @@ -122,6 +125,20 @@ def video_model_group(decoded: DecodedVideoId, llm_router: "Router | None") -> s return creating_group +async def authorized_video_model( + decoded: DecodedVideoId, llm_router: "Router | None", user_api_key_dict: UserAPIKeyAuth +) -> str | None: + model: Final = video_model_group(decoded, llm_router) + if model is None: + return None + await authorize_model_for_key( + model_id=video_access_group(decoded, llm_router, resolved_model=model) or model, + llm_router=llm_router, + user_api_key_dict=user_api_key_dict, + ) + return model + + def fallback_requested_model(requested_model: str | None, response: object, llm_router: "Router | None") -> str | None: """The group a create asked for when one of that group's router fallbacks served it.""" from litellm.proxy.auth.auth_utils import router_fallback_model_names diff --git a/litellm/rag/ingestion/bedrock_ingestion.py b/litellm/rag/ingestion/bedrock_ingestion.py index 5520912dd41..1d1ec421292 100644 --- a/litellm/rag/ingestion/bedrock_ingestion.py +++ b/litellm/rag/ingestion/bedrock_ingestion.py @@ -14,11 +14,15 @@ from __future__ import annotations import asyncio import json import uuid -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, ClassVar, Final from litellm._logging import verbose_logger from litellm.litellm_core_utils.aws_partition import get_aws_arn_prefix from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.web_identity_session_policy import ( + BEDROCK_KNOWLEDGE_BASE_INGESTION_SESSION_POLICY, + WebIdentitySessionPolicy, +) from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion, RAGIngestBadRequestError if TYPE_CHECKING: @@ -85,6 +89,8 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): - aws_web_identity_token, aws_sts_endpoint, aws_external_id """ + _web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = BEDROCK_KNOWLEDGE_BASE_INGESTION_SESSION_POLICY + def __init__( self, ingest_options: RAGIngestOptions, diff --git a/litellm/rag/ingestion/s3_vectors_ingestion.py b/litellm/rag/ingestion/s3_vectors_ingestion.py index f31b068e709..a0f29d26d26 100644 --- a/litellm/rag/ingestion/s3_vectors_ingestion.py +++ b/litellm/rag/ingestion/s3_vectors_ingestion.py @@ -18,7 +18,7 @@ from __future__ import annotations import hashlib import uuid from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Final, TypedDict +from typing import TYPE_CHECKING, ClassVar, Final, TypedDict import litellm from litellm._logging import verbose_logger @@ -29,6 +29,10 @@ from litellm.constants import ( ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.web_identity_session_policy import ( + S3_VECTORS_INGESTION_SESSION_POLICY, + WebIdentitySessionPolicy, +) from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -85,6 +89,8 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): - non_filterable_metadata_keys: List of metadata keys to exclude from filtering """ + _web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = S3_VECTORS_INGESTION_SESSION_POLICY + def __init__( self, ingest_options: RAGIngestOptions, diff --git a/litellm/redis/stores/spend_counters/__init__.py b/litellm/redis/stores/spend_counters/__init__.py index 3ad06589f53..6a6dc85586e 100644 --- a/litellm/redis/stores/spend_counters/__init__.py +++ b/litellm/redis/stores/spend_counters/__init__.py @@ -6,15 +6,17 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Final from litellm._internal_context import with_service_target -from litellm.caching.redis_batch import Batch +from litellm.caching.redis_batch import Batch, PostCallSettlement from litellm.redis.stores._route import post_call_request, ready, route from litellm.redis.stores.spend_counters.handle import SpendCountersHandle from litellm.redis.stores.spend_counters.types import ( + COMMITTED_COUNTER_PREFIX, END_USER_COUNTER_PREFIX, SPEND_COUNTER_PREFIX, SPEND_COUNTERS_TARGET, SpendIncrement, SpendReading, + committed_counter, end_user_counter, is_unavailable, key_counter, @@ -34,12 +36,14 @@ if TYPE_CHECKING: from litellm.redis.stores.connection import RedisStores __all__ = [ + "COMMITTED_COUNTER_PREFIX", "END_USER_COUNTER_PREFIX", "SPEND_COUNTERS_TARGET", "SPEND_COUNTER_PREFIX", "SpendCounterStore", "SpendIncrement", "SpendReading", + "committed_counter", "end_user_counter", "is_unavailable", "key_counter", @@ -81,6 +85,11 @@ class SpendCounterStore: def has_redis(self) -> bool: return self._handle.has_redis + def set_committed_ttl(self, ttl: float) -> None: + """The least TTL a Redis write gives a committed twin, which keeps it past the DB write that holds its + charges.""" + self._handle.set_committed_ttl(ttl) + @with_service_target(SPEND_COUNTERS_TARGET) async def read(self, counter: str) -> SpendReading: reading: Final = await self._handle.read(counter) @@ -104,12 +113,13 @@ class SpendCounterStore: return tuple(await self._handle.increment_many(increments)) @with_service_target(SPEND_COUNTERS_TARGET) - def increment_many_post_call(self, increments: Sequence[SpendIncrement]) -> bool: - """Declares the increments on the request's post-call batch, False when there is no Redis or request. - Each total lands in memory once the batch goes out; a counter whose increment failed is dropped everywhere.""" + def increment_many_post_call(self, increments: Sequence[SpendIncrement]) -> PostCallSettlement | None: + """Declares the increments on the request's post-call batch, None when there is no Redis or request. + Each total lands in memory once the batch goes out; a counter whose increment failed is dropped everywhere. + The settlement reports when every reply is back, whoever sent the batch.""" request: Final = post_call_request() if self.has_redis and increments else None if request is None: - return False + return None return self._handle.increment_many_post_call(request, increments) @with_service_target(SPEND_COUNTERS_TARGET) diff --git a/litellm/redis/stores/spend_counters/handle.py b/litellm/redis/stores/spend_counters/handle.py index e4bd789c341..5ee29b7ad8e 100644 --- a/litellm/redis/stores/spend_counters/handle.py +++ b/litellm/redis/stores/spend_counters/handle.py @@ -3,7 +3,7 @@ from __future__ import annotations from collections.abc import Awaitable, Mapping, Sequence from typing import Protocol -from litellm.caching.redis_batch import NativePending +from litellm.caching.redis_batch import NativePending, PostCallSettlement from litellm.redis.stores.spend_counters.types import SpendIncrement from litellm.rust_bridge._native import NativeRedisClient, NativeRequestRedis, NativeSend, SpendRead @@ -13,6 +13,7 @@ class SpendCountersHandle(Protocol): def has_redis(self) -> bool: ... @property def redis_ttl_seconds(self) -> float | None: ... + def set_committed_ttl(self, ttl_seconds: float, /) -> None: ... def attach(self, client: NativeRedisClient, /) -> bool: ... def read(self, counter: str) -> Awaitable[SpendRead]: ... def read_many( @@ -20,7 +21,9 @@ class SpendCountersHandle(Protocol): ) -> NativePending[Sequence[float | None] | None] | None: ... def increment(self, counter: str, amount: float) -> Awaitable[float]: ... def increment_many(self, increments: Sequence[SpendIncrement]) -> Awaitable[Sequence[float]]: ... - def increment_many_post_call(self, request: NativeRequestRedis, increments: Sequence[SpendIncrement]) -> bool: ... + def increment_many_post_call( + self, request: NativeRequestRedis, increments: Sequence[SpendIncrement] + ) -> PostCallSettlement | None: ... def set_max(self, counter: str, spend: float) -> Awaitable[float | None]: ... def seed(self, counter: str, spend: float, *, increment_if_lost: bool) -> Awaitable[float]: ... def set(self, counter: str, spend: float, ttl_seconds: float | None = None) -> Awaitable[None]: ... diff --git a/litellm/redis/stores/spend_counters/types.py b/litellm/redis/stores/spend_counters/types.py index 7dc846dabaa..9be029b6953 100644 --- a/litellm/redis/stores/spend_counters/types.py +++ b/litellm/redis/stores/spend_counters/types.py @@ -7,6 +7,7 @@ from litellm.rust_bridge._native import NativeRedisStoreError SPEND_COUNTERS_TARGET: Final = "spend_counters" SPEND_COUNTER_PREFIX: Final = "spend:" END_USER_COUNTER_PREFIX: Final = "spend:end_user:" +COMMITTED_COUNTER_PREFIX: Final = "spend:committed:" SpendIncrement: TypeAlias = tuple[str, float] SpendReading: TypeAlias = tuple[float | None, bool] @@ -61,6 +62,11 @@ def shadow_eval_counter(job_id: str) -> str: return f"spend:shadow_eval:{job_id}" +def committed_counter(counter_key: str) -> str: + """The charges-only twin of a key, team, user, org or end user lifetime counter.""" + return f"{COMMITTED_COUNTER_PREFIX}{counter_key.removeprefix(SPEND_COUNTER_PREFIX)}" + + def is_unavailable(error: BaseException) -> bool: """Redis timed out or its breaker is open, the failures a post-call increment drops.""" return isinstance(error, NativeRedisStoreError) and error.args[:1] in (("timeout",), ("breaker_open",)) diff --git a/litellm/repositories/accounting_delivery_repository.py b/litellm/repositories/accounting_delivery_repository.py index d5bbe891089..624c2b8851a 100644 --- a/litellm/repositories/accounting_delivery_repository.py +++ b/litellm/repositories/accounting_delivery_repository.py @@ -60,12 +60,17 @@ class DeliveryBatchResult: error: str | None = None -class DeliveryMetrics(BaseModel): +class DeliveryBacklog(BaseModel): model_config = ConfigDict(frozen=True) pending_count: int parked_count: int oldest_pending_seconds: float + + +class DeliveryStorage(BaseModel): + model_config = ConfigDict(frozen=True) + batch_count: int table_bytes: int @@ -88,6 +93,14 @@ SET attempts = attempts + 1, last_error = $3, state = $4::"AccountingDeliveryState" WHERE delivery_id = $1 AND state = 'pending' AND attempts = $2::integer RETURNING delivery_id, attempts""" +_BACKLOG: Final = """SELECT COUNT(*) FILTER (WHERE state = 'pending') AS pending_count, + COUNT(*) FILTER (WHERE state = 'parked') AS parked_count, + COALESCE(EXTRACT(EPOCH FROM CURRENT_TIMESTAMP - MIN(created_at) FILTER (WHERE state = 'pending')), 0) + AS oldest_pending_seconds +FROM "LiteLLM_SpendLogDelivery" +WHERE state IN ('pending', 'parked')""" +_STORAGE: Final = """SELECT (SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery" WHERE is_batch) AS batch_count, + pg_total_relation_size('"LiteLLM_SpendLogDelivery"') AS table_bytes""" _UNDECODABLE: Final = """UPDATE "LiteLLM_SpendLogDelivery" SET state = 'parked', attempts = attempts + 1, last_error = 'UndecodablePayload' WHERE delivery_id = ANY($1::text[]) AND state = 'pending'""" @@ -114,17 +127,13 @@ class AccountingDeliveryRepository: def __init__(self, database: LedgerDatabase) -> None: self._database = database - async def metrics(self) -> DeliveryMetrics: + async def backlog(self) -> DeliveryBacklog: async with self._database.tx(timeout=_TIMEOUT) as transaction: - rows: Final = await transaction.query_raw( - "SELECT COUNT(*) FILTER (WHERE state = 'pending') AS pending_count, " - "COUNT(*) FILTER (WHERE state = 'parked') AS parked_count, " - "COALESCE(EXTRACT(EPOCH FROM CURRENT_TIMESTAMP - MIN(created_at) FILTER (WHERE state = 'pending')), 0) " - "AS oldest_pending_seconds, COUNT(*) FILTER (WHERE is_batch) AS batch_count, " - "pg_total_relation_size('\"LiteLLM_SpendLogDelivery\"') AS table_bytes " - 'FROM "LiteLLM_SpendLogDelivery"' - ) - return DeliveryMetrics.model_validate(rows[0]) + return DeliveryBacklog.model_validate((await transaction.query_raw(_BACKLOG))[0]) + + async def storage(self) -> DeliveryStorage: + async with self._database.tx(timeout=_TIMEOUT) as transaction: + return DeliveryStorage.model_validate((await transaction.query_raw(_STORAGE))[0]) async def cleanup_applied(self, *, retention_days: int = 1, batch_size: int = 1000) -> int: if retention_days < 1 or batch_size < 1: diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 9da89d96e73..900cfa3d597 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -647,12 +647,21 @@ class ResponsesSessionHandler: SELECT 1 FROM matching_session WHERE matching_session.session_id = logs.session_id - AND matching_session."user" IS NOT DISTINCT FROM logs."user" - AND matching_session.team_id IS NOT DISTINCT FROM logs.team_id + AND ( + matching_session."user" = logs."user" + OR (matching_session."user" IS NULL AND logs."user" IS NULL) + ) + AND ( + matching_session.team_id = logs.team_id + OR (matching_session.team_id IS NULL AND logs.team_id IS NULL) + ) AND ( COALESCE(matching_session."user", '') <> '' OR COALESCE(matching_session.team_id, '') <> '' - OR matching_session.api_key = logs.api_key + OR ( + COALESCE(matching_session.api_key, '') <> '' + AND matching_session.api_key = logs.api_key + ) ) ) ORDER BY "endTime" ASC; diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 4273647567e..08b2c518632 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -486,9 +486,18 @@ class BaseResponsesAPIStreamingIterator: if _is_json_object(parsed_chunk): if self.responses_api_provider_config is None: raise ValueError("responses_api_provider_config is required to process live streaming chunks") + typed_chunk: Final = ( + parsed_chunk + if parsed_chunk.get("type") is not None or "error" not in parsed_chunk + else { + "sequence_number": SEQUENCE_NUMBER_STAMPED_ON_EMISSION, + **parsed_chunk, + "type": ResponsesAPIStreamEvents.ERROR.value, + } + ) openai_responses_api_chunk: Final = self.responses_api_provider_config.transform_streaming_response( model=self.model, - parsed_chunk=parsed_chunk, + parsed_chunk=typed_chunk, logging_obj=self.logging_obj, ) diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index eef454862f9..55155a96fa3 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -47,6 +47,12 @@ from litellm.litellm_core_utils.add_retry_fallback_headers import ( prepare_response_for_header_attachment, ) from litellm.litellm_core_utils.agentic_followup_failure import is_agentic_followup_failure +from litellm.litellm_core_utils.attempt_settlement import ( + ATTEMPT_SETTLEMENT_GRACE_SECONDS, + RoutedAttempt, + released, + routed_attempt, +) from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type from litellm.litellm_core_utils.core_helpers import ( as_str_mapping, @@ -63,6 +69,11 @@ from litellm.litellm_core_utils.routing_common_utils import ( ) from litellm.litellm_core_utils.routing_groups import validate_routing_strategy from litellm.litellm_core_utils.secret_redaction import redact_string +from litellm.litellm_core_utils.total_deadline import ( + TotalTimeoutNotApplicable, + current_total_deadline, + enforce_total_timeout_scope, +) from litellm.redis.stores.connection import connect from litellm.redis.stores.router_cooldowns import ROUTER_COOLDOWNS_TARGET, RouterCooldownStore from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET, RouterUsage @@ -209,7 +220,14 @@ from litellm.router_rust.surface.settings import ( request_timeout, settings_view, ) -from litellm.router_strategy.least_busy import CALL_IDENTITY, call_id_of, owning_picks, pick_reserves, picks_reserve +from litellm.router_strategy.least_busy import ( + CALL_IDENTITY, + call_id_of, + owning_picks, + pick_reserves, + picks_reserve, + release_started_attempt, +) from litellm.router_utils.cooldown_handlers import get_cooldown_deployment_info, mark_deployment_selection_failure from litellm.router_utils.endpoint_eligibility import endpoint_for_call, selected_endpoint from litellm.router_utils.kubernetes_pod_discovery import ( @@ -218,7 +236,7 @@ from litellm.router_utils.kubernetes_pod_discovery import ( resolve_pods_after_bound, ) from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import EncryptedContentAffinityCheck -from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import async_io_token_refund_failure +from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import release_io_token_reservation from litellm.router_utils.router_callbacks.track_deployment_metrics import find_deployment_metadata from litellm.router_utils.routing_read_batch import RoutingReadBatch from litellm.types.router import ( @@ -2902,6 +2920,9 @@ class RustCore( case "attempt": outcome = await driver.attempt(step) case "sleep": + if (outlasted := _error_outlasting_total_timeout(driver, _sleep_seconds(step))) is not None: + driver.close_route_phase() + raise outlasted await self._host.sleep(_sleep_seconds(step)) outcome = None case "allow_fallback": @@ -3860,6 +3881,15 @@ def _origin_level(origin: Mapping[str, object]) -> int: return level if isinstance(level, int) else 0 +def _error_outlasting_total_timeout(driver: "_Driver", seconds: float) -> BaseException | None: + """The error a retry would have answered when its backoff outlasts the request's total_timeout: the + call ends with it now, rather than with a timeout after waiting out the rest of the budget.""" + deadline: Final = current_total_deadline() + if deadline is None or deadline.allows_wait(seconds): + return None + return driver.last_error + + def _sleep_seconds(step: Mapping[str, object]) -> float: """A retry's sleep: ``_time_to_sleep_before_retry`` returns a plain ``0`` when another deployment is healthy, else a backoff in float seconds.""" @@ -5150,6 +5180,28 @@ class _Driver: async def attempt(self, step: Mapping[str, object]) -> Mapping[str, object]: if self._passthrough is not None and not step.get("deployment_id"): return await self._passthrough_attempt(step) + with routed_attempt() as attempt: + return await self._routed_attempt(step, attempt) + + def _release( + self, call: AttemptCall, attempt: RoutedAttempt, *, refund_io_tokens: bool + ) -> "tuple[asyncio.Task[None], ...]": + """What a failed or cancelled attempt still holds, whatever its logging object (shared across + attempts, it logs only the first failure). Tokens a provider already spent are kept counted.""" + core: Final = self._core + releases: Final = ( + release_started_attempt(core._selectors.reservations, core._cache, attempt.attempt_id), # pyright: ignore[reportPrivateUsage] # the driver serves its own core + release_io_token_reservation( + core._cache, # pyright: ignore[reportPrivateUsage] # see above + call.kwargs, # pyright: ignore[reportArgumentType] # the call's kwargs dict + parent_otel_span=get_parent_otel_span_from_kwargs(call.kwargs), + ) + if refund_io_tokens + else None, + ) + return tuple(release for release in releases if release is not None) + + async def _routed_attempt(self, step: Mapping[str, object], attempt: RoutedAttempt) -> Mapping[str, object]: step = await self._resolve_pods(step) # rebind-ok: the picked deployment, resolved to a pod deployment: Mapping[str, object] = as_str_mapping(step.get("deployment")) or {} # rebind-ok: the picked one call: AttemptCall | None = None # rebind-ok: built inside the try @@ -5163,6 +5215,7 @@ class _Driver: return {**self._failed(step, self._default_copy(step), None, refused_call), "recorded": True} try: deployment, call = self._call(step) # rebind-ok: rebound as the Python router rebinds it + enforce_total_timeout_scope(call.kwargs) if call.refused is not None: await log_refused_call(call, call.refused) raise call.refused @@ -5215,16 +5268,18 @@ class _Driver: # `self.success_calls[model_name] += 1`, inside the try: a counter replaced with a # plain dict fails the attempt. self._core._endpoints.counters.succeeded(self._plan, call) # pyright: ignore[reportPrivateUsage] # see above + except asyncio.CancelledError: + if call is not None: + await released( + self._release(call, attempt, refund_io_tokens=not attempt.answered), + ATTEMPT_SETTLEMENT_GRACE_SECONDS, + ) + raise + except TotalTimeoutNotApplicable: + raise except Exception as error: # noqa: BLE001 # every attempt failure is classified and handed to the Rust call loop if call is not None: - # `make_call`: a failed call hands back the io tokens it reserved, whatever its - # logging object (shared across attempts, it logs only the first failure). - with contextlib.suppress(Exception): - await async_io_token_refund_failure( - self._core._cache, # pyright: ignore[reportPrivateUsage] # the driver serves its own core - call.kwargs, # pyright: ignore[reportArgumentType] # the call's kwargs dict - parent_otel_span=get_parent_otel_span_from_kwargs(call.kwargs), - ) + await released(self._release(call, attempt, refund_io_tokens=True)) failed: Final = self._failed(step, deployment, call, error) if reading_stream is None or not logs_failure: return { diff --git a/litellm/router_rust/steps/callbacks.py b/litellm/router_rust/steps/callbacks.py index cd2d69d13a8..48cb998df9b 100644 --- a/litellm/router_rust/steps/callbacks.py +++ b/litellm/router_rust/steps/callbacks.py @@ -452,6 +452,10 @@ class OptionalChecks: *self._policies_not_set_globally(), ) + def has_filters(self) -> bool: + """Whether ``filter_deployments`` would run any callback: without one it keeps every deployment.""" + return any(filters_deployments(callback) for callback in self.routing_callbacks()) + async def filter_deployments( self, model: str, @@ -464,12 +468,10 @@ class OptionalChecks: the Rust core takes only if its later filters keep it; a filter's error is raised for the Rust core to decide on.""" remaining: list[object] = list(deployments) # mutable-ok: filter contract # rebind-ok: narrowed - base: Final[object] = CustomLogger.async_filter_deployments # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # untyped deployment dicts request_kwargs: Final = dict(request) callbacks: Final = tuple(self.routing_callbacks()) for callback in callbacks: - # The base filter returns its input unchanged: only overriding filters are run. - if not isinstance(callback, CustomLogger) or type(callback).async_filter_deployments is base: # pyright: ignore[reportUnknownMemberType] # see above + if not isinstance(callback, CustomLogger) or not filters_deployments(callback): continue if isinstance(callback, PromptCachingDeploymentCheck): continue @@ -572,14 +574,22 @@ def _ids_of(value: object) -> Sequence[object]: return value if isinstance(value, list | tuple) else () # pyright: ignore[reportUnknownVariableType] # the step's id list +def filters_deployments(callback: object) -> bool: + """A routing callback that overrides ``async_filter_deployments``: the base filter returns its + input unchanged, so only overriding filters are run.""" + base: Final[object] = CustomLogger.async_filter_deployments # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # untyped deployment dicts + return isinstance(callback, CustomLogger) and type(callback).async_filter_deployments is not base # pyright: ignore[reportUnknownMemberType] # see above + + def filtered_ids(remaining: Iterable[object], candidates: Sequence[str]) -> list[str]: # mutable-ok: a JSON list """The ids of the deployments the filters left, in their order, among the candidates.""" - ids: Final[list[str]] = [] # mutable-ok: built here + allowed: Final = frozenset(candidates) + ids: Final[dict[str, None]] = {} # mutable-ok: built here, insertion-ordered for deployment in remaining: found = (as_str_mapping((as_str_mapping(deployment) or {}).get("model_info")) or {}).get("id") - if isinstance(found, str) and found in candidates and found not in ids: - ids.append(found) - return ids + if isinstance(found, str) and found in allowed: + ids.setdefault(found) + return list(ids) class RustOptionalChecks: @@ -698,6 +708,8 @@ class RustOptionalChecks: from litellm.router_rust.core import attempt_failure candidates: Final = [str(id) for id in step.get("candidates") or ()] # pyright: ignore[reportGeneralTypeIssues, reportUnknownVariableType, reportUnknownArgumentType] # the step's id list + if not self._opt.has_filters(): + return (list(dict.fromkeys(candidates)), None), None records: Final = self._records_of(candidates) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownVariableType] # RustCore's deployment dicts narrowed: Final = _request_with_failed_ids(request, step) try: diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index dfcb4cccf54..9d3767268ca 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -1,3 +1,4 @@ +import asyncio import contextlib import logging import threading @@ -15,14 +16,16 @@ from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import log_redis_failure +from litellm.constants import LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.attempt_settlement import current_attempt_id, release_in_background, released from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET from litellm.types.utils import CallTypes IN_FLIGHT_COUNT_TTL_SECONDS: Final = 60 * 60 SELECTION_RESERVATION_TTL_SECONDS: Final = 30.0 -SELECTION_RESERVATION_MAX_KEYS: Final = 10_000 +SELECTION_RESERVATION_MAX_KEYS: Final = 20_000 _RESERVATIONS_LOCK: Final = threading.Lock() _PICK_RESERVES: Final[ContextVar[bool]] = ContextVar("least_busy_pick_reserves") _PICK_OWNER: Final[ContextVar[str | None]] = ContextVar("least_busy_pick_owner", default=None) @@ -76,7 +79,9 @@ _CALL_KWARGS: Final = TypeAdapter(_CallKwargs) CALL_IDENTITY: Final = TypeAdapter(CallIdentity) _DEPLOYMENTS: Final = TypeAdapter(list[_Deployment]) _MEMORY_COUNTS: Final[TypeAdapter[tuple[float | None, ...] | None]] = TypeAdapter(tuple[float | None, ...] | None) -_ATTEMPT_TIMES: Final[TypeAdapter[tuple[float, ...] | None]] = TypeAdapter(tuple[float, ...] | None) +_Start: TypeAlias = tuple[float, str | None] +_STARTS: Final[TypeAdapter[tuple[_Start, ...] | None]] = TypeAdapter(tuple[_Start, ...] | None) +_STARTED_ATTEMPT: Final[TypeAdapter[tuple[str, str, str] | None]] = TypeAdapter(tuple[str, str, str] | None) _Pick: TypeAlias = tuple[float, str | None] _PICKS: Final[TypeAdapter[tuple[_Pick, ...] | None]] = TypeAdapter(tuple[_Pick, ...] | None) _PICKED_DEPLOYMENTS: Final[TypeAdapter[tuple[str, ...] | None]] = TypeAdapter(tuple[str, ...] | None) @@ -100,7 +105,11 @@ def _routed_deployment_id(metadata: _Metadata | None) -> str | int | None: def _deployment_ref(kwargs: Mapping[str, object]) -> tuple[str, str] | None: call_type: Final = kwargs.get("call_type") - if is_batch_retrieve_call_type(call_type) or call_type == CallTypes.arealtime.value: + if ( + is_batch_retrieve_call_type(call_type) + or call_type == CallTypes.arealtime.value + or kwargs.get(LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL) + ): return None try: call: Final = _CALL_KWARGS.validate_python(kwargs) @@ -133,6 +142,10 @@ def _started_attempts_key(call_id: str, ref: tuple[str, str]) -> str: return f"least_busy_started:{call_id}:{_request_count_key(*ref)}" +def _started_attempt_key(attempt_id: str) -> str: + return f"least_busy_attempt:{attempt_id}" + + def _picked_deployments_key(call_id: str) -> str: return f"least_busy_picked:{call_id}" @@ -211,8 +224,8 @@ def _picks(store: _ReservationStore, deployment_id: str) -> tuple[_Pick, ...]: return _PICKS.validate_python(store.get_cache(_reservation_key(deployment_id))) or () -def _attempt_times(store: _ReservationStore, key: str) -> tuple[float, ...]: - return _ATTEMPT_TIMES.validate_python(store.get_cache(key)) or () +def _starts(store: _ReservationStore, key: str) -> tuple[_Start, ...]: + return _STARTS.validate_python(store.get_cache(key)) or () def _picked_deployments(store: _ReservationStore, call_id: str) -> tuple[str, ...]: @@ -231,6 +244,60 @@ def _release_call_on(store: _ReservationStore, deployment_id: str, call_id: str) _put(store, _reservation_key(deployment_id), kept, SELECTION_RESERVATION_TTL_SECONDS) +def _end_started_attempt(store: _ReservationStore, attempt_id: str, *, every: bool) -> tuple[str, str, int] | None: + """Under ``_RESERVATIONS_LOCK``: end one of the attempt's own starts (a call retried inside one + attempt starts more than once), or ``every`` one still open. Returns the deployment and how many ended.""" + started: Final = _STARTED_ATTEMPT.validate_python(store.get_cache(_started_attempt_key(attempt_id))) + if started is None: + return None + call_id, model_group, deployment_id = started + key: Final = _started_attempts_key(call_id, (model_group, deployment_id)) + starts: Final = _starts(store, key) + own: Final = tuple(index for index, start in enumerate(starts) if start[1] == attempt_id) + ending: Final = frozenset(own if every else own[:1]) + _put( + store, + key, + tuple(start for index, start in enumerate(starts) if index not in ending), + IN_FLIGHT_COUNT_TTL_SECONDS, + ) + if len(ending) == len(own): + store.delete_cache(_started_attempt_key(attempt_id)) + return (model_group, deployment_id, len(ending)) if ending else None + + +def _count_locally(cache: DualCache, key: str, delta: int) -> None: + try: + local: Final = cache.increment_cache(key, delta, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) + if local < 0: + cache.set_cache(key, 0, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) + except TypeError as e: + _warn_unwritable(key, e) + + +@with_service_target(ROUTER_USAGE_TARGET) +async def _async_count_shared(cache: DualCache, key: str, delta: int) -> None: + redis_cache: Final = cache.redis_cache + if redis_cache is None: + return + try: + await redis_cache.async_increment_with_floor(key, delta, IN_FLIGHT_COUNT_TTL_SECONDS) + except Exception as e: + _warn_unwritable(key, e) + + +def release_started_attempt(store: _ReservationStore, cache: DualCache, attempt_id: str) -> "asyncio.Task[None] | None": + """End a routed attempt that never reached its success or failure callback (refused, failed + without logging, or cancelled): its in-flight count drops now, the shared one in the background.""" + with _RESERVATIONS_LOCK: + ended: Final = _end_started_attempt(store, attempt_id, every=True) + if ended is None: + return None + model_group, deployment_id, count = ended + _count_locally(cache, _request_count_key(model_group, deployment_id), -count) + return release_in_background(_async_count_shared(cache, _request_count_key(model_group, deployment_id), -count)) + + def release_unstarted_picks(store: _ReservationStore, call_id: str | None) -> None: if call_id is None: return @@ -312,29 +379,47 @@ class LeastBusyLoggingHandler(CustomLogger): return healthy_deployments[index] def _start_attempt(self, ref: tuple[str, str], call_id: str | None) -> None: + attempt_id: Final = current_attempt_id() with _RESERVATIONS_LOCK: now: Final = self._clock() self._take_reservation(ref[1], call_id, now) - self._count_locally(_request_count_key(*ref), 1) + _count_locally(self.router_cache, _request_count_key(*ref), 1) if call_id is not None: key: Final = _started_attempts_key(call_id, ref) - started: Final = (*_attempt_times(self._reservations, key), now) + started: Final = (*_starts(self._reservations, key), (now, attempt_id)) _put(self._reservations, key, started, IN_FLIGHT_COUNT_TTL_SECONDS) + if attempt_id is not None: + _put( + self._reservations, + _started_attempt_key(attempt_id), + (call_id, *ref), + IN_FLIGHT_COUNT_TTL_SECONDS, + ) - def _end_attempt(self, ref: tuple[str, str], call_id: str | None) -> bool: + def _end_attempt(self, ref: tuple[str, str], call_id: str | None) -> tuple[str, str] | None: if call_id is None: with _RESERVATIONS_LOCK: - self._count_locally(_request_count_key(*ref), -1) - return True + _count_locally(self.router_cache, _request_count_key(*ref), -1) + return ref + attempt_id: Final = current_attempt_id() + if attempt_id is not None: + with _RESERVATIONS_LOCK: + ended: Final = _end_started_attempt(self._reservations, attempt_id, every=False) + if ended is not None: + _count_locally(self.router_cache, _request_count_key(ended[0], ended[1]), -1) + return None if ended is None else (ended[0], ended[1]) key: Final = _started_attempts_key(call_id, ref) with _RESERVATIONS_LOCK: - started: Final = _attempt_times(self._reservations, key) + started: Final = _starts(self._reservations, key) if started: + oldest_attempt: Final = started[0][1] + if oldest_attempt is not None: + self._reservations.delete_cache(_started_attempt_key(oldest_attempt)) _put(self._reservations, key, started[1:], IN_FLIGHT_COUNT_TTL_SECONDS) - self._count_locally(_request_count_key(*ref), -1) + _count_locally(self.router_cache, _request_count_key(*ref), -1) else: self._take_reservation(ref[1], call_id, self._clock()) - return bool(started) + return ref if started else None @with_service_target(ROUTER_USAGE_TARGET) def log_pre_api_call(self, model: str, messages: object, kwargs: Mapping[str, object]) -> None: @@ -418,15 +503,17 @@ class LeastBusyLoggingHandler(CustomLogger): def _settle(self, kwargs: Mapping[str, object]) -> None: ref: Final = _deployment_ref(kwargs) - if ref is None or not self._end_attempt(ref, call_id_of(kwargs)): - return - self._count_shared(_request_count_key(*ref), -1) + ended: Final = None if ref is None else self._end_attempt(ref, call_id_of(kwargs)) + if ended is not None: + self._count_shared(_request_count_key(*ended), -1) async def _async_settle(self, kwargs: Mapping[str, object]) -> None: ref: Final = _deployment_ref(kwargs) - if ref is None or not self._end_attempt(ref, call_id_of(kwargs)): - return - await self._async_count_shared(_request_count_key(*ref), -1) + ended: Final = None if ref is None else self._end_attempt(ref, call_id_of(kwargs)) + if ended is not None: + await released( + (release_in_background(_async_count_shared(self.router_cache, _request_count_key(*ended), -1)),) + ) def _this_workers_counts(self, keys: tuple[str, ...]) -> tuple[int, ...]: return _local_counts(self.router_cache.batch_get_cache(list(keys), local_only=True), keys) @@ -434,16 +521,6 @@ class LeastBusyLoggingHandler(CustomLogger): def _at_least_this_workers(self, keys: tuple[str, ...], shared: tuple[int, ...]) -> tuple[int, ...]: return tuple(max(counts) for counts in zip(shared, self._this_workers_counts(keys), strict=True)) - def _count_locally(self, key: str, delta: int) -> None: - try: - local: Final = self.router_cache.increment_cache( - key, delta, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS - ) - if local < 0: - self.router_cache.set_cache(key, 0, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) - except TypeError as e: - _warn_unwritable(key, e) - def _count_shared(self, key: str, delta: int) -> None: redis_cache: Final = self.router_cache.redis_cache if redis_cache is None: @@ -452,12 +529,3 @@ class LeastBusyLoggingHandler(CustomLogger): redis_cache.increment_with_floor(key, delta, IN_FLIGHT_COUNT_TTL_SECONDS) except Exception as e: _warn_unwritable(key, e) - - async def _async_count_shared(self, key: str, delta: int) -> None: - redis_cache: Final = self.router_cache.redis_cache - if redis_cache is None: - return - try: - await redis_cache.async_increment_with_floor(key, delta, IN_FLIGHT_COUNT_TTL_SECONDS) - except Exception as e: - _warn_unwritable(key, e) diff --git a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py index 846cc9f3d61..4b618aba476 100644 --- a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py @@ -11,9 +11,13 @@ Used by ModelRateLimitingCheck when a deployment sets itpm/otpm. from __future__ import annotations +import asyncio import contextlib import contextvars -from collections.abc import Mapping, MutableMapping +import threading +import time +from collections import OrderedDict +from collections.abc import Callable, Mapping, MutableMapping from dataclasses import dataclass from datetime import datetime from types import MappingProxyType @@ -26,11 +30,15 @@ import litellm from litellm import token_counter from litellm._internal_context import with_service_target from litellm._logging import verbose_router_logger +from litellm._uuid import uuid from litellm.caching.dual_cache import DualCache from litellm.caching.redis_cache import RedisCache, RedisCircuitBreakerOpenError from litellm.constants import NO_MAX_TOKENS_TPM_RESERVATION_FRACTION +from litellm.litellm_core_utils.attempt_settlement import current_attempt_id, release_in_background, released from litellm.litellm_core_utils.core_helpers import coerce_token_limit from litellm.litellm_core_utils.internal_call_metadata import ( + IO_TOKEN_ATTEMPT_KEY, + IO_TOKEN_RESERVATION_ID_KEY, IO_TOKEN_RESERVATION_KEYS, ITPM_CACHE_KEY, ITPM_RESERVED_KEY, @@ -55,6 +63,31 @@ else: RoutingArgsTTL: Final = 60 +class _SettledReservations: + """The ids of reservations settled within the counters' window, so a copy of a stash that outlives + the dict it was claimed from (the logging object copies litellm_metadata) is never settled twice. + Ids leave only once their window has passed, never to make room.""" + + def __init__(self, window: float, clock: Callable[[], float] = time.monotonic) -> None: + self._window: Final = window + self._clock: Final = clock + self._expiries: Final[OrderedDict[str, float]] = OrderedDict() # mutable-ok: pruned oldest first as ids expire + self._lock: Final = threading.Lock() + + def claim(self, reservation_id: str) -> bool: + with self._lock: + now: Final = self._clock() + while self._expiries and next(iter(self._expiries.values())) <= now: + _ = self._expiries.popitem(last=False) + if reservation_id in self._expiries: + return False + self._expiries[reservation_id] = now + self._window + return True + + +_SETTLED_RESERVATIONS: Final = _SettledReservations(window=RoutingArgsTTL * 2) + + @dataclass(frozen=True, slots=True) class _IoTokenRateLimitRequest: kwargs: dict[str, object] | None @@ -355,9 +388,10 @@ def _resolve_reconcile_usage_tokens( def _reservation( itpm_reserved: int, otpm_reserved: int, itpm_cache_key: str | None, otpm_cache_key: str | None -) -> Mapping[str, object]: +) -> Mapping[str, int | str | None]: return MappingProxyType( { + IO_TOKEN_RESERVATION_ID_KEY: uuid.uuid4().hex, ITPM_RESERVED_KEY: itpm_reserved, OTPM_RESERVED_KEY: otpm_reserved, ITPM_CACHE_KEY: itpm_cache_key, @@ -370,6 +404,7 @@ def _stash_reservation_in_metadata(request: _IoTokenRateLimitRequest | None, res metadata: Final = None if request is None or not request.kwargs else request.kwargs.get(request.metadata_channel) if isinstance(metadata, dict): metadata.update(reservation) + metadata[IO_TOKEN_ATTEMPT_KEY] = current_attempt_id() def _key_for_this_minute(reserved_key: str) -> str: @@ -387,12 +422,16 @@ def _reconcile_increments( ) -def _extract_reservation(reservation: Mapping[str, int | str | None]) -> tuple[int, int, str | None, str | None]: +def _reserved_tokens(value: object) -> int: + return value if isinstance(value, int) and not isinstance(value, bool) else 0 + + +def _extract_reservation(reservation: Mapping[str, object]) -> tuple[int, int, str | None, str | None]: itpm_cache_key: Final = reservation.get(ITPM_CACHE_KEY) otpm_cache_key: Final = reservation.get(OTPM_CACHE_KEY) return ( - int(reservation.get(ITPM_RESERVED_KEY, 0) or 0), - int(reservation.get(OTPM_RESERVED_KEY, 0) or 0), + _reserved_tokens(reservation.get(ITPM_RESERVED_KEY)), + _reserved_tokens(reservation.get(OTPM_RESERVED_KEY)), itpm_cache_key if isinstance(itpm_cache_key, str) else None, otpm_cache_key if isinstance(otpm_cache_key, str) else None, ) @@ -421,11 +460,31 @@ def _reservation_channels(kwargs: Mapping[str, object] | None) -> tuple[object, return top_level + from_params + from_logging_object -def _read_reservation_from_kwargs(kwargs: Mapping[str, object] | None) -> tuple[int, int, str | None, str | None]: - for channel_dict in _reservation_channels(kwargs): - if isinstance(channel_dict, dict) and ITPM_RESERVED_KEY in channel_dict: - return _extract_reservation(channel_dict) - return 0, 0, None, None +def _stashed_reservation(kwargs: Mapping[str, object] | None) -> Mapping[str, object] | None: + channels: Final = (_as_mutable_mapping(channel) for channel in _reservation_channels(kwargs)) + return next((channel for channel in channels if channel is not None and ITPM_RESERVED_KEY in channel), None) + + +def _claim_reservation_from_kwargs( + kwargs: Mapping[str, object] | None, *, any_attempt: bool = False +) -> tuple[int, int, str | None, str | None]: + """ + Take the stashed reservation so nothing else settles it. Under a routed attempt only that attempt's + own reservation is taken, unless ``any_attempt``: a late callback of an earlier attempt sharing these + kwargs leaves the next attempt's reservation alone. + """ + stashed: Final = _stashed_reservation(kwargs) + if stashed is None: + return 0, 0, None, None + attempt: Final = current_attempt_id() + if not any_attempt and attempt is not None and stashed.get(IO_TOKEN_ATTEMPT_KEY) not in (None, attempt): + return 0, 0, None, None + reservation: Final = _extract_reservation(stashed) + reservation_id: Final = stashed.get(IO_TOKEN_RESERVATION_ID_KEY) + _clear_reservation_from_kwargs(kwargs) + if isinstance(reservation_id, str) and not _SETTLED_RESERVATIONS.claim(reservation_id): + return 0, 0, None, None + return reservation def _clear_reservation_from_kwargs(kwargs: Mapping[str, object] | None) -> None: @@ -626,11 +685,50 @@ async def async_io_token_pre_call_check( return deployment itpm_key, otpm_key = cache_keys - itpm_reserved = 0 - otpm_reserved = 0 + # The reservation runs in a task of its own: a request cancelled mid-reservation never knows how much + # it took, so the task finishes and then hands back whatever it reserved. + reserving: Final = asyncio.get_running_loop().create_task( + _async_reserve( + dual_cache, + itpm_key, + otpm_key, + _reservation_value(estimated_input, itpm_limit) if itpm_limit is not None else 0, + otpm_to_reserve, + itpm_limit, + otpm_limit, + parent_otel_span, + ) + ) + try: + reservation: Final = await asyncio.shield(reserving) + except asyncio.CancelledError: + _ = release_in_background(_refund_once_reserved(dual_cache, reserving, parent_otel_span)) + raise + _stash_reservation_in_metadata(request, reservation) + return deployment + +async def _refund_once_reserved( + dual_cache: DualCache, reserving: asyncio.Task[Mapping[str, int | str | None]], parent_otel_span: Span | None +) -> None: + try: + reservation: Final = await reserving + except Exception: # noqa: BLE001 # a refused or failed reservation holds nothing + return + await _async_increment_all(dual_cache, _refunds(_extract_reservation(reservation)), parent_otel_span) + + +async def _async_reserve( + dual_cache: DualCache, + itpm_key: str, + otpm_key: str, + itpm_reserved: int, + otpm_reserved: int, + itpm_limit: int | None, + otpm_limit: int | None, + parent_otel_span: Span | None, +) -> Mapping[str, int | str | None]: if itpm_limit is not None: - itpm_reserved = _reservation_value(estimated_input, itpm_limit) await _increment_with_rollback( dual_cache, itpm_key, @@ -641,7 +739,6 @@ async def async_io_token_pre_call_check( ) if otpm_limit is not None: - otpm_reserved = otpm_to_reserve try: await _increment_with_rollback( dual_cache, @@ -664,16 +761,12 @@ async def async_io_token_pre_call_check( ) raise - _stash_reservation_in_metadata( - request, - _reservation( - itpm_reserved, - otpm_reserved, - itpm_key if itpm_limit is not None else None, - otpm_key if otpm_limit is not None else None, - ), + return _reservation( + itpm_reserved, + otpm_reserved if otpm_limit is not None else 0, + itpm_key if itpm_limit is not None else None, + otpm_key if otpm_limit is not None else None, ) - return deployment async def _async_redis_usage(redis_cache: RedisCache | None, key: str, parent_otel_span: Span | None) -> int | None: @@ -732,25 +825,22 @@ def io_token_reconcile_success( ) -> None: request_kwargs: Final[Mapping[str, object] | None] = kwargs response: Final[object] = response_obj - reservation: Final = _read_reservation_from_kwargs(request_kwargs) + reservation: Final = _claim_reservation_from_kwargs(request_kwargs) itpm_reserved, otpm_reserved, itpm_key, otpm_key = reservation if itpm_key is None and otpm_key is None: return billable_input, completion_tokens, usage_resolved = _resolve_reconcile_usage_tokens(request_kwargs, response) - try: - if usage_resolved: - for key, delta in _reconcile_increments(reservation, billable_input, completion_tokens): - dual_cache.increment_cache(key=key, value=delta, ttl=RoutingArgsTTL) - else: - verbose_router_logger.debug( - "[IO TOKEN LIMIT] usage missing; keeping reservation (itpm_reserved=%s, otpm_reserved=%s)", - itpm_reserved, - otpm_reserved, - ) - finally: - _clear_reservation_from_kwargs(request_kwargs) + if usage_resolved: + for key, delta in _reconcile_increments(reservation, billable_input, completion_tokens): + dual_cache.increment_cache(key=key, value=delta, ttl=RoutingArgsTTL) + else: + verbose_router_logger.debug( + "[IO TOKEN LIMIT] usage missing; keeping reservation (itpm_reserved=%s, otpm_reserved=%s)", + itpm_reserved, + otpm_reserved, + ) verbose_router_logger.debug( "[IO TOKEN LIMIT] reconciled (usage_resolved=%s, itpm_reserved=%s, billable_input=%s, otpm_reserved=%s, output=%s)", @@ -772,7 +862,7 @@ async def async_io_token_reconcile_success( ) -> None: request_kwargs: Final[Mapping[str, object] | None] = kwargs response: Final[object] = response_obj - reservation: Final = _read_reservation_from_kwargs(request_kwargs) + reservation: Final = _claim_reservation_from_kwargs(request_kwargs) itpm_reserved, otpm_reserved, itpm_key, otpm_key = reservation if itpm_key is None and otpm_key is None: return @@ -781,22 +871,25 @@ async def async_io_token_reconcile_success( # A refund goes back to the minute that held the reservation, while usage above the reservation is charged # to the minute the response finished in, so a call that runs past its minute still counts against a live - # window. Always clear the stash afterwards (even if an increment throws) so a retry or a duplicate success - # event can't re-process it. - try: - if usage_resolved: - for key, delta in _reconcile_increments(reservation, billable_input, completion_tokens): - await dual_cache.async_increment_cache( - key=key, value=delta, ttl=RoutingArgsTTL, parent_otel_span=parent_otel_span - ) - else: - verbose_router_logger.debug( - "[IO TOKEN LIMIT] usage missing; keeping reservation (itpm_reserved=%s, otpm_reserved=%s)", - itpm_reserved, - otpm_reserved, + # window. + if usage_resolved: + await released( + ( + release_in_background( + _async_increment_all( + dual_cache, + _reconcile_increments(reservation, billable_input, completion_tokens), + parent_otel_span, + ) + ), ) - finally: - _clear_reservation_from_kwargs(request_kwargs) + ) + else: + verbose_router_logger.debug( + "[IO TOKEN LIMIT] usage missing; keeping reservation (itpm_reserved=%s, otpm_reserved=%s)", + itpm_reserved, + otpm_reserved, + ) verbose_router_logger.debug( "[IO TOKEN LIMIT] reconciled (usage_resolved=%s, itpm_reserved=%s, billable_input=%s, otpm_reserved=%s, output=%s)", @@ -812,24 +905,14 @@ async def async_io_token_reconcile_success( def io_token_refund_failure( dual_cache: DualCache, kwargs: Mapping[str, object] | None, + *, + any_attempt: bool = False, ) -> None: request_kwargs: Final[Mapping[str, object] | None] = kwargs - itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(request_kwargs) - if itpm_key is None and otpm_key is None: - return - if itpm_key is not None and itpm_reserved > 0: - dual_cache.increment_cache( - key=itpm_key, - value=-itpm_reserved, - ttl=RoutingArgsTTL, - ) - if otpm_key is not None and otpm_reserved > 0: - dual_cache.increment_cache( - key=otpm_key, - value=-otpm_reserved, - ttl=RoutingArgsTTL, - ) - _clear_reservation_from_kwargs(request_kwargs) + reservation: Final = _claim_reservation_from_kwargs(request_kwargs, any_attempt=any_attempt) + itpm_reserved, otpm_reserved, _, _ = reservation + for key, delta in _refunds(reservation): + dual_cache.increment_cache(key=key, value=delta, ttl=RoutingArgsTTL) verbose_router_logger.debug("[IO TOKEN LIMIT] refunded ITPM=%s OTPM=%s", itpm_reserved, otpm_reserved) @@ -857,7 +940,7 @@ def refund_stale_reservation_before_retry(dual_cache: DualCache, kwargs: Mapping """ if not kwargs: return - io_token_refund_failure(dual_cache, kwargs) + io_token_refund_failure(dual_cache, kwargs, any_attempt=True) @with_service_target(ROUTER_USAGE_TARGET) @@ -867,26 +950,38 @@ async def async_io_token_refund_failure( *, parent_otel_span: Span | None = None, ) -> None: - request_kwargs: Final[Mapping[str, object] | None] = kwargs - itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(request_kwargs) - if itpm_key is None and otpm_key is None: - return - if itpm_key is not None and itpm_reserved > 0: + refund: Final = release_io_token_reservation(dual_cache, kwargs, parent_otel_span=parent_otel_span) + if refund is not None: + await released((refund,)) + + +def release_io_token_reservation( + dual_cache: DualCache, + kwargs: Mapping[str, object] | None, + *, + parent_otel_span: Span | None = None, +) -> asyncio.Task[None] | None: + reservation: Final = _claim_reservation_from_kwargs(kwargs) + refunds: Final = _refunds(reservation) + if not refunds: + return None + verbose_router_logger.debug("[IO TOKEN LIMIT] refunding ITPM=%s OTPM=%s", reservation[0], reservation[1]) + return release_in_background(_async_increment_all(dual_cache, refunds, parent_otel_span)) + + +def _refunds(reservation: tuple[int, int, str | None, str | None]) -> tuple[tuple[str, int], ...]: + itpm_reserved, otpm_reserved, itpm_key, otpm_key = reservation + held: Final = ((itpm_key, itpm_reserved), (otpm_key, otpm_reserved)) + return tuple((key, -reserved) for key, reserved in held if key is not None and reserved > 0) + + +async def _async_increment_all( + dual_cache: DualCache, increments: tuple[tuple[str, int], ...], parent_otel_span: Span | None +) -> None: + for key, delta in increments: await dual_cache.async_increment_cache( - key=itpm_key, - value=-itpm_reserved, - ttl=RoutingArgsTTL, - parent_otel_span=parent_otel_span, + key=key, value=delta, ttl=RoutingArgsTTL, parent_otel_span=parent_otel_span ) - if otpm_key is not None and otpm_reserved > 0: - await dual_cache.async_increment_cache( - key=otpm_key, - value=-otpm_reserved, - ttl=RoutingArgsTTL, - parent_otel_span=parent_otel_span, - ) - _clear_reservation_from_kwargs(request_kwargs) - verbose_router_logger.debug("[IO TOKEN LIMIT] refunded ITPM=%s OTPM=%s", itpm_reserved, otpm_reserved) def build_io_token_rate_limit_headers( diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 77af87f5eb3..114d2df026e 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -358,6 +358,11 @@ class SpendRead: @property def from_redis(self) -> bool: ... +@final +class PostCallSettlement: + @property + def settled(self) -> bool: ... + @final class NativeSpendCounters: def __new__(cls, client: NativeRedisClient | None) -> NativeSpendCounters: ... @@ -366,6 +371,7 @@ class NativeSpendCounters: def has_redis(self) -> bool: ... @property def redis_ttl_seconds(self) -> float | None: ... + def set_committed_ttl(self, ttl_seconds: float) -> None: ... def read(self, counter: str) -> Future[SpendRead]: ... def read_many( self, counters: Sequence[str], send: NativeSend @@ -374,7 +380,7 @@ class NativeSpendCounters: def increment_many(self, increments: Sequence[tuple[str, float]]) -> Future[list[float]]: ... def increment_many_post_call( self, request: NativeRequestRedis, increments: Sequence[tuple[str, float]] - ) -> bool: ... + ) -> PostCallSettlement | None: ... def set_max(self, counter: str, spend: float) -> Future[float | None]: ... def seed(self, counter: str, spend: float, *, increment_if_lost: bool) -> Future[float]: ... def set(self, counter: str, spend: float, ttl_seconds: float | None = None) -> Future[None]: ... @@ -596,6 +602,7 @@ __all__ = [ "NativeSend", "NativeSensitiveRoutePins", "NativeSpendCounters", + "PostCallSettlement", "ProcessReservedForForking", "ResponsesWebSocketConnection", "RouterCall", diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index 9adfcb60d74..78040ba7225 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -95,7 +95,7 @@ class Feature(str, Enum): CACHE_CONTROL_INJECTION = "cache_control_injection" """Prompt caching injection that Python applies to messages and system.""" STREAM_COST_INJECTION = "stream_cost_injection" - """litellm.include_cost_in_streaming_usage, where Python adds cost to streamed usage.""" + """litellm.include_cost_in_streaming_usage or the x-litellm-include-cost-in-streaming-usage header, where Python adds cost to streamed usage.""" THINKING_SIGNATURE_RETRY = "thinking_signature_retry" """Replayed thinking blocks, which Python strips and retries once if Anthropic rejects their signature.""" TOOL_CUSTOM_FIELD = "tool_custom_field" diff --git a/litellm/rust_bridge/host/model_capabilities.py b/litellm/rust_bridge/host/model_capabilities.py index cbf59dfb97b..efb8e10022a 100644 --- a/litellm/rust_bridge/host/model_capabilities.py +++ b/litellm/rust_bridge/host/model_capabilities.py @@ -1,6 +1,27 @@ from __future__ import annotations +from collections.abc import Mapping +from functools import lru_cache +from typing import Final, TypeAlias + import litellm +from litellm.litellm_core_utils.core_helpers import as_str_mapping + + +class _Held: + __slots__ = ("value",) + + def __init__(self, value: object) -> None: + self.value = value + + def __hash__(self) -> int: + return id(self.value) + + def __eq__(self, other: object) -> bool: + return isinstance(other, _Held) and other.value is self.value + + +_Fingerprint: TypeAlias = tuple[str, str | None, _Held, int, int, tuple[_Held, ...]] def _resolved_provider(model: str, custom_llm_provider: str | None) -> tuple[str, str]: @@ -11,11 +32,50 @@ def _resolved_provider(model: str, custom_llm_provider: str | None) -> tuple[str return resolved_model, provider +def _fingerprint(model: str, custom_llm_provider: str | None, resolved_model: str, provider: str) -> _Fingerprint: + """Changes whenever an answer could: ``register_model`` bumps the generation, and replacing + any entry the probes look up (``litellm.model_cost[key] = ...``) changes that entry's identity. + The key holds the entries it names, so a cached answer's entries are never freed and their + addresses never reused.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + from litellm.utils import get_model_cost_mutation_generation + + model_cost: Final = as_str_mapping(litellm.model_cost) or {} # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # the model map is an untyped dict of rows + keys: Final = dict.fromkeys( + ( + *AnthropicModelInfo._model_map_lookup_candidates(model), # pyright: ignore[reportPrivateUsage] # the keys the capability walk reads + resolved_model, + f"{provider}/{resolved_model}", + ) + ) + return ( + model, + custom_llm_provider, + _Held(model_cost), + get_model_cost_mutation_generation(), + len(model_cost), + tuple(_Held(model_cost.get(key)) for key in keys), + ) + + def anthropic_model_capabilities(model: str, custom_llm_provider: str | None) -> dict[str, object]: + """The capability probes the Rust messages route reads, once per model-map state: each probe + walks the model map (``get_model_info`` plus generalization matching), and the route asks on + every request.""" + resolved_model, provider = _resolved_provider(model, custom_llm_provider) + cached: Final = _capabilities( + _fingerprint(model, custom_llm_provider, resolved_model, provider), resolved_model, provider + ) + tiers: Final = as_str_mapping(cached["effort_tiers"]) + return {**cached, "effort_tiers": dict(tiers) if tiers is not None else cached["effort_tiers"]} + + +@lru_cache(maxsize=256) +def _capabilities(fingerprint: _Fingerprint, resolved_model: str, provider: str) -> Mapping[str, object]: from litellm.llms.anthropic.chat.transformation import AnthropicConfig from litellm.llms.anthropic.common_utils import AnthropicModelInfo - resolved_model, provider = _resolved_provider(model, custom_llm_provider) + model: Final = fingerprint[0] def supports(flag: str) -> bool: return AnthropicModelInfo._supports_model_capability(model, flag, provider) # pyright: ignore[reportPrivateUsage] # same probes the Python transform runs; forking them would drift @@ -27,6 +87,7 @@ def anthropic_model_capabilities(model: str, custom_llm_provider: str | None) -> "supports_reasoning": supports("supports_reasoning"), "supports_adaptive_thinking": supports("supports_adaptive_thinking"), "thinking_always_on": supports("thinking_always_on"), + "thinking_on_by_default": AnthropicModelInfo.thinks_by_default(model, provider), "supports_legacy_thinking": supports("supports_legacy_thinking"), "supports_output_config": supports("supports_output_config"), "supports_sampling_params": AnthropicModelInfo._supports_sampling_params(resolved_model), # pyright: ignore[reportPrivateUsage] # same gate the handler applies diff --git a/litellm/rust_bridge/messages/detectors.py b/litellm/rust_bridge/messages/detectors.py index 6a0c867979b..99555bb4852 100644 --- a/litellm/rust_bridge/messages/detectors.py +++ b/litellm/rust_bridge/messages/detectors.py @@ -6,7 +6,11 @@ from typing import Final import litellm from litellm.anthropic_beta_headers_manager import is_foundry_api_base -from litellm.litellm_core_utils.core_helpers import drops_unsupported_params +from litellm.litellm_core_utils.core_helpers import ( + drops_unsupported_params, + proxy_request_headers, + stream_cost_injection_requested, +) from litellm.llms.opencode.common_utils import is_opencode_api_base from litellm.rust_bridge.detectors import api_base, overridden, present, streaming from litellm.rust_bridge.public_call import NativeCall, optional_mapping, optional_sequence, optional_str @@ -64,8 +68,8 @@ def strict_tools_retry(request: NativeCall) -> bool: def stream_cost_injection(request: NativeCall) -> bool: - """Python adds response cost to streaming usage when the global setting is truthy.""" - return streaming(request) and bool(litellm.include_cost_in_streaming_usage) + """Python adds response cost to streaming usage when the global setting or the request header asks for it.""" + return streaming(request) and stream_cost_injection_requested(proxy_request_headers(request.resolved)) def pre_request_hook(_request: NativeCall) -> bool: diff --git a/litellm/secret_managers/aws_secret_manager_v2.py b/litellm/secret_managers/aws_secret_manager_v2.py index de8566e6fdc..37aabefb2ed 100644 --- a/litellm/secret_managers/aws_secret_manager_v2.py +++ b/litellm/secret_managers/aws_secret_manager_v2.py @@ -18,7 +18,7 @@ import json import os from collections.abc import Mapping from types import MappingProxyType -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, ClassVar, Final import httpx from pydantic import ConfigDict, TypeAdapter, with_config @@ -29,6 +29,7 @@ from litellm._logging import verbose_logger from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.optional_imports import ensure_optional_import from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.web_identity_session_policy import SECRETS_MANAGER_SESSION_POLICY, WebIdentitySessionPolicy from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, get_httpx_client, @@ -53,6 +54,8 @@ _AWS_SECRET_MANAGER_RESPONSE: Final = TypeAdapter(dict[str, object], config=Conf class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): + _web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = SECRETS_MANAGER_SESSION_POLICY + def __init__( self, aws_region_name: str | None = None, diff --git a/litellm/types/proxy/auth/user_api_key_auth.py b/litellm/types/proxy/auth/user_api_key_auth.py index 1183c36eda7..e08459ea14d 100644 --- a/litellm/types/proxy/auth/user_api_key_auth.py +++ b/litellm/types/proxy/auth/user_api_key_auth.py @@ -16,6 +16,7 @@ from litellm.types.agents import AgentCaller, AgentResponse from litellm.types.proxy.agent_identity import ManagedAgentContext from litellm.types.proxy.auth.user_roles import LitellmUserRoles from litellm.types.proxy.carried_budget_state import ( + CommittedBudgetSpend, OrgBudgetSnapshot, TeamBudgetSnapshot, UserBudgetSnapshot, @@ -206,6 +207,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob team_budget_snapshot: TeamBudgetSnapshot | None = Field(default=None, exclude=True) user_budget_snapshot: UserBudgetSnapshot | None = Field(default=None, exclude=True) org_budget_snapshot: OrgBudgetSnapshot | None = Field(default=None, exclude=True) + committed_budget_spend: CommittedBudgetSpend | None = Field(default=None, exclude=True, repr=False) matched_model_access_groups: list[str] | None = Field(default=None, exclude=True) budget_throttle_pct: float | None = Field(default=None, exclude=True) user: object | None = None # Expanded user object when expand=user is used @@ -329,6 +331,21 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob ) +COMMITTED_BUDGET_SPEND_PARAM: Final = "committed_budget_spend" + + +def committed_budget_spend_of_row( + metadata: Mapping[str, object], litellm_params: Mapping[str, object] +) -> CommittedBudgetSpend | None: + auth: Final = metadata.get("user_api_key_auth") + if isinstance(auth, UserAPIKeyAuth) and auth.committed_budget_spend is not None: + return auth.committed_budget_spend + if "user_api_key_auth" in metadata: + return CommittedBudgetSpend() + carried: Final = litellm_params.get(COMMITTED_BUDGET_SPEND_PARAM) + return carried if isinstance(carried, CommittedBudgetSpend) else None + + def user_api_key_has_admin_view(user_api_key_dict: UserAPIKeyAuth) -> bool: """Return True if the caller's role grants unscoped read access to all tenant resources (managed files, batches, vector stores, spend rows, etc). diff --git a/litellm/types/proxy/carried_budget_state.py b/litellm/types/proxy/carried_budget_state.py index 1a0c55a38ff..83ef6a21188 100644 --- a/litellm/types/proxy/carried_budget_state.py +++ b/litellm/types/proxy/carried_budget_state.py @@ -9,10 +9,18 @@ reads them back with ``from_metadata``. ``None`` means this request never carrie from collections.abc import Mapping from datetime import datetime from types import MappingProxyType +from typing import Final from pydantic import ConfigDict, Field, ValidationError from typing_extensions import Self +from litellm.redis.stores.spend_counters.types import ( + end_user_counter, + key_counter, + org_counter, + team_counter, + user_counter, +) from litellm.types.llms.base import LiteLLMBaseModel @@ -25,7 +33,7 @@ class _BudgetSnapshot(LiteLLMBaseModel): @classmethod def from_metadata(cls, metadata: Mapping[str, object]) -> Self | None: try: - return cls.model_validate(metadata) + return cls.model_validate(metadata, by_name=False) except ValidationError: return None @@ -64,3 +72,51 @@ class OrgBudgetSnapshot(_BudgetSnapshot): max_budget: float | None = Field( validation_alias="user_api_key_org_max_budget", serialization_alias="user_api_key_org_max_budget" ) + + +class EntitySpendCounterKeys(LiteLLMBaseModel): + """Spend counter of each entity whose remaining budget a gauge reports. ``None`` for an entity with no budget.""" + + model_config = ConfigDict(frozen=True) + + key: str | None = None + team: str | None = None + user: str | None = None + org: str | None = None + end_user: str | None = None + + @classmethod + def of( + cls, + token: str | None, + team_id: str | None, + user_id: str | None, + org_id: str | None, + end_user_id: str | None, + ) -> Self: + return cls( + key=key_counter(token) if token else None, + team=team_counter(team_id) if team_id else None, + user=user_counter(user_id) if user_id else None, + org=org_counter(org_id) if org_id else None, + end_user=end_user_counter(end_user_id) if end_user_id else None, + ) + + def tracked(self) -> tuple[str, ...]: + counter_keys: Final = (self.key, self.team, self.user, self.org, self.end_user) + return tuple(counter_key for counter_key in counter_keys if counter_key is not None) + + +class CommittedBudgetSpend(LiteLLMBaseModel): + """Spend each entity was charged, budget holds excluded, read from the committed spend counters for the + remaining budget gauges only. ``None`` when the entity has no committed counter. ``counter_keys`` names the + budgeted entities whose committed counters this request charges and reads.""" + + model_config = ConfigDict(frozen=True) + + counter_keys: EntitySpendCounterKeys = Field(default=EntitySpendCounterKeys(), repr=False) + key: float | None = Field(default=None, repr=False) + team: float | None = Field(default=None, repr=False) + user: float | None = Field(default=None, repr=False) + org: float | None = Field(default=None, repr=False) + end_user: float | None = Field(default=None, repr=False) diff --git a/litellm/utils.py b/litellm/utils.py index b9bf872e853..66b5b274f97 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -86,9 +86,12 @@ from litellm.constants import ( from litellm.litellm_core_utils.agentic_followup_failure import is_agentic_followup_failure from litellm.litellm_core_utils.allowed_tools_tool_choice import ( CHAT_ALLOWED_TOOLS_LAYOUT, + HOSTED_TOOL_CHOICE_TYPES, parse_chat_allowed_tools_tool_choice, ) +from litellm.litellm_core_utils.attempt_settlement import mark_attempt_answered from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, bind_budget_reservation_to_callbacks, has_attributes, iter_items, @@ -2515,7 +2518,9 @@ def client(original_function): exception=deployment_error, call_type=call_type, ) - except BaseException: # noqa: S110, BLE001 # hook dispatch - including cancellation mid-await - must never replace the real deployment failure, so there is nothing to do with what it raises + except asyncio.CancelledError: + raise + except BaseException: # noqa: S110, BLE001 # a failing hook must never replace the real deployment failure, so there is nothing to do with what it raises pass raise end_time = datetime.datetime.now() @@ -2557,6 +2562,7 @@ def client(original_function): ) elif call_type in (CallTypes.arealtime.value, CallTypes.aresponses_websocket.value): return result + mark_attempt_answered() ### POST-CALL RULES ### post_call_processing( original_response=result, @@ -5375,18 +5381,30 @@ def pre_process_non_default_params( return non_default_params -def remove_sensitive_keys_from_dict(d: dict) -> dict: +def remove_sensitive_keys_from_dict(d: Mapping[str, object]) -> dict[str, object]: """ Remove sensitive keys from a dictionary """ - sensitive_key_phrases: Final = ["key", "secret", "access", "credential"] - remove_keys: Final = [] - for key in d: - if any(phrase in key.lower() for phrase in sensitive_key_phrases): - remove_keys.append(key) - for key in remove_keys: - d.pop(key) - return d + sensitive_key_phrases: Final = ("key", "secret", "access", "credential") + result: Final = dict[str, object]() + pending: Final = [(d, result)] + while pending: + source, target = pending.pop() + for key, value in source.items(): + if ( + any(phrase in key.lower() for phrase in sensitive_key_phrases) + or key.lower() == "token" + or key.lower().endswith("_token") + ): + continue + mapping = as_str_mapping(value) if isinstance(value, dict) else None + if mapping is not None: + nested = dict[str, object]() + target[key] = nested + pending.append((mapping, nested)) + else: + target[key] = value + return result _TOOL_LIST: Final = TypeAdapter(list[object]) @@ -9757,7 +9775,14 @@ def validate_chat_completion_tool_choice( ) # Standard OpenAI format: {"type": "function", "function": {...}} - if tool_choice_type != "allowed_tools" and (tool_choice_type is None or tool_choice.get("function") is None): + forces_hosted_tool: Final[bool] = ( + isinstance(tool_choice_type, str) and tool_choice_type in HOSTED_TOOL_CHOICE_TYPES + ) + if ( + tool_choice_type != "allowed_tools" + and not forces_hosted_tool + and (tool_choice_type is None or tool_choice.get("function") is None) + ): raise BadRequestError( message=f"Invalid tool choice, tool_choice={tool_choice}. Please ensure tool_choice follows the OpenAI spec", model=model, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index ac5e99c9313..79e4d92aa9d 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -1969,6 +1969,9 @@ "prompt_cache_min_tokens": 512, "source": "https://aws.amazon.com/bedrock/pricing/" }, + "bedrock/claude-fable-latest": { + "alias_of": "global.anthropic.claude-fable-5-1" + }, "us.anthropic.claude-fable-5": { "cache_creation_input_token_cost": 1.375e-05, "cache_creation_input_token_cost_above_1hr": 2.2e-05, @@ -2305,6 +2308,9 @@ "region": "us-east-1" } }, + "bedrock/claude-opus-latest": { + "alias_of": "global.anthropic.claude-opus-5-5" + }, "us.anthropic.claude-opus-5": { "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, @@ -4063,6 +4069,9 @@ "deprecation_date": "2027-12-05", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement" }, + "azure_ai/claude-fable-latest": { + "alias_of": "azure_ai/claude-fable-5-1" + }, "azure_ai/claude-mythos-5-1": { "input_cost_per_token": 1e-05, "output_cost_per_token": 5e-05, @@ -4165,6 +4174,9 @@ "prompt_cache_min_tokens": 512, "source": "https://management.azure.com/subscriptions/c873328e-b572-4770-8dff-aaeb6f1f0e79/providers/Microsoft.CognitiveServices/locations/eastus2/models?api-version=2024-10-01" }, + "azure_ai/claude-opus-latest": { + "alias_of": "azure_ai/claude-opus-5-5" + }, "azure_ai/claude-opus-4-8": { "deprecation_date": "2027-09-01", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement", @@ -4382,6 +4394,9 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure_ai/gpt-astra-latest": { + "alias_of": "azure_ai/gpt-6-astra" + }, "azure_ai/gpt-6-luna": { "cache_creation_input_token_cost": 1.25e-07, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-07, @@ -4430,6 +4445,9 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure_ai/gpt-luna-latest": { + "alias_of": "azure_ai/gpt-6-luna" + }, "azure_ai/gpt-6-sol": { "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, @@ -4526,6 +4544,9 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure_ai/gpt-sol-latest": { + "alias_of": "azure_ai/gpt-6.1-sol" + }, "azure_ai/gpt-6.1-terra": { "cache_creation_input_token_cost": 1.25e-06, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-06, @@ -4574,6 +4595,9 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure_ai/gpt-terra-latest": { + "alias_of": "azure_ai/gpt-6.1-terra" + }, "azure_ai/gpt-5.5": { "deprecation_date": "2027-10-26", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement", @@ -5036,6 +5060,9 @@ "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false }, + "azure_ai/gpt-mini-latest": { + "alias_of": "azure_ai/gpt-5.4-mini" + }, "azure_ai/gpt-5.4-mini-2026-03-17": { "deprecation_date": "2027-09-21", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement", @@ -5316,7 +5343,7 @@ "supports_tool_choice": true }, "azure/eu/gpt-4o-realtime-preview-2024-12-17": { - "cache_read_input_audio_token_cost": 2.5e-06, + "cache_read_input_audio_token_cost": 2.75e-06, "cache_read_input_token_cost": 2.75e-06, "input_cost_per_audio_token": 4.4e-05, "input_cost_per_token": 5.5e-06, @@ -5326,7 +5353,7 @@ "mode": "realtime", "output_cost_per_audio_token": 8.8e-05, "output_cost_per_token": 2.2e-05, - "source": "https://azure.microsoft.com/en-us/pricing/details/azure-openai/", + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_modalities": [ "text", "audio" @@ -6664,6 +6691,7 @@ "supports_tool_choice": true }, "azure/gpt-realtime": { + "cache_read_input_image_token_cost": 5e-07, "cache_creation_input_audio_token_cost": 4e-06, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, @@ -6699,9 +6727,11 @@ "supports_parallel_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure" + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/" }, "azure/gpt-realtime-2025-08-28": { + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "cache_read_input_image_token_cost": 5e-07, "cache_creation_input_audio_token_cost": 4e-06, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, @@ -6739,6 +6769,7 @@ "supports_tool_choice": true }, "azure/gpt-realtime-1.5": { + "cache_read_input_image_token_cost": 5e-07, "cache_creation_input_audio_token_cost": 4e-06, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, @@ -6775,9 +6806,11 @@ "supports_parallel_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true, - "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure" + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/" }, "azure/gpt-realtime-1.5-2026-02-23": { + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "cache_read_input_image_token_cost": 5e-07, "cache_creation_input_audio_token_cost": 4e-06, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, @@ -6898,6 +6931,8 @@ "supports_tool_choice": true }, "azure/gpt-realtime-mini": { + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", + "cache_read_input_image_token_cost": 8e-08, "cache_creation_input_audio_token_cost": 3e-07, "cache_read_input_audio_token_cost": 3e-07, "cache_read_input_token_cost": 6e-08, @@ -7006,6 +7041,7 @@ ] }, "azure/gpt-4o-realtime-preview-2024-12-17": { + "cache_read_input_audio_token_cost": 2.5e-06, "cache_read_input_token_cost": 2.5e-06, "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 5e-06, @@ -7015,7 +7051,7 @@ "mode": "realtime", "output_cost_per_audio_token": 8e-05, "output_cost_per_token": 2e-05, - "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure", + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_modalities": [ "text", "audio" @@ -9080,6 +9116,9 @@ "cache_read_input_token_cost_flex": 5e-07, "cache_creation_input_token_cost_flex": 6.25e-06 }, + "azure/gpt-astra-latest": { + "alias_of": "azure/gpt-6-astra" + }, "azure/gpt-6-astra-2026-09-03": { "cache_creation_input_token_cost": 1.25e-05, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, @@ -9184,6 +9223,9 @@ "output_cost_per_token_above_272k_tokens_flex": 3.75e-07, "output_cost_per_token_flex": 2.5e-07 }, + "azure/gpt-luna-latest": { + "alias_of": "azure/gpt-6-luna" + }, "azure/gpt-6-luna-2026-09-22": { "cache_creation_input_token_cost": 1.25e-07, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-07, @@ -9403,6 +9445,9 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "azure/gpt-terra-latest": { + "alias_of": "azure/gpt-6.1-terra" + }, "azure/gpt-6.1-sol": { "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, @@ -9455,6 +9500,9 @@ "cache_read_input_token_cost_above_272k_tokens_flex": 1e-07, "cache_creation_input_token_cost_above_272k_tokens_flex": 2.5e-06 }, + "azure/gpt-sol-latest": { + "alias_of": "azure/gpt-6.1-sol" + }, "azure/gpt-6.1-sol-2026-09-29": { "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, @@ -11146,6 +11194,9 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_xhigh_reasoning_effort": true }, + "azure/gpt-mini-latest": { + "alias_of": "azure/gpt-5.4-mini" + }, "azure/gpt-5.4-mini-2026-03-17": { "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_flex": 3.75e-08, @@ -12388,7 +12439,7 @@ "supports_tool_choice": true }, "azure/us/gpt-4o-realtime-preview-2024-12-17": { - "cache_read_input_audio_token_cost": 2.5e-06, + "cache_read_input_audio_token_cost": 2.75e-06, "cache_read_input_token_cost": 2.75e-06, "input_cost_per_audio_token": 4.4e-05, "input_cost_per_token": 5.5e-06, @@ -12398,7 +12449,7 @@ "mode": "realtime", "output_cost_per_audio_token": 8.8e-05, "output_cost_per_token": 2.2e-05, - "source": "https://azure.microsoft.com/en-us/pricing/details/azure-openai/", + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_modalities": [ "text", "audio" @@ -13212,6 +13263,9 @@ "supports_tool_choice": true, "supports_vision": true }, + "azure_ai/kimi-latest": { + "alias_of": "azure_ai/FW-Kimi-K3" + }, "azure_ai/FW-MiniMax-M2.5": { "deprecation_date": "2027-07-01", "cache_read_input_token_cost": 3.3e-08, @@ -13251,6 +13305,9 @@ "supports_tool_choice": true, "supports_vision": true }, + "azure_ai/minimax-latest": { + "alias_of": "azure_ai/FW-MiniMax-M3" + }, "azure_ai/FW-Nemotron-Lightning-3.5-30B-A3B": { "cache_read_input_token_cost": 1e-08, "input_cost_per_token": 6e-08, @@ -13656,6 +13713,9 @@ "cache_read_input_token_cost": 1.45e-07, "supports_prompt_caching": true }, + "azure_ai/deepseek-pro-latest": { + "alias_of": "azure_ai/deepseek-v4-pro" + }, "azure_ai/deepseek-v4-flash": { "deprecation_date": "2028-02-20", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/concepts/model-lifecycle-retirement", @@ -13674,6 +13734,9 @@ "cache_read_input_token_cost": 2.8e-08, "supports_prompt_caching": true }, + "azure_ai/deepseek-flash-latest": { + "alias_of": "azure_ai/deepseek-v4-flash" + }, "azure_ai/DeepSeek-V4-Flash-0731": { "cache_read_input_token_cost": 1.4e-08, "deprecation_date": "2026-12-03", @@ -16623,6 +16686,9 @@ "source": "https://platform.claude.com/docs/en/about-claude/pricing", "supports_web_search": true }, + "claude-sonnet-latest": { + "alias_of": "claude-sonnet-5-5" + }, "claude-sonnet-4-6": { "supports_web_search": true, "supports_anthropic_compaction": true, @@ -17024,6 +17090,9 @@ "supports_native_structured_output": true, "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, + "claude-fable-latest": { + "alias_of": "claude-fable-5-1" + }, "claude-opus-5-5": { "supports_web_search": true, "supports_anthropic_compaction": true, @@ -17069,6 +17138,9 @@ "prompt_cache_min_tokens": 512, "source": "https://platform.claude.com/docs/en/about-claude/pricing" }, + "claude-opus-latest": { + "alias_of": "claude-opus-5-5" + }, "claude-opus-5": { "supports_web_search": true, "supports_anthropic_compaction": true, @@ -18868,6 +18940,9 @@ "supports_tool_choice": true, "supports_vision": true }, + "dashscope/qwen-max-latest": { + "alias_of": "dashscope/qwen3.8-max" + }, "dashscope/qwen3.8-flash": { "cache_creation_input_token_cost": 2e-07, "cache_read_input_token_cost": 1.6e-08, @@ -18887,6 +18962,9 @@ "supports_vision": true, "supports_web_search": true }, + "dashscope/qwen-flash-latest": { + "alias_of": "dashscope/qwen3.8-flash" + }, "dashscope/qwen3.8-omni-flash": { "cache_read_input_token_cost": 1.6e-08, "input_cost_per_token": 1.5e-07, @@ -21767,6 +21845,9 @@ "supports_vision": true, "thinking_always_on": true }, + "databricks/claude-fable-latest": { + "alias_of": "databricks/databricks-claude-fable-5-1" + }, "databricks/databricks-claude-haiku-4-5": { "cache_creation_input_token_cost": 1.24999e-06, "cache_creation_input_token_cost_above_1hr": 1.99997e-06, @@ -22007,6 +22088,9 @@ "supports_xhigh_reasoning_effort": true, "thinking_always_on": true }, + "databricks/claude-opus-latest": { + "alias_of": "databricks/databricks-claude-opus-5-5" + }, "databricks/databricks-claude-sonnet-4": { "cache_creation_input_token_cost": 3.74997e-06, "cache_read_input_token_cost": 3.0002e-07, @@ -22724,6 +22808,9 @@ "cache_read_input_token_cost_priority": 4.55e-07, "cache_creation_input_token_cost_priority": 2.45e-06 }, + "databricks/glm-latest": { + "alias_of": "databricks/databricks-glm-5-3" + }, "databricks/databricks-glm-5-3-flash": { "cache_creation_input_token_cost": 1.5001e-07, "cache_read_input_token_cost": 3.003e-08, @@ -22757,6 +22844,9 @@ "cache_read_input_token_cost_priority": 5.25e-08, "cache_creation_input_token_cost_priority": 2.625e-07 }, + "databricks/glm-flash-latest": { + "alias_of": "databricks/databricks-glm-5-3-flash" + }, "databricks/databricks-gpt-5": { "cache_creation_input_token_cost": 1.24999e-06, "cache_read_input_token_cost": 1.2502e-07, @@ -22919,6 +23009,9 @@ "cache_read_input_token_cost_priority": 1.5001e-07, "cache_creation_input_token_cost_priority": 1.50003e-06 }, + "databricks/gpt-mini-latest": { + "alias_of": "databricks/databricks-gpt-5-4-mini" + }, "databricks/databricks-gpt-5-4-nano": { "cache_creation_input_token_cost": 1.9999e-07, "cache_read_input_token_cost": 2.002e-08, @@ -23315,6 +23408,9 @@ "cache_read_input_token_cost_priority": 5.25e-07, "cache_creation_input_token_cost_priority": 5.25e-06 }, + "databricks/kimi-latest": { + "alias_of": "databricks/databricks-kimi-k3" + }, "databricks/system.ai.kimi-k3": { "litellm_provider": "databricks", "mode": "chat", @@ -28843,6 +28939,9 @@ "supports_tool_choice": true, "supports_vision": true }, + "fireworks_ai/minimax-latest": { + "alias_of": "fireworks_ai/accounts/fireworks/models/minimax-m3" + }, "fireworks_ai/accounts/fireworks/models/mixtral-8x22b-instruct-hf": { "input_cost_per_token": 1.2e-06, "litellm_provider": "fireworks_ai", @@ -29273,52 +29372,6 @@ "comment": "Open flagship GLM for long-horizon coding agents and million-token context work", "source": "https://api.friendli.ai/serverless/v1/models" }, - "friendliai/deepseek-ai/DeepSeek-V3.2": { - "litellm_provider": "friendliai", - "max_input_tokens": 163840, - "max_output_tokens": 163840, - "input_cost_per_token": 5e-07, - "output_cost_per_token": 1.5e-06, - "cache_read_input_token_cost": 2.5e-07, - "supports_prompt_caching": true, - "supports_reasoning": true, - "reasoning_effort_levels": [], - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_native_structured_output": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": false, - "supports_image_input": false, - "supports_video_input": false, - "mode": "chat", - "comment": "DeepSeek chat model for instruction following, coding, and analysis", - "source": "https://api.friendli.ai/serverless/v1/models" - }, - "friendliai/MiniMaxAI/MiniMax-M2.5": { - "litellm_provider": "friendliai", - "max_input_tokens": 196608, - "max_output_tokens": 196608, - "input_cost_per_token": 3e-07, - "output_cost_per_token": 1.2e-06, - "cache_read_input_token_cost": 6e-08, - "supports_prompt_caching": true, - "supports_reasoning": true, - "reasoning_effort_levels": [], - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_response_schema": true, - "supports_native_structured_output": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": false, - "supports_image_input": false, - "supports_video_input": false, - "mode": "chat", - "comment": "Prior MiniMax coding model for agent workflows, office edits, and automation", - "source": "https://api.friendli.ai/serverless/v1/models" - }, "friendliai/zai-org/GLM-5.1": { "litellm_provider": "friendliai", "max_input_tokens": 202752, @@ -30485,7 +30538,6 @@ "successor_model": "gemini-3.8-flash", "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "cache_read_input_token_cost_above_200k_tokens_priority": 4.5e-07, "cache_read_input_token_cost_flex": 1.25e-07, "cache_read_input_token_cost_priority": 2.25e-07, @@ -30544,7 +30596,6 @@ "output_cost_per_token_priority": 1.8e-05 }, "gemini-3.1-pro-preview": { - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, @@ -30608,7 +30659,6 @@ "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, "input_cost_per_token_batches": 1e-06, @@ -30660,7 +30710,6 @@ "vertex_ai/gemini-3-pro-preview": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, "input_cost_per_token_batches": 1e-06, @@ -31104,7 +31153,6 @@ "supports_web_search": true }, "vertex_ai/gemini-3.1-pro-preview": { - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, @@ -31168,7 +31216,6 @@ "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, "input_cost_per_token_batches": 1e-06, @@ -32947,6 +32994,7 @@ "input_cost_per_token_flex": 1e-06, "output_cost_per_token_flex": 6e-06, "cache_read_input_token_cost_above_200k_tokens_batches": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_flex": 4e-07, "input_cost_per_token_above_200k_tokens_batches": 2e-06, "input_cost_per_token_above_200k_tokens_flex": 2e-06, "output_cost_per_token_above_200k_tokens_batches": 9e-06, @@ -33015,6 +33063,7 @@ "input_cost_per_token_flex": 1e-06, "output_cost_per_token_flex": 6e-06, "cache_read_input_token_cost_above_200k_tokens_batches": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_flex": 4e-07, "input_cost_per_token_above_200k_tokens_batches": 2e-06, "input_cost_per_token_above_200k_tokens_flex": 2e-06, "output_cost_per_token_above_200k_tokens_batches": 9e-06, @@ -33789,6 +33838,9 @@ "supports_vision": true, "thinking_always_on": true }, + "github_copilot/claude-fable-latest": { + "alias_of": "github_copilot/claude-fable-5.1" + }, "github_copilot/claude-haiku-4.5": { "cache_creation_input_token_cost": 1.25e-06, "cache_read_input_token_cost": 1e-07, @@ -33832,6 +33884,9 @@ "supports_vision": true, "thinking_always_on": false }, + "github_copilot/claude-haiku-latest": { + "alias_of": "github_copilot/claude-haiku-5.5" + }, "github_copilot/claude-opus-4.8": { "cache_creation_input_token_cost": 6.25e-06, "cache_read_input_token_cost": 5e-07, @@ -33920,6 +33975,9 @@ "supports_vision": true, "thinking_always_on": true }, + "github_copilot/claude-opus-latest": { + "alias_of": "github_copilot/claude-opus-5.5" + }, "github_copilot/claude-sonnet-5": { "cache_creation_input_token_cost": 2.5e-06, "cache_read_input_token_cost": 2e-07, @@ -33963,6 +34021,9 @@ "supports_vision": true, "thinking_always_on": true }, + "github_copilot/claude-sonnet-latest": { + "alias_of": "github_copilot/claude-sonnet-5.5" + }, "github_copilot/gemini-3.7-flash": { "cache_read_input_token_cost": 7.5e-08, "deprecation_date": "2026-10-19", @@ -34308,6 +34369,9 @@ "supports_response_schema": true, "supports_vision": true }, + "github_copilot/gpt-astra-latest": { + "alias_of": "github_copilot/gpt-6-astra" + }, "github_copilot/gpt-6-luna": { "cache_creation_input_token_cost": 1.25e-07, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-07, @@ -34331,6 +34395,9 @@ "supports_response_schema": true, "supports_vision": true }, + "github_copilot/gpt-luna-latest": { + "alias_of": "github_copilot/gpt-6-luna" + }, "github_copilot/gpt-6-sol": { "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, @@ -34377,6 +34444,9 @@ "supports_response_schema": true, "supports_vision": true }, + "github_copilot/gpt-sol-latest": { + "alias_of": "github_copilot/gpt-6.1-sol" + }, "github_copilot/kimi-k3": { "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, @@ -34394,6 +34464,9 @@ "supports_response_schema": true, "supports_vision": true }, + "github_copilot/kimi-latest": { + "alias_of": "github_copilot/kimi-k3" + }, "github_copilot/mai-code-1.1-flash": { "cache_read_input_token_cost": 2e-08, "input_cost_per_token": 2e-07, @@ -37687,6 +37760,9 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "gpt-astra-latest": { + "alias_of": "gpt-6-astra" + }, "gpt-6-sol": { "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_12hr": 5e-06, @@ -37884,6 +37960,9 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "gpt-luna-latest": { + "alias_of": "gpt-6-luna" + }, "gpt-5.6": { "cache_creation_input_token_cost": 5e-06, "cache_creation_input_token_cost_above_12hr": 1e-05, @@ -39067,6 +39146,9 @@ "supports_xhigh_reasoning_effort": true, "supports_minimal_reasoning_effort": false }, + "gpt-mini-latest": { + "alias_of": "gpt-5.4-mini" + }, "gpt-5.4-mini-2026-03-17": { "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_batches": 3.75e-08, @@ -43025,6 +43107,9 @@ "max_input_tokens": 1000000, "max_output_tokens": 128000 }, + "minimax/minimax-latest": { + "alias_of": "minimax/MiniMax-M3" + }, "minimax/MiniMax-H3": { "litellm_provider": "minimax", "mode": "video_generation", @@ -44535,6 +44620,9 @@ "supports_video_input": true, "supports_vision": true }, + "moonshot/kimi-latest": { + "alias_of": "moonshot/kimi-k3" + }, "moonshot/moonshot-v1-128k": { "deprecation_date": "2026-08-31", "input_cost_per_token": 2e-06, @@ -50047,13 +50135,13 @@ "supports_reasoning": true }, "together_ai/Qwen/Qwen3.7-Max": { - "cache_read_input_token_cost": 3e-07, - "input_cost_per_token": 1.5e-06, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 2.5e-06, "litellm_provider": "together_ai", "max_input_tokens": 1000000, "max_output_tokens": 1000000, "mode": "chat", - "output_cost_per_token": 4.5e-06, + "output_cost_per_token": 7.5e-06, "source": "https://www.together.ai/pricing", "supports_prompt_caching": true }, @@ -50136,6 +50224,9 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "together_ai/deepseek-pro-latest": { + "alias_of": "together_ai/deepseek-ai/DeepSeek-V4-Pro-0813" + }, "together_ai/meta-models/Muse-Glimmer-30B": { "cache_read_input_token_cost": 4e-08, "input_cost_per_token": 3.5e-07, @@ -50182,6 +50273,9 @@ "supports_tool_choice": true, "supports_vision": true }, + "together_ai/kimi-latest": { + "alias_of": "together_ai/moonshotai/Kimi-K3" + }, "together_ai/thinkingmachines/Inkling": { "cache_read_input_token_cost": 1.7e-07, "deprecation_date": "2026-11-06", @@ -50232,6 +50326,9 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "together_ai/glm-latest": { + "alias_of": "together_ai/zai-org/GLM-5.3" + }, "together_ai/zai-org/GLM-5.3-Flash": { "cache_read_input_token_cost": 3e-08, "input_cost_per_token": 1.5e-07, @@ -50249,6 +50346,9 @@ "supports_tool_choice": true, "supports_vision": true }, + "together_ai/glm-flash-latest": { + "alias_of": "together_ai/zai-org/GLM-5.3-Flash" + }, "tts-1": { "deprecation_date": "2027-01-06", "deprecation_source": "https://developers.openai.com/api/docs/deprecations", @@ -53019,6 +53119,9 @@ "output_cost_per_token_batches": 2.5e-05, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" }, + "vertex_ai/claude-fable-latest": { + "alias_of": "vertex_ai/claude-fable-5-1" + }, "vertex_ai/claude-fable-5@default": { "deprecation_date": "2027-06-08", "regional_endpoint_uplift_multiplier": 1.1, @@ -53217,6 +53320,9 @@ "prompt_cache_min_tokens": 512, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" }, + "vertex_ai/claude-opus-latest": { + "alias_of": "vertex_ai/claude-opus-5-5" + }, "vertex_ai/claude-opus-5-5@default": { "regional_endpoint_uplift_multiplier": 1.1, "supports_mid_conversation_system": true, @@ -55765,6 +55871,9 @@ "supports_vision": true, "supports_web_search": true }, + "xai/grok-latest": { + "alias_of": "xai/grok-4.7" + }, "zai.glm-4.7": { "deprecation_date": "2027-03-30", "deprecation_source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-lifecycle-legacy.html", @@ -55866,6 +55975,9 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "zai/glm-latest": { + "alias_of": "zai/glm-5.3" + }, "zai/glm-5.3-flash": { "cache_creation_input_token_cost": 0, "cache_read_input_token_cost": 3e-08, @@ -55883,6 +55995,9 @@ "source": "https://docs.z.ai/guides/overview/pricing", "supports_vision": true }, + "zai/glm-flash-latest": { + "alias_of": "zai/glm-5.3-flash" + }, "zai/glm-5v-turbo": { "cache_read_input_token_cost": 2.4e-07, "input_cost_per_token": 1.2e-06, @@ -61519,6 +61634,9 @@ "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-6-astra.html", "supports_web_search": true }, + "bedrock_mantle/gpt-astra-latest": { + "alias_of": "bedrock_mantle/openai.gpt-6-astra" + }, "bedrock_mantle/openai.gpt-6-sol": { "input_cost_per_token": 2.2e-06, "input_cost_per_token_above_272k_tokens": 4.4e-06, @@ -61599,6 +61717,9 @@ "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-cards-openai.html", "supports_web_search": true }, + "bedrock_mantle/gpt-luna-latest": { + "alias_of": "bedrock_mantle/openai.gpt-6-luna" + }, "us.openai.gpt-6-astra": { "supports_bedrock_runtime_chat_completions_response_format": true, "input_cost_per_token": 1.1e-05, @@ -61767,6 +61888,9 @@ "/v1/responses" ] }, + "bedrock/gpt-astra-latest": { + "alias_of": "global.openai.gpt-6-astra" + }, "openai.gpt-6-sol": { "input_cost_per_token": 2e-06, "input_cost_per_token_above_272k_tokens": 4e-06, @@ -61909,6 +62033,9 @@ "/v1/responses" ] }, + "bedrock/gpt-luna-latest": { + "alias_of": "global.openai.gpt-6-luna" + }, "bedrock_mantle/openai.gpt-5.5": { "input_cost_per_token": 5.5e-06, "input_cost_per_token_above_272k_tokens": 1.1e-05, @@ -62186,6 +62313,9 @@ "supports_forced_tool_use": false, "token_count_model": {"model":"anthropic.claude-opus-5-5","litellm_provider":"bedrock_mantle","region":"us-east-1"} }, + "bedrock_mantle/claude-opus-latest": { + "alias_of": "bedrock_mantle/anthropic.claude-opus-5-5" + }, "bedrock_mantle/anthropic.claude-sonnet-5-5": { "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, @@ -62226,6 +62356,9 @@ "prompt_cache_min_tokens": 512, "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" }, + "bedrock_mantle/claude-sonnet-latest": { + "alias_of": "bedrock_mantle/anthropic.claude-sonnet-5-5" + }, "us.xai.grok-4.6": { "supports_regex_lookaround": false, "input_cost_per_token": 2.2e-06, @@ -63796,6 +63929,9 @@ "token_count_framing": 23, "token_count_tool_framing": 12 }, + "deepseek/deepseek-flash-latest": { + "alias_of": "deepseek/deepseek-v4-flash" + }, "deepseek/deepseek-v4-flash-vision-exp": { "cache_creation_input_token_cost": 0.0, "cache_read_input_token_cost": 6e-09, @@ -63904,6 +64040,9 @@ "token_count_framing": 76, "token_count_tool_framing": 75 }, + "deepseek/deepseek-pro-latest": { + "alias_of": "deepseek/deepseek-v4-pro" + }, "tencent/deepseek-v4-pro": { "cache_creation_input_token_cost": 0.0, "cache_read_input_token_cost": 1.45e-07, @@ -65327,6 +65466,9 @@ "supports_tool_choice": true, "supports_vision": true }, + "fireworks_ai/kimi-latest": { + "alias_of": "fireworks_ai/accounts/fireworks/models/kimi-k3" + }, "fireworks_ai/deepseek-v4-flash-0731": { "cache_read_input_token_cost": 7e-09, "cache_read_input_token_cost_priority": 8.75e-09, @@ -67576,7 +67718,7 @@ "deepinfra/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B": { "max_output_tokens": 262144, "max_input_tokens": 262144, - "input_cost_per_token": 8.5e-08, + "input_cost_per_token": 7.5e-08, "output_cost_per_token": 4e-07, "litellm_provider": "deepinfra", "mode": "chat", @@ -69000,6 +69142,9 @@ "output_cost_per_token": 4.7e-07, "source": "https://api.together.ai/v1/models" }, + "together_ai/qwen-flash-latest": { + "alias_of": "together_ai/Qwen/Qwen3.8-Flash" + }, "together_ai/moonshotai/Kimi-K2.5-fp4": { "input_cost_per_token": 5e-07, "output_cost_per_token": 2.8e-06, @@ -69216,6 +69361,9 @@ "supports_reasoning": true, "supports_tool_choice": true }, + "azure_ai/glm-latest": { + "alias_of": "azure_ai/FW-GLM-5.3" + }, "azure_ai/FW-GLM-5.3-Flash": { "cache_read_input_token_cost": 3.8e-08, "input_cost_per_token": 1.88e-07, @@ -69231,6 +69379,9 @@ "supports_reasoning": true, "supports_tool_choice": true }, + "azure_ai/glm-flash-latest": { + "alias_of": "azure_ai/FW-GLM-5.3-Flash" + }, "azure_ai/FW-GPT-OSS-120B": { "deprecation_date": "2027-07-01", "cache_read_input_token_cost": 8.2e-08, @@ -71236,6 +71387,9 @@ "supports_tool_choice": true, "supports_vision": true }, + "baseten/glm-latest": { + "alias_of": "baseten/zai-org/GLM-5.3" + }, "baseten/zai-org/GLM-5.3-Fast": { "cache_read_input_token_cost": 2.1e-07, "input_cost_per_token": 2.1e-06, @@ -71579,6 +71733,13 @@ "output_cost_per_token": 1.7e-06, "source": "https://api.together.ai/v1/models" }, + "together_ai/essentialai/rnj-1-instruct": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "together_ai", + "mode": "chat", + "output_cost_per_token": 1.5e-07, + "source": "https://www.together.ai/models/rnj-1-instruct" + }, "together_ai/google/gemma-4-31B-it": { "deprecation_date": "2026-09-15", "deprecation_source": "https://docs.together.ai/docs/deprecations", @@ -72720,6 +72881,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/gpt-realtime-2": { + "cache_read_input_image_token_cost": 5e-07, "cache_creation_input_audio_token_cost": 4e-07, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, @@ -72738,7 +72900,7 @@ "provider_specific_entry": { "realtime_ga_only": 1 }, - "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "source": "https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/", "supported_endpoints": [ "/v1/realtime" ], @@ -74425,6 +74587,9 @@ "supports_tool_choice": true, "supports_vision": true }, + "bedrock/kimi-latest": { + "alias_of": "global.moonshotai.kimi-k3" + }, "us.moonshotai.kimi-k3": { "supports_regex_lookaround": false, "cache_creation_input_token_cost": 4.125e-06, @@ -74975,6 +75140,9 @@ "supports_tool_choice": true, "supports_vision": true }, + "baseten/kimi-latest": { + "alias_of": "baseten/moonshotai/Kimi-K3" + }, "baseten/nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B": { "cache_read_input_token_cost": 1.2e-07, "input_cost_per_token": 6e-07, @@ -75057,6 +75225,9 @@ "supports_tool_choice": true, "supports_vision": true }, + "baseten/glm-flash-latest": { + "alias_of": "baseten/zai-org/GLM-5.3-Flash" + }, "baseten/deepseek-ai/DeepSeek-V4-Flash-0731": { "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 1.3e-07, @@ -75073,6 +75244,9 @@ "supports_tool_choice": true, "supports_vision": false }, + "baseten/deepseek-flash-latest": { + "alias_of": "baseten/deepseek-ai/DeepSeek-V4-Flash-0731" + }, "baseten/deepseek-ai/DeepSeek-V4-Pro": { "deprecation_date": "2026-10-16", "cache_read_input_token_cost": 1.45e-07, @@ -75106,6 +75280,9 @@ "supports_tool_choice": true, "supports_vision": false }, + "baseten/deepseek-pro-latest": { + "alias_of": "baseten/deepseek-ai/DeepSeek-V4-Pro-0813" + }, "baseten/zai-org/GLM-5.2-Fast": { "cache_read_input_token_cost": 2.1e-07, "input_cost_per_token": 2.1e-06, @@ -76097,6 +76274,9 @@ "prompt_cache_min_tokens": 512, "source": "https://management.azure.com/subscriptions/c873328e-b572-4770-8dff-aaeb6f1f0e79/providers/Microsoft.CognitiveServices/locations/eastus2/models?api-version=2024-10-01" }, + "azure_ai/claude-sonnet-latest": { + "alias_of": "azure_ai/claude-sonnet-5-5" + }, "bedrock/us-gov-east-1/anthropic.claude-sonnet-5-5": { "bedrock_converse_supports_strict_tools": false, "bedrock_output_config_effort_ceiling": "xhigh", @@ -76261,6 +76441,9 @@ "prompt_cache_min_tokens": 512, "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" }, + "bedrock/claude-sonnet-latest": { + "alias_of": "global.anthropic.claude-sonnet-5-5" + }, "jp.anthropic.claude-sonnet-5-5": { "bedrock_converse_supports_strict_tools": false, "cache_creation_input_token_cost": 2.75e-06, @@ -76520,6 +76703,9 @@ "prompt_cache_min_tokens": 512, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" }, + "vertex_ai/claude-sonnet-latest": { + "alias_of": "vertex_ai/claude-sonnet-5-5" + }, "vertex_ai/claude-sonnet-5-5@default": { "regional_endpoint_uplift_multiplier": 1.1, "supports_mid_conversation_system": true, @@ -76632,6 +76818,9 @@ "supports_tool_choice": true, "supports_vision": true }, + "bedrock/grok-latest": { + "alias_of": "global.xai.grok-4.7" + }, "us.xai.grok-4.7": { "supports_regex_lookaround": false, "supported_endpoints": ["/v1/chat/completions"], @@ -76804,6 +76993,9 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "gpt-sol-latest": { + "alias_of": "gpt-6.1-sol" + }, "global.openai.gpt-6.1-sol": { "cache_creation_input_token_cost": 2.5e-06, "cache_creation_input_token_cost_above_272k_tokens": 5e-06, @@ -76850,6 +77042,9 @@ "supports_sampling_params": false, "supports_xhigh_reasoning_effort": true }, + "bedrock/gpt-sol-latest": { + "alias_of": "global.openai.gpt-6.1-sol" + }, "bedrock_mantle/openai.gpt-6.1-sol": { "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_272k_tokens": 5.5e-06, @@ -76898,6 +77093,9 @@ "supports_xhigh_reasoning_effort": true, "use_openai_responses_path": true }, + "bedrock_mantle/gpt-sol-latest": { + "alias_of": "bedrock_mantle/openai.gpt-6.1-sol" + }, "us.openai.gpt-6.1-sol": { "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_272k_tokens": 5.5e-06, @@ -76989,6 +77187,9 @@ "supports_tool_choice": true, "supports_vision": true }, + "vertex_ai/grok-latest": { + "alias_of": "vertex_ai/xai/grok-4.7" + }, "azure_ai/kimi-k2-thinking": { "deprecation_date": "2026-03-29", "deprecation_source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", @@ -77046,6 +77247,9 @@ "supports_tool_choice": true, "supports_vision": false }, + "bedrock/glm-latest": { + "alias_of": "global.zai.glm-5.3" + }, "us.zai.glm-5.3": { "cache_creation_input_token_cost": 2.31e-06, "cache_creation_input_token_cost_batches": 1.155e-06, @@ -77269,6 +77473,9 @@ "supports_vision": true, "supports_web_search": true }, + "azure_ai/grok-latest": { + "alias_of": "azure_ai/grok-4.7" + }, "in.moonshotai.kimi-k3": { "cache_creation_input_token_cost": 4.125e-06, "cache_creation_input_token_cost_flex": 2.0625e-06, @@ -77355,6 +77562,9 @@ "cache_creation_input_token_cost_above_1hr_above_100k_tokens": 1e-06, "cache_read_input_token_cost_above_100k_tokens": 5e-08 }, + "claude-haiku-latest": { + "alias_of": "claude-haiku-5-5" + }, "bedrock_mantle/anthropic.claude-haiku-5-5": { "input_cost_per_token_above_100k_tokens": 5.5e-07, "output_cost_per_token_above_100k_tokens": 2.75e-06, @@ -77401,6 +77611,9 @@ "prompt_cache_min_tokens": 512, "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-haiku-5-5.html" }, + "bedrock_mantle/claude-haiku-latest": { + "alias_of": "bedrock_mantle/anthropic.claude-haiku-5-5" + }, "bedrock_mantle/us-gov-west-1/anthropic.claude-haiku-5-5": { "input_cost_per_token_above_100k_tokens": 6e-07, "output_cost_per_token_above_100k_tokens": 3e-06, @@ -77632,6 +77845,9 @@ "prompt_cache_min_tokens": 512, "source": "https://platform.claude.com/docs/en/models/haiku-5-5/migration-guide" }, + "azure_ai/claude-haiku-latest": { + "alias_of": "azure_ai/claude-haiku-5-5" + }, "bedrock/us-gov-east-1/anthropic.claude-haiku-5-5": { "input_cost_per_token_above_100k_tokens": 6e-07, "output_cost_per_token_above_100k_tokens": 3e-06, @@ -77822,6 +78038,9 @@ "prompt_cache_min_tokens": 512, "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" }, + "bedrock/claude-haiku-latest": { + "alias_of": "global.anthropic.claude-haiku-5-5" + }, "jp.anthropic.claude-haiku-5-5": { "input_cost_per_token_above_100k_tokens": 5.5e-07, "output_cost_per_token_above_100k_tokens": 2.75e-06, @@ -78029,6 +78248,9 @@ "prompt_cache_min_tokens": 512, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" }, + "vertex_ai/claude-haiku-latest": { + "alias_of": "vertex_ai/claude-haiku-5-5" + }, "vertex_ai/claude-haiku-5-5@default": { "input_cost_per_token_above_100k_tokens": 5e-07, "output_cost_per_token_above_100k_tokens": 2.5e-06, @@ -78933,6 +79155,9 @@ "supports_response_schema": true, "supports_vision": true }, + "github_copilot/grok-latest": { + "alias_of": "github_copilot/grok-4.7" + }, "together_ai/deepseek-ai/DeepSeek-V4.1-Flash-1017": { "cache_read_input_token_cost": 6e-09, "input_cost_per_token": 3e-07, @@ -80356,6 +80581,9 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "gpt-terra-latest": { + "alias_of": "gpt-6.1-terra" + }, "chatgpt/gpt-6.1-terra": { "litellm_provider": "chatgpt", "max_input_tokens": 922000, @@ -80414,6 +80642,9 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true }, + "bedrock/gpt-terra-latest": { + "alias_of": "global.openai.gpt-6.1-terra" + }, "us.openai.gpt-6.1-terra": { "cache_creation_input_token_cost": 1.375e-06, "cache_creation_input_token_cost_above_272k_tokens": 2.75e-06, @@ -80493,6 +80724,9 @@ "supports_response_schema": true, "supports_vision": true }, + "github_copilot/gpt-terra-latest": { + "alias_of": "github_copilot/gpt-6.1-terra" + }, "refuseless/glm-5.3-abliterated": { "cache_read_input_token_cost": 1.5e-07, "input_cost_per_token": 1e-06, @@ -80647,6 +80881,9 @@ "supports_xhigh_reasoning_effort": true, "use_openai_responses_path": true }, + "bedrock_mantle/gpt-terra-latest": { + "alias_of": "bedrock_mantle/openai.gpt-6.1-terra" + }, "clf_ai_gateway/glm-5.3": { "metadata": { "comment": "Billed price checked 2026-10-29, which the gateway shows as 60% off a list price 2.5x higher. Update the rows when the discount ends." @@ -82196,6 +82433,9 @@ "supports_vision": true, "thinking_always_on": false }, + "databricks/claude-haiku-latest": { + "alias_of": "databricks/databricks-claude-haiku-5-5" + }, "databricks/databricks-claude-sonnet-5-5": { "cache_creation_input_token_cost": 2.49998e-06, "cache_creation_input_token_cost_above_1hr": 4.00001e-06, @@ -82225,6 +82465,9 @@ "supports_vision": false, "thinking_always_on": true }, + "databricks/claude-sonnet-latest": { + "alias_of": "databricks/databricks-claude-sonnet-5-5" + }, "databricks/databricks-deepseek-v4-1-flash": { "cache_creation_input_token_cost": 3.0002e-07, "cache_read_input_token_cost": 3.003e-08, @@ -82457,6 +82700,9 @@ "cache_read_input_token_cost_above_272k_tokens_priority": 3.9998e-07, "cache_creation_input_token_cost_above_272k_tokens_priority": 9.99999e-06 }, + "databricks/gpt-sol-latest": { + "alias_of": "databricks/databricks-gpt-6-1-sol" + }, "databricks/databricks-gpt-6-1-terra": { "cache_creation_input_token_cost": 1.24999e-06, "cache_creation_input_token_cost_above_272k_tokens": 2.49998e-06, @@ -82497,6 +82743,9 @@ "cache_read_input_token_cost_above_272k_tokens_priority": 1.9999e-07, "cache_creation_input_token_cost_above_272k_tokens_priority": 5.00003e-06 }, + "databricks/gpt-terra-latest": { + "alias_of": "databricks/databricks-gpt-6-1-terra" + }, "databricks/databricks-gpt-6-astra": { "cache_creation_input_token_cost": 1.249997e-05, "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, @@ -82537,6 +82786,9 @@ "cache_read_input_token_cost_above_272k_tokens_priority": 4.00001e-06, "cache_creation_input_token_cost_above_272k_tokens_priority": 5.000002e-05 }, + "databricks/gpt-astra-latest": { + "alias_of": "databricks/databricks-gpt-6-astra" + }, "databricks/databricks-gpt-6-luna": { "cache_creation_input_token_cost": 1.2502e-07, "cache_creation_input_token_cost_above_272k_tokens": 2.4997e-07, @@ -82577,6 +82829,9 @@ "cache_read_input_token_cost_above_272k_tokens_priority": 3.997e-08, "cache_creation_input_token_cost_above_272k_tokens_priority": 5.0001e-07 }, + "databricks/gpt-luna-latest": { + "alias_of": "databricks/databricks-gpt-6-luna" + }, "databricks/databricks-gpt-6-sol": { "cache_creation_input_token_cost": 2.49998e-06, "cache_creation_input_token_cost_above_272k_tokens": 5.00003e-06, @@ -82707,6 +82962,9 @@ "supports_tool_choice": true, "supports_vision": true }, + "databricks/grok-latest": { + "alias_of": "databricks/databricks-grok-4-7" + }, "databricks/databricks-kimi-k2-7": { "cache_creation_input_token_cost": 9.4997e-07, "cache_read_input_token_cost": 1.8998e-07, @@ -86639,6 +86897,9 @@ "supports_vision": true, "use_openai_responses_path": true }, + "bedrock_mantle/grok-latest": { + "alias_of": "bedrock_mantle/xai.grok-4.7" + }, "together_ai/ricursive-ai/Ripple-65B-Preview": { "cache_read_input_token_cost": 2e-08, "input_cost_per_token": 2e-07, @@ -87049,5 +87310,484 @@ "supports_parallel_function_calling": true, "supports_system_messages": true, "supports_tool_choice": true + }, + "anthropic.claude-fable-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "eu.anthropic.claude-fable-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.1e-05, + "cache_creation_input_token_cost_above_1hr": 1.76e-05, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 8.8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output_by_default": true, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "eu.anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.1e-05, + "cache_creation_input_token_cost_above_1hr": 1.76e-05, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 8.8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "global.anthropic.claude-fable-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output_by_default": true, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "global.anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "us-gov.anthropic.claude-fable-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.2e-05, + "cache_creation_input_token_cost_above_1hr": 1.92e-05, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 9.6e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.8e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "us-gov.anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.2e-05, + "cache_creation_input_token_cost_above_1hr": 1.92e-05, + "cache_read_input_token_cost": 2.4e-07, + "input_cost_per_token": 9.6e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.8e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "us.anthropic.claude-fable-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.1e-05, + "cache_creation_input_token_cost_above_1hr": 1.76e-05, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 8.8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_bedrock_native_structured_output_by_default": true, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": false, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "us.anthropic.claude-mythos-5-2": { + "bedrock_output_config_effort_ceiling": "xhigh", + "cache_creation_input_token_cost": 1.1e-05, + "cache_creation_input_token_cost_above_1hr": 1.76e-05, + "cache_read_input_token_cost": 2.2e-07, + "input_cost_per_token": 8.8e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4.4e-05, + "prompt_cache_min_tokens": 512.0, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_parallel_tool_use_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "claude-fable-5-2": { + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_creation_input_token_cost_batches": 5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, + "input_cost_per_token": 8e-06, + "input_cost_per_token_batches": 4e-06, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "output_cost_per_token_batches": 2e-05, + "prompt_cache_min_tokens": 512.0, + "provider_specific_entry": {"us":1.1,"fast":2}, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://platform.claude.com/docs/en/about-claude/pricing", + "supports_adaptive_thinking": true, + "supports_anthropic_compaction": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_fast_mode": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_speed": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "bedrock/claude_platform/claude-fable-5-2": { + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_creation_input_token_cost_batches": 5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, + "input_cost_per_token": 8e-06, + "input_cost_per_token_batches": 4e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "output_cost_per_token_batches": 2e-05, + "prompt_cache_min_tokens": 512.0, + "provider_specific_entry": {"us":1.1}, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing", + "supports_adaptive_thinking": true, + "supports_anthropic_compaction": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_output_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "azure_ai/claude-fable-5-2": { + "cache_creation_input_token_cost": 1e-05, + "cache_creation_input_token_cost_above_1hr": 1.6e-05, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 4e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://platform.claude.com/docs/en/about-claude/pricing", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": false, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true + }, + "azure_ai/claude-fable-5-2-fast": { + "cache_creation_input_token_cost": 2e-05, + "cache_creation_input_token_cost_above_1hr": 3.2e-05, + "cache_read_input_token_cost": 4e-07, + "input_cost_per_token": 1.6e-05, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 8e-05, + "prompt_cache_min_tokens": 512.0, + "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "source": "https://platform.claude.com/docs/en/about-claude/pricing", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_max_reasoning_effort": false, + "supports_mid_conversation_system": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "thinking_always_on": true } } diff --git a/tests/e2e/AGENTS.md b/tests/e2e/AGENTS.md index a4c51a53098..89bf24016da 100644 --- a/tests/e2e/AGENTS.md +++ b/tests/e2e/AGENTS.md @@ -50,7 +50,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `gateway/` - proxy configuration only (`litellm-config.yml`). No tests - `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture. Ships its own driver/builder/publisher, covered by `tests/e2e_harness/claude_code/`. The HTTP probes ride the shared transport (`ProxyClient.count_tokens` / `ProxyClient.messages`). The CLI-driving path stays bespoke. `pr_gate_version_resolver.py` only forwards to `cli_harness/npm_version_resolver.py`, because the Buildkite runner images still run that path and install @latest without it - `cli_harness/` - no tests, only the plumbing every suite that drives a real coding-agent CLI shares: `npm_version_resolver.py` picks the CLI release to install (the newest npm release at least three days old, run by CI as `python tests/e2e/cli_harness/npm_version_resolver.py `, stdlib only), `lockdown.py` gives the CLI subprocess an allowlisted env and a fresh empty HOME so it never sees the runner's credentials or dotfiles, and `rate_limiter.py` is the cross-process token bucket per upstream that keeps parallel CLI runs under one request rate. Covered by `tests/e2e_harness/cli_harness/`. Code only one CLI suite needs stays in that suite -- `codex/` - the Codex compatibility suite: drives the real `codex exec --json` CLI through the gateway, one class per row (the command round trip, a file edit with the patch tool, a long session that follows a chain of files one command at a time, auto-compaction, where Codex reads a random value, compacts through the gateway and is asked for the value after the file is deleted, and an MCP tool, where Codex calls the tool of `mcp_server.py`, a stdio MCP server the test hands it with `add_mcp_server`, on a first and a resumed turn), each row parametrized over `SETUPS` in `setups.py`, one `pytest.param` per setup (GPT on Bedrock Mantle, Bedrock Runtime, Azure OpenAI and OpenAI, plus Claude on Bedrock). A row with one setup broken by a known bug parametrizes over `with_expected_failure(, reason=...)`, which is `SETUPS` with a strict xfail on that one setup. The `setup` fixture registers that setup's deployment under a fresh name via `/model/new` and the `codex` fixture makes a key scoped to it, and both are deleted on teardown. `harness.py` gives each test a throwaway HOME and CODEX_HOME whose `config.toml` points a custom `responses` provider at the gateway and has Codex read its model list from the gateway's `/v1/models` (`model_catalog_url` on the provider), with the sandbox on (`workspace-write`) and approvals off, and runs the CLI under `cli_harness`'s env lockdown and rate limiter. `run.py` reads the JSONL events into a `CodexRun`: the commands with their output, the patched files, the MCP tool calls with the text they returned, the last agent message as `answer`, the usage and the errors, where an error's `status_code` is what lets the suite's 5xx rerun rule catch a turn Codex gave up on. `codex exec --json` prints nothing when Codex compacts, so `rollout.py` reads the compactions from the thread's rollout file under CODEX_HOME/sessions, and a compaction is remote when its replacement history kept the provider's `compaction` item. `exec` and `resume` take `compact_after_tokens`, which forces Codex's remote compaction (`compaction_trigger` items over `/v1/responses`, the path Codex otherwise only takes for OpenAI and Azure) on the gateway provider and compacts once the context grows past that many tokens. A test passes on a random value the model can only learn through the gateway, never on the model's wording or the exact command it typed. It needs `codex` on PATH and fails without it. Covered by `tests/e2e_harness/codex/` +- `codex/` - the Codex compatibility suite: drives the real `codex exec --json` CLI through the gateway, one class per row (the command round trip, a file edit with the patch tool, a long session that follows a chain of files one command at a time, auto-compaction, where Codex reads a random value, compacts through the gateway and is asked for the value after the file is deleted, and an MCP tool, where Codex calls the tool of `mcp_server.py`, a stdio MCP server the test hands it with `add_mcp_server`, on a first and a resumed turn, and a sub-agent, where Codex's multi-agent v2 hands a helper a random codeword as its task with shell tools off so the helper can only learn it from the task, and the helper's own final reply is read from its rollout), each row parametrized over `SETUPS` in `setups.py`, one `pytest.param` per setup (GPT on Bedrock Mantle, Bedrock Runtime, Azure OpenAI and OpenAI, plus Claude on Bedrock). A row with one setup broken by a known bug parametrizes over `with_expected_failure(, reason=...)`, which is `SETUPS` with a strict xfail on that one setup. The `setup` fixture registers that setup's deployment under a fresh name via `/model/new` and the `codex` fixture makes a key scoped to it, and both are deleted on teardown. `harness.py` gives each test a throwaway HOME and CODEX_HOME whose `config.toml` points a custom `responses` provider at the gateway and has Codex read its model list from the gateway's `/v1/models` (`model_catalog_url` on the provider), with the sandbox on (`workspace-write`) and approvals off, and runs the CLI under `cli_harness`'s env lockdown and rate limiter. `run.py` reads the JSONL events into a `CodexRun`: the commands with their output, the patched files, the MCP tool calls with the text they returned, the last agent message as `answer`, the usage and the errors, where an error's `status_code` is what lets the suite's 5xx rerun rule catch a turn Codex gave up on. `codex exec --json` prints nothing when Codex compacts, so `rollout.py` reads the compactions from the thread's rollout file under CODEX_HOME/sessions, and a compaction is remote when its replacement history kept the provider's `compaction` item. `subagents` reads the sub-agents a thread spawned from their own rollouts, whose `session_meta` names that thread as their parent, each with its final replies and whether it forked the parent's conversation (`forked_from_id`), so the sub-agent row can reject a helper that saw the codeword in inherited history instead of its task. `exec` takes `enable` and `disable`, Codex feature names it passes as `--enable` and `--disable`. `exec` and `resume` take `compact_after_tokens`, which forces Codex's remote compaction (`compaction_trigger` items over `/v1/responses`, the path Codex otherwise only takes for OpenAI and Azure) on the gateway provider and compacts once the context grows past that many tokens. A test passes on a random value the model can only learn through the gateway, never on the model's wording or the exact command it typed. It needs `codex` on PATH and fails without it. Covered by `tests/e2e_harness/codex/` - `ui/` - the Admin UI browser suite: Playwright in TypeScript, driving the dashboard served by a live proxy on port 4000 (seeded postgres + mock LLM upstream, see its `run_e2e.sh`). It is a self-contained npm package with its own lockfile and does not use the Python harness, pytest markers, or the shared transport. The Python rules in this file (typed models, `Result` unions, basedpyright zero-error gate) do not apply inside it. Its only Python file, `fixtures/mock_llm_server/server.py`, is excluded from the e2e basedpyright gate via the root `pyrightconfig.json`. Its specs follow the dashboard's rule against testing what shadcn and Base UI already guarantee, in `ui/litellm-dashboard/AGENTS.md` ## MCP suite: real Datadog only diff --git a/tests/e2e/codex/harness.py b/tests/e2e/codex/harness.py index df005b5b5e7..cc29581f2ce 100644 --- a/tests/e2e/codex/harness.py +++ b/tests/e2e/codex/harness.py @@ -12,7 +12,7 @@ from typing import Final from cli_harness.lockdown import isolated_home, locked_down_env from cli_harness.rate_limiter import RateBucket, RateLimiter -from codex.rollout import Compaction, rollout_compactions +from codex.rollout import Compaction, Subagent, rollout_compactions, rollout_parent_thread, rollout_subagent from codex.run import CodexRun, codex_run from e2e_metadata import step @@ -71,8 +71,28 @@ class CodexHarness: _ = config.write(_mcp_server_table(name, server)) @step('Run Codex on {model} through the proxy with the prompt "{prompt}"') - def exec(self, prompt: str, *, model: str, cwd: Path, compact_after_tokens: int | None = None) -> CodexRun: - return self._run(("exec",), (*_compaction_flags(compact_after_tokens), "--model", model, "--", prompt), cwd) + def exec( + self, + prompt: str, + *, + model: str, + cwd: Path, + compact_after_tokens: int | None = None, + enable: tuple[str, ...] = (), + disable: tuple[str, ...] = (), + ) -> CodexRun: + return self._run( + ("exec",), + ( + *_compaction_flags(compact_after_tokens), + *_feature_flags(enable=enable, disable=disable), + "--model", + model, + "--", + prompt, + ), + cwd, + ) @step('Resume Codex thread {thread_id} on {model} through the proxy with the prompt "{prompt}"') def resume( @@ -88,6 +108,12 @@ class CodexHarness: rollouts: Final = sorted(self.codex_home.glob(f"sessions/**/rollout-*-{thread_id}.jsonl")) return tuple(chain.from_iterable(rollout_compactions(rollout.read_text()) for rollout in rollouts)) + def subagents(self, thread_id: str) -> tuple[Subagent, ...]: + rollouts: Final = tuple( + path.read_text() for path in sorted(self.codex_home.glob("sessions/**/rollout-*.jsonl")) + ) + return tuple(rollout_subagent(rollout) for rollout in rollouts if rollout_parent_thread(rollout) == thread_id) + def _run(self, subcommand: tuple[str, ...], args: tuple[str, ...], cwd: Path) -> CodexRun: _ = self.limiter.acquire(self.rate_bucket) completed: Final = subprocess.run( @@ -116,6 +142,13 @@ def _mcp_server_table(name: str, server: McpServer) -> str: ) +def _feature_flags(*, enable: tuple[str, ...], disable: tuple[str, ...]) -> tuple[str, ...]: + return ( + *chain.from_iterable(("--enable", feature) for feature in enable), + *chain.from_iterable(("--disable", feature) for feature in disable), + ) + + def _compaction_flags(after_tokens: int | None) -> tuple[str, ...]: match after_tokens: case None: diff --git a/tests/e2e/codex/rollout.py b/tests/e2e/codex/rollout.py index fcaeaa7ca42..5b8253fab10 100644 --- a/tests/e2e/codex/rollout.py +++ b/tests/e2e/codex/rollout.py @@ -1,19 +1,23 @@ -"""The compactions Codex recorded in a thread's rollout file, `CODEX_HOME/sessions/**/rollout-*-.jsonl`. +"""What Codex recorded in a thread's rollout file, `CODEX_HOME/sessions/**/rollout-*-.jsonl`. -`codex exec --json` prints no event for a compaction, so the rollout is the only place it shows. Each line is -`{"timestamp", "type", "payload"}` (`codex-rs/history/src/rollout_payload.rs`) and only `compacted` lines are read. +`codex exec --json` prints no event for a compaction or for a sub-agent's own turn, so the rollout is the only place +they show. Each line is `{"timestamp", "type", "payload"}` (`codex-rs/history/src/rollout_payload.rs`). A remote compaction is one whose replacement history carries the provider's `compaction` item, while a local one -replaces the history with a summary Codex wrote itself. +replaces the history with a summary Codex wrote itself. A sub-agent gets a rollout of its own whose `session_meta` +names the thread that spawned it, and carries `forked_from_id` when it inherited that thread's conversation. Each turn +it finishes ends in a `task_complete` event with its last message. """ from __future__ import annotations from dataclasses import dataclass -from typing import Final +from typing import Final, Literal from pydantic import BaseModel, ConfigDict _COMPACTED: Final = "compacted" +_SESSION_META: Final = "session_meta" +_EVENT: Final = "event_msg" _PROVIDER_COMPACTION_ITEM: Final = "compaction" @@ -34,14 +38,58 @@ class _CompactedPayload(_Frozen): replacement_history: tuple[_HistoryItem, ...] | None = None +class _SessionMeta(_Frozen): + parent_thread_id: str | None = None + forked_from_id: str | None = None + + +class _Event(_Frozen): + type: str + + +class _TaskComplete(_Frozen): + type: Literal["task_complete"] + last_agent_message: str | None = None + + @dataclass(frozen=True, slots=True) class Compaction: remote: bool +@dataclass(frozen=True, slots=True) +class Subagent: + forked: bool + replies: tuple[str, ...] + + def rollout_compactions(rollout: str) -> tuple[Compaction, ...]: - lines: Final = tuple(_Line.model_validate_json(line) for line in rollout.splitlines() if line.strip()) - return tuple(_compaction(line.payload) for line in lines if line.type == _COMPACTED) + return tuple(_compaction(line.payload) for line in _lines(rollout) if line.type == _COMPACTED) + + +def rollout_parent_thread(rollout: str) -> str | None: + return _session_meta(rollout).parent_thread_id + + +def rollout_subagent(rollout: str) -> Subagent: + return Subagent(forked=_session_meta(rollout).forked_from_id is not None, replies=rollout_final_messages(rollout)) + + +def rollout_final_messages(rollout: str) -> tuple[str, ...]: + events: Final = (line.payload for line in _lines(rollout) if line.type == _EVENT) + completions: Final = ( + _TaskComplete.model_validate(event) for event in events if _Event.model_validate(event).type == "task_complete" + ) + return tuple(done.last_agent_message for done in completions if done.last_agent_message is not None) + + +def _session_meta(rollout: str) -> _SessionMeta: + meta: Final = next((line.payload for line in _lines(rollout) if line.type == _SESSION_META), None) + return _SessionMeta.model_validate(meta) if meta is not None else _SessionMeta() + + +def _lines(rollout: str) -> tuple[_Line, ...]: + return tuple(_Line.model_validate_json(line) for line in rollout.splitlines() if line.strip()) def _compaction(payload: object) -> Compaction: diff --git a/tests/e2e/codex/test_subagent_e2e.py b/tests/e2e/codex/test_subagent_e2e.py new file mode 100644 index 00000000000..0ad86c8b49a --- /dev/null +++ b/tests/e2e/codex/test_subagent_e2e.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +from typing import Final + +import pytest +from codex.harness import CodexHarness +from codex.setups import SETUPS, ModelSetup +from e2e_config import unique_marker + +pytestmark = pytest.mark.e2e + + +@pytest.mark.parametrize("deployment", SETUPS) +class TestCodexSubagent: + def test_a_subagent_reads_the_task_its_parent_hands_it(self, codex: CodexHarness, setup: ModelSetup) -> None: + codeword: Final = unique_marker() + run: Final = codex.exec( + 'Spawn exactly one sub-agent with the spawn_agent tool, with fork_turns set to "none" so it does not' + f" inherit this conversation. Its task: reply with exactly this codeword and nothing else: {codeword}." + " Wait for its reply, then tell me the exact reply it sent you, verbatim.", + model=setup.model, + cwd=codex.new_workspace(files={}), + enable=("multi_agent_v2",), + disable=("shell_tool", "unified_exec"), + ) + assert run.exit_code == 0, run + assert run.thread_id is not None, run + subagents: Final = codex.subagents(run.thread_id) + assert subagents, run + assert not any(subagent.forked for subagent in subagents), (subagents, run) + assert any(any(codeword in reply for reply in subagent.replies) for subagent in subagents), (subagents, run) diff --git a/tests/e2e/llm_translation/test_together_ai_e2e.py b/tests/e2e/llm_translation/test_together_ai_e2e.py index 2d5d5f2d93a..7775c5c6ed2 100644 --- a/tests/e2e/llm_translation/test_together_ai_e2e.py +++ b/tests/e2e/llm_translation/test_together_ai_e2e.py @@ -7,9 +7,9 @@ backends are pinned because the registry has no flag for what they prove: ``enab and the ``{"reasoning": {"enabled": false}}`` toggle that ``reasoning_effort="none"`` maps to are Qwen hybrid-model contracts, and the structured-output case runs on that hybrid model with reasoning off, since a reasoning-only model can spend the whole token budget -thinking and return no content. MiniMax-M3 is the serverless model whose +thinking and return no content. MiniMax-M3.1-Flash is the serverless model whose template renders a replayed ``reasoning_content`` back into the prompt (Qwen and -DeepSeek silently drop it). MiniMax-M3 honors that replayed field on nearly every call, not +DeepSeek silently drop it). MiniMax-M3.1-Flash honors that replayed field on nearly every call, not every call (one miss in dozens of otherwise identical calls), so the replay case asks up to ``REPLAY_ATTEMPTS`` times and fails only when no answer carries the secret, which a proxy that strips the field guarantees. Requires TOGETHER_API_KEY on the proxy; no @@ -56,7 +56,7 @@ from pydantic import BaseModel pytestmark = pytest.mark.e2e HYBRID_REASONING_BACKEND = "together_ai/Qwen/Qwen3.5-9B" -REASONING_REPLAY_BACKEND = "together_ai/MiniMaxAI/MiniMax-M3" +REASONING_REPLAY_BACKEND = "together_ai/MiniMaxAI/MiniMax-M3.1-Flash" SECRET_PROMPT = "Remember this for later and reply with just OK." SECRET_REASONING = "The user told me their favorite color is chartreuse. I must remember it." diff --git a/tests/e2e_harness/codex/test_codex_rollout.py b/tests/e2e_harness/codex/test_codex_rollout.py index 2ef27a8a9c5..3db39066204 100644 --- a/tests/e2e_harness/codex/test_codex_rollout.py +++ b/tests/e2e_harness/codex/test_codex_rollout.py @@ -8,7 +8,14 @@ from typing import Final from cli_harness.rate_limiter import RateBucket, RateLimiter from codex.harness import CodexHarness -from codex.rollout import Compaction, rollout_compactions +from codex.rollout import ( + Compaction, + Subagent, + rollout_compactions, + rollout_final_messages, + rollout_parent_thread, + rollout_subagent, +) THREAD_ID: Final = "01a18fea-b2e8-7ed2-a29b-1ce4d2b138c3" USER_PROMPT: Final = { @@ -60,13 +67,83 @@ class TestRolloutCompactions: assert rollout_compactions(_rollout()) == () -def test_the_harness_reads_only_the_rollout_of_the_thread_it_is_asked_about(tmp_path: Path) -> None: - codex: Final = CodexHarness( - home=tmp_path, +def _harness(home: Path) -> CodexHarness: + return CodexHarness( + home=home, api_key="unused", rate_bucket=RateBucket.OPENAI, - limiter=RateLimiter({}, tmp_path, time.monotonic, time.sleep), + limiter=RateLimiter({}, home, time.monotonic, time.sleep), ) + + +def _subagent_rollout( + *, thread_id: str, parent: str | None, final_messages: tuple[str | None, ...], forked: bool = False +) -> str: + meta: Final = ( + {"id": thread_id, "cli_version": "0.165.0"} + | ({"parent_thread_id": parent} if parent else {}) + | ({"forked_from_id": parent} if forked else {}) + ) + inherited: Final = (_line("session_meta", {"id": parent, "cli_version": "0.165.0"}),) if forked else () + return "\n".join( + ( + _line("session_meta", meta), + *inherited, + _line("event_msg", {"type": "task_started", "turn_id": "t1"}), + _line( + "response_item", + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "thinking"}]}, + ), + *( + _line("event_msg", {"type": "task_complete", "turn_id": "t1", "last_agent_message": message}) + for message in final_messages + ), + ) + ) + + +class TestRolloutSubagents: + def test_a_subagent_rollout_names_the_thread_that_spawned_it(self) -> None: + rollout: Final = _subagent_rollout(thread_id="child", parent=THREAD_ID, final_messages=()) + assert rollout_parent_thread(rollout) == THREAD_ID + + def test_a_top_level_rollout_has_no_parent(self) -> None: + assert rollout_parent_thread(_rollout()) is None + + def test_a_subagent_that_inherited_its_parents_conversation_is_forked(self) -> None: + rollout: Final = _subagent_rollout(thread_id="child", parent=THREAD_ID, final_messages=("hi",), forked=True) + assert rollout_subagent(rollout) == Subagent(forked=True, replies=("hi",)) + + def test_a_subagent_that_started_fresh_is_not_forked(self) -> None: + rollout: Final = _subagent_rollout(thread_id="child", parent=THREAD_ID, final_messages=("hi",)) + assert rollout_subagent(rollout) == Subagent(forked=False, replies=("hi",)) + + def test_the_final_messages_are_the_last_message_of_each_finished_turn(self) -> None: + rollout: Final = _subagent_rollout(thread_id="child", parent=THREAD_ID, final_messages=("first", None, "last")) + assert rollout_final_messages(rollout) == ("first", "last") + + +def test_the_harness_reads_only_the_subagents_of_the_thread_it_is_asked_about(tmp_path: Path) -> None: + codex: Final = _harness(tmp_path) + sessions: Final = codex.codex_home / "sessions" / "2026" / "11" / "06" + sessions.mkdir(parents=True) + rollouts: Final = { + THREAD_ID: _subagent_rollout(thread_id=THREAD_ID, parent=None, final_messages=("the parent's answer",)), + "child-a": _subagent_rollout(thread_id="child-a", parent=THREAD_ID, final_messages=("reply a",)), + "child-b": _subagent_rollout(thread_id="child-b", parent=THREAD_ID, final_messages=("reply b",), forked=True), + "stranger": _subagent_rollout(thread_id="stranger", parent="other-thread", final_messages=("not ours",)), + } + for thread_id, rollout in rollouts.items(): + _ = (sessions / f"rollout-2026-11-06T17-01-05-{thread_id}.jsonl").write_text(rollout) + + assert codex.subagents(THREAD_ID) == ( + Subagent(forked=False, replies=("reply a",)), + Subagent(forked=True, replies=("reply b",)), + ) + + +def test_the_harness_reads_only_the_rollout_of_the_thread_it_is_asked_about(tmp_path: Path) -> None: + codex: Final = _harness(tmp_path) sessions: Final = codex.codex_home / "sessions" / "2026" / "10" / "30" sessions.mkdir(parents=True) _ = (sessions / f"rollout-2026-10-30T20-03-32-{THREAD_ID}.jsonl").write_text(_rollout(REMOTE)) diff --git a/tests/integration/_support/database_relay.py b/tests/integration/_support/database_relay.py index 2e9f2ec4a41..0ac060dc867 100644 --- a/tests/integration/_support/database_relay.py +++ b/tests/integration/_support/database_relay.py @@ -6,7 +6,6 @@ from contextlib import contextmanager from typing import Final from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit - from tests.integration._support.async_server import LoopbackServer, close_writer from tests.integration._support.ports import allocate_port @@ -74,16 +73,36 @@ class DatabaseRelay: class HeldStatementRelay: - def __init__(self, upstream_host: str, upstream_port: int, trigger: bytes) -> None: + """Holds the first client statement carrying ``trigger`` until ``release``. With ``after``, only a + trigger statement that follows a statement carrying ``after`` is held, so a test can pin the hold to + one point of a longer sequence. Starts armed unless ``armed`` is false, then ``arm`` starts it.""" + + def __init__( + self, + upstream_host: str, + upstream_port: int, + trigger: bytes, + *, + after: bytes | None = None, + armed: bool = True, + ) -> None: self.port: Final = allocate_port() self._upstream_host: Final = upstream_host self._upstream_port: Final = upstream_port self._trigger: Final = trigger + self._after: Final = after self._server: Final = LoopbackServer(self.port, self._serve) self._loop: Final = self._server.loop self._released: Final = asyncio.Event() + self._armed: Final = threading.Event() + if armed: + self._armed.set() + self._after_seen: Final = threading.Event() self.held: Final = threading.Event() + def arm(self) -> None: + self._armed.set() + def release(self) -> None: if not self._loop.is_closed(): self._loop.call_soon_threadsafe(self._released.set) @@ -95,21 +114,25 @@ class HeldStatementRelay: self.release() self._server.stop() - def _holds(self, window: bytes) -> bool: - return not self.held.is_set() and self._trigger in window + def _holds(self, markers: frozenset[bytes]) -> bool: + if not self._armed.is_set() or self.held.is_set(): + return False + if self._after is not None and not self._after_seen.is_set(): + if self._after in markers: + self._after_seen.set() + return False + return self._trigger in markers async def _serve(self, client_reader: asyncio.StreamReader, client_writer: asyncio.StreamWriter) -> None: server_reader, server_writer = await asyncio.open_connection(self._upstream_host, self._upstream_port) async def forward(reader: asyncio.StreamReader, writer: asyncio.StreamWriter, inspect: bool) -> None: - tail = b"" # rebind-ok: carries the previous read's end so a trigger split across reads still matches + scanner: Final = StatementMarkerScanner((self._trigger,) + ((self._after,) if self._after else ())) try: while chunk := await reader.read(65536): - window: Final = tail + chunk - if inspect and self._holds(window): + if inspect and any(self._holds(markers) for markers in scanner.feed(chunk)): self.held.set() await self._released.wait() - tail = window[-(len(self._trigger) - 1) :] writer.write(chunk) await writer.drain() except (ConnectionError, asyncio.IncompleteReadError): @@ -123,6 +146,47 @@ class HeldStatementRelay: ) +class StatementMarkerScanner: + def __init__(self, markers: tuple[bytes, ...]) -> None: + self._markers: Final = markers + self._statements: Final[dict[bytes, frozenset[bytes]]] = {} + self._buffer = b"" + self._startup = True + + def feed(self, chunk: bytes) -> Generator[frozenset[bytes]]: + self._buffer += chunk + while len(self._buffer) >= (4 if self._startup else 5): + offset: Final = 0 if self._startup else 1 + length: Final = int.from_bytes(self._buffer[offset : offset + 4], "big") + if length < 4: + raise ValueError("Invalid PostgreSQL message length") + if len(self._buffer) < length + offset: + return + message: Final = self._buffer[: length + offset] + self._buffer = self._buffer[length + offset :] + if self._startup: + self._startup = message[4:8] == b"\x04\xd2\x16/" + continue + payload: Final = message[5:] + if message[:1] == b"P": + name, _, query = payload.partition(b"\x00") + matches: Final = self._matches(query.split(b"\x00", 1)[0]) + if name: + self._statements[name] = matches + yield matches + elif message[:1] == b"B": + fields: Final = payload.split(b"\x00", 2) + prepared: Final = self._statements.get(fields[1], frozenset()) if len(fields) > 1 else frozenset() + yield prepared | self._matches(payload) + elif message[:1] == b"Q": + yield self._matches(payload) + elif message[:1] == b"C" and payload[:1] == b"S": + self._statements.pop(payload[1:].split(b"\x00", 1)[0], None) + + def _matches(self, query: bytes) -> frozenset[bytes]: + return frozenset(marker for marker in self._markers if marker in query) + + class TriggerScanner: def __init__(self, trigger: bytes) -> None: self._trigger: Final = trigger @@ -295,10 +359,12 @@ def database_relay(database_url: str, trigger: bytes) -> Generator[tuple[Databas @contextmanager -def held_statement_relay(database_url: str, trigger: bytes) -> Generator[tuple[HeldStatementRelay, str]]: +def held_statement_relay( + database_url: str, trigger: bytes, *, after: bytes | None = None, armed: bool = True +) -> Generator[tuple[HeldStatementRelay, str]]: parts: Final = urlsplit(database_url) assert parts.hostname is not None and parts.port is not None, database_url - relay: Final = HeldStatementRelay(parts.hostname, parts.port, trigger) + relay: Final = HeldStatementRelay(parts.hostname, parts.port, trigger, after=after, armed=armed) try: relay.start() yield relay, _relayed_url(database_url, relay.port) diff --git a/tests/integration/database/test_harness_lifecycle.py b/tests/integration/database/test_harness_lifecycle.py index d600a92c5a2..4cb01dfeef6 100644 --- a/tests/integration/database/test_harness_lifecycle.py +++ b/tests/integration/database/test_harness_lifecycle.py @@ -219,6 +219,49 @@ async def _unused_peer(reader: asyncio.StreamReader, writer: asyncio.StreamWrite await reader.read() +async def _echo_protocol(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + while chunk := await reader.read(65536): + writer.write(chunk) + await writer.drain() + + +def _protocol_message(kind: bytes, payload: bytes) -> bytes: + return kind + (len(payload) + 4).to_bytes(4, "big") + payload + + +@pytest.mark.parametrize("prepared_trigger", (False, True)) +def test_relay_holds_cached_statements_parsed_before_arming(prepared_trigger: bool) -> None: + upstream_port: Final = allocate_port() + upstream: Final = LoopbackServer(upstream_port, _echo_protocol) + relay: Final = HeldStatementRelay("127.0.0.1", upstream_port, b"trigger", after=b"snapshot", armed=False) + startup: Final = b"\x00\x00\x00\x08\x00\x03\x00\x00" + snapshot: Final = _protocol_message(b"P", b"s1\x00SELECT snapshot\x00\x00\x00") + trigger: Final = _protocol_message(b"P", b"s2\x00SELECT trigger\x00\x00\x00") + bind_snapshot: Final = _protocol_message(b"B", b"\x00s1\x00\x00\x00\x00\x00\x00\x00") + bind_trigger: Final = _protocol_message(b"B", b"\x00s2\x00\x00\x00\x00\x00\x00\x00") + upstream.start() + try: + relay.start() + with ( + socket.create_connection(("127.0.0.1", relay.port), timeout=5) as client, + client.makefile("rb") as replies, + ): + before: Final = startup + snapshot + bind_snapshot + (trigger if prepared_trigger else b"") + client.sendall(before) + assert replies.read(len(before)) == before + relay.arm() + client.sendall(bind_snapshot) + assert replies.read(len(bind_snapshot)) == bind_snapshot + held: Final = bind_trigger if prepared_trigger else _protocol_message(b"Q", b"SELECT trigger\x00") + client.sendall(held) + assert relay.held.wait(timeout=5), "cached snapshot did not open the relay barrier" + relay.release() + assert replies.read(len(held)) == held + finally: + relay.stop() + upstream.stop() + + def test_relay_bind_failure_closes_its_loop_and_allows_repeated_cleanup() -> None: port: Final = allocate_port() with socket.socket() as occupied: diff --git a/tests/integration/database/test_spendlogs_delivery.py b/tests/integration/database/test_spendlogs_delivery.py index d1227e85565..2b0be742cd4 100644 --- a/tests/integration/database/test_spendlogs_delivery.py +++ b/tests/integration/database/test_spendlogs_delivery.py @@ -373,15 +373,45 @@ async def test_receipt_cleanup_keeps_pending_and_parked_rows_and_reports_lag( ) repository = AccountingDeliveryRepository(delivery_client) assert await repository.cleanup_applied() == 1 - metrics = await repository.metrics() - assert metrics.pending_count == 1 - assert metrics.parked_count == 1 - assert metrics.oldest_pending_seconds >= 40 * 86400 + backlog = await repository.backlog() + assert backlog.pending_count == 1 + assert backlog.parked_count == 1 + assert backlog.oldest_pending_seconds >= 40 * 86400 assert delivery_database.execute( 'SELECT request_id FROM "LiteLLM_SpendLogDelivery" ORDER BY request_id' ).fetchall() == [("parked",), ("pending",)] +@pytest.mark.asyncio +async def test_backlog_reads_only_unapplied_deliveries_however_many_receipts_are_retained( + delivery_database: psycopg.Connection[tuple[object, ...]], delivery_client: Prisma +) -> None: + from litellm.repositories.accounting_delivery_repository import _BACKLOG + + delivery_database.execute( + 'INSERT INTO "LiteLLM_SpendLogDelivery" (delivery_id, request_id, accounting_payload, is_batch, state, applied_at) ' + "SELECT 'applied-' || n, 'applied-' || n, jsonb_build_object('payload', repeat(md5(n::text), 32)), " + "n % 10 = 0, 'applied', CURRENT_TIMESTAMP FROM generate_series(1, 20000) AS n" + ) + _insert(delivery_database, tuple(_row(state) for state in ("pending", "parked"))) + delivery_database.execute( + 'UPDATE "LiteLLM_SpendLogDelivery" SET state = request_id::"AccountingDeliveryState" ' + "WHERE request_id IN ('pending', 'parked')" + ) + delivery_database.execute('ANALYZE "LiteLLM_SpendLogDelivery"') + explained = delivery_database.execute(f"EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON) {_BACKLOG}").fetchone() + assert explained is not None + plan = explained[0][0]["Plan"] + table_pages = delivery_database.execute( + "SELECT relpages FROM pg_class WHERE oid = '\"LiteLLM_SpendLogDelivery\"'::regclass" + ).fetchone() + assert table_pages is not None and table_pages[0] > 1000 + assert plan["Shared Hit Blocks"] + plan["Shared Read Blocks"] < table_pages[0] / 20 + + backlog = await AccountingDeliveryRepository(delivery_client).backlog() + assert (backlog.pending_count, backlog.parked_count) == (1, 1) + + @pytest.mark.parametrize("delivery_database", ["public", "ledger_tenant"], indirect=True) @pytest.mark.parametrize("state", [None, "pending", "parked", "applied"]) def test_delivery_migration_replay_preserves_ledger_receipts(delivery_database, state): @@ -467,9 +497,9 @@ async def test_retention_keeps_batch_receipts_indefinitely(delivery_database, de assert await repository.cleanup_applied() == 1 assert await repository.cleanup_applied() == 0 assert delivery_database.execute('SELECT request_id FROM "LiteLLM_SpendLogDelivery"').fetchall() == [("batch",)] - metrics = await repository.metrics() - assert metrics.batch_count == 1 - assert metrics.table_bytes > 0 + storage = await repository.storage() + assert storage.batch_count == 1 + assert storage.table_bytes > 0 @pytest.mark.parametrize("disabled", [False, True]) diff --git a/tests/integration/database/test_spendlogs_ledger_parity.py b/tests/integration/database/test_spendlogs_ledger_parity.py index 60f3c074cca..ef87712b788 100644 --- a/tests/integration/database/test_spendlogs_ledger_parity.py +++ b/tests/integration/database/test_spendlogs_ledger_parity.py @@ -127,6 +127,66 @@ async def test_worker_removes_more_expired_receipts_than_one_cleanup_batch_in_on await client.db.disconnect() +@pytest.mark.asyncio +async def test_worker_checks_storage_and_repeats_delay_alerts_once_per_interval( + parity_database: str, monkeypatch: pytest.MonkeyPatch +): + from litellm.accounting.worker import flush_accounting_deliveries + + class _Alerts: + def __init__(self) -> None: + self.messages: list[str] = [] # mutable-ok: alert sink the worker appends to + + async def alerting_handler(self, message: str, level: str, alert_type: object) -> None: + self.messages.append(message) + + class _Clock: + def __init__(self) -> None: + self.now = 0.0 + + def __call__(self) -> float: + return self.now + + monkeypatch.setenv("DATABASE_URL", parity_database) + monkeypatch.setenv("LITELLM_ACCOUNTING_DELIVERY_ALERT_BYTES", "1") + with psycopg.connect(parity_database, autocommit=True) as connection: + connection.execute( + 'INSERT INTO "LiteLLM_SpendLogDelivery" ' + "(delivery_id, request_id, accounting_payload, is_batch, attempts, created_at, next_attempt_at) " + "VALUES ('stuck', 'stuck', '{}'::jsonb, false, 1, CURRENT_TIMESTAMP - INTERVAL '1 hour', " + "CURRENT_TIMESTAMP + INTERVAL '1 hour')" + ) + client: Final = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + clock: Final = _Clock() + writer: Final = DBSpendUpdateWriter(accounting_clock=clock) + alerts: Final = _Alerts() + monkeypatch.setenv("LITELLM_ACCOUNTING_DELIVERY_CHECK_INTERVAL_SECONDS", "100") + + async def flush_at(moment: float) -> tuple[str, ...]: + clock.now = moment + before: Final = len(alerts.messages) + await flush_accounting_deliveries(writer, client, alerts) + return tuple(message.split(",")[0] for message in alerts.messages[before:]) + + await client.db.connect() + try: + first: Final = await flush_at(0.0) + with psycopg.connect(parity_database, autocommit=True) as connection: + connection.execute( + 'INSERT INTO "LiteLLM_SpendLogDelivery" (delivery_id, request_id, accounting_payload, is_batch, state) ' + "VALUES ('batch', 'batch', '{}'::jsonb, true, 'applied')" + ) + within: Final = (await flush_at(10.0), await flush_at(99.0)) + after: Final = await flush_at(100.0) + finally: + await client.db.disconnect() + + delayed: Final = "Accounting deliveries are delayed. Pending=1" + assert first == ("Accounting delivery storage requires review. Retained batches=0", delayed) + assert within == ((), ()) + assert after == ("Accounting delivery storage requires review. Retained batches=1", delayed) + + @pytest.mark.parametrize("ledger", [False, True]) @pytest.mark.parametrize("exclude_messages", [False, True]) @pytest.mark.asyncio diff --git a/tests/integration/management/test_credential_writes_survive_reload.py b/tests/integration/management/test_credential_writes_survive_reload.py new file mode 100644 index 00000000000..9fb7d7eb681 --- /dev/null +++ b/tests/integration/management/test_credential_writes_survive_reload.py @@ -0,0 +1,119 @@ +import time +import uuid +from concurrent.futures import ThreadPoolExecutor, wait +from pathlib import Path +from typing import Final, Literal + +import pytest +import yaml +from pydantic import JsonValue + +from tests.integration._support.client import JSON_OBJECT, Gateway +from tests.integration._support.database import read_rows, scratch_database +from tests.integration._support.database_relay import held_statement_relay, plaintext_database_url +from tests.integration._support.process import owned_proxy + +RELOAD_INTERVAL_SECONDS: Final = 6 +CREDENTIAL_SNAPSHOT: Final = b'FROM "public"."LiteLLM_CredentialsTable"' +CONFIG_ROW_READ: Final = b"environment_variables" +CredentialWrite = Literal["create", "update", "delete"] + + +def _reloading_config(directory: Path) -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["general_settings"] = { + **config.get("general_settings", {}), + "store_model_in_db": True, + "proxy_config_reload_interval_seconds": RELOAD_INTERVAL_SECONDS, + } + config["litellm_settings"] = { + name: value + for name, value in config.get("litellm_settings", {}).items() + if name not in ("cache", "cache_params", "enable_redis_auth_cache") + } + path: Final = directory / "credential-reload.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +def _credential_body(name: str, display_name: str) -> dict[str, JsonValue]: + return { + "credential_name": name, + "display_name": display_name, + "credential_values": {"api_key": "synthetic-credential"}, + "credential_info": {"custom_llm_provider": "openai"}, + } + + +def _write_credential(proxy: Gateway, write: CredentialWrite, name: str) -> int: + match write: + case "create": + return proxy.request("POST", "/credentials", _credential_body(name, "created")).status_code + case "update": + return proxy.request( + "PATCH", f"/credentials/{name}", {"display_name": "updated", "credential_info": {}} + ).status_code + case "delete": + return proxy.request("DELETE", f"/credentials/{name}").status_code + + +def _served_display_name(proxy: Gateway, name: str) -> JsonValue: + response: Final = proxy.request("GET", f"/credentials/by_name/{name}") + if response.status_code == 404: + return None + assert response.status_code == 200, response.text + return JSON_OBJECT.validate_json(response.content)["display_name"] + + +def _stored_display_name(name: str, database_url: str) -> JsonValue: + rows: Final = read_rows( + 'SELECT display_name FROM "LiteLLM_CredentialsTable" WHERE credential_name = %s', + (name,), + database_url=database_url, + ) + return rows[0]["display_name"] if rows else None + + +def _served_display_names_until_the_next_reload(proxy: Gateway, name: str) -> frozenset[JsonValue]: + deadline: Final = time.monotonic() + RELOAD_INTERVAL_SECONDS / 3 + observed: Final = tuple(_served_display_name(proxy, name) for _ in iter(lambda: time.monotonic() < deadline, False)) + return frozenset(observed) + + +@pytest.mark.timeout(300) +@pytest.mark.parametrize("write", ("create", "update", "delete")) +def test_credential_write_is_not_undone_by_an_in_flight_config_reload( + gateway: Gateway, tmp_path: Path, write: CredentialWrite +) -> None: + name: Final = f"credential-{uuid.uuid4().hex}" + expected: Final = {"create": "created", "update": "updated", "delete": None}[write] + with ( + scratch_database() as database_url, + held_statement_relay(database_url, CONFIG_ROW_READ, after=CREDENTIAL_SNAPSHOT, armed=False) as ( + relay, + relayed_url, + ), + ): + overrides: Final = { + "DATABASE_URL": plaintext_database_url(relayed_url), + "LITELLM_CONFIG_PARAM_CACHE_TTL_SECONDS": "0", + } + with owned_proxy( + gateway, + tmp_path, + overrides, + config=_reloading_config(tmp_path), + remove_environment=("REDIS_HOST", "REDIS_PORT", "REDIS_URL"), + workers=1, + ) as proxy: + if write != "create": + proxy.post("/credentials", _credential_body(name, "original")) + relay.arm() + assert relay.held.wait(timeout=RELOAD_INTERVAL_SECONDS * 3), "no config reload reached the database" + with ThreadPoolExecutor(max_workers=1) as pool: + writing: Final = pool.submit(_write_credential, proxy, write, name) + _ = wait((writing,), timeout=2) + relay.release() + assert writing.result(timeout=30) == 200 + assert _stored_display_name(name, database_url) == expected + assert _served_display_names_until_the_next_reload(proxy, name) == frozenset({expected}) diff --git a/tests/integration/observability/test_azure_content_safety_dispatch_resilience.py b/tests/integration/observability/test_azure_content_safety_dispatch_resilience.py index 14eb687f69f..d184112f1ec 100644 --- a/tests/integration/observability/test_azure_content_safety_dispatch_resilience.py +++ b/tests/integration/observability/test_azure_content_safety_dispatch_resilience.py @@ -8,11 +8,13 @@ from typing import Final import httpx import pytest +from integration._support.budget import readiness_seconds from integration._support.client import Gateway, Scenario, eventually from integration._support.database import read_rows from integration._support.redis_process import owned_redis from integration._support.wire import Wire, wire_server from integration._support.ports import allocate_port +from integration._support.process import graceful_stop_seconds from integration.observability.azure_dispatch_support import ( ATTACK_MARKER, AzureBehavior, @@ -406,7 +408,7 @@ def test_x3_killing_one_worker_leaves_the_other_serving( _assert_spend(response.headers["x-litellm-call-id"], response.text) -@pytest.mark.timeout(180) +@pytest.mark.timeout(2 * readiness_seconds(2) + 2 * graceful_stop_seconds() + 120) def test_x4_proxy_restart_preserves_completed_spend_rows( gateway: Gateway, tmp_path: Path ) -> None: diff --git a/tests/integration/observability/test_azure_storage_chaos.py b/tests/integration/observability/test_azure_storage_chaos.py index acc2eb1439d..e78f1a4a1db 100644 --- a/tests/integration/observability/test_azure_storage_chaos.py +++ b/tests/integration/observability/test_azure_storage_chaos.py @@ -13,8 +13,9 @@ from _azure_storage_support import ( collect_files, ) from _s3_v2_support import matched_ids, mixed_burst, surface_reply +from integration._support.budget import readiness_seconds from integration._support.client import Gateway, JsonValue, eventually -from integration._support.process import group_members, owned_proxy_process +from integration._support.process import graceful_stop_seconds, group_members, owned_proxy_process from integration._support.tls import server_context, write_self_signed_cert from integration._support.wire import wire_server @@ -175,7 +176,7 @@ def test_killing_one_worker_keeps_the_other_serving_and_uploading(gateway: Gatew ) -@pytest.mark.timeout(180) +@pytest.mark.timeout(2 * readiness_seconds(WORKERS) + 2 * graceful_stop_seconds() + 120) def test_restarting_the_proxy_before_the_queue_flushes_bounds_the_loss_to_the_unflushed_queue_and_recovers( gateway: Gateway, tmp_path: Path ) -> None: diff --git a/tests/integration/observability/test_otel_auto_router_model.py b/tests/integration/observability/test_otel_auto_router_model.py new file mode 100644 index 00000000000..d056112daca --- /dev/null +++ b/tests/integration/observability/test_otel_auto_router_model.py @@ -0,0 +1,180 @@ +import json +import re +import uuid +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import pytest +import yaml +from integration._support.client import Gateway, eventually, gateway_from_environment +from integration._support.otlp_sink import SpanSinks, recorded_spans +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(300) + +MARKER: Final = re.compile(rb"autorouter-[0-9a-f]{32}") +TIERS: Final = ("SIMPLE", "MEDIUM", "COMPLEX", "REASONING") +USAGE: Final = {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6} + + +@dataclass(frozen=True, slots=True) +class AutoRouterRig: + proxy: Gateway + sink: str + router: str + alias: str + routed: str + + +def _upstream(request: Request) -> Reply: + found: Final = MARKER.search(request.body) + if found is None: + return Reply(body=json.dumps({"object": "list", "data": []}).encode()) + if request.target.endswith("/responses"): + output: Final = { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "ok", "annotations": []}], + } + return Reply( + body=json.dumps( + { + "id": "resp_" + found.group(0).decode(), + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini-2024-07-18", + "output": [output], + "usage": {"input_tokens": 5, "output_tokens": 1, "total_tokens": 6}, + } + ).encode() + ) + head: Final = {"id": "chatcmpl-" + found.group(0).decode(), "created": 1, "model": "gpt-4o-mini-2024-07-18"} + if json.loads(request.body).get("stream") is not True: + message: Final = {"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"} + return Reply( + body=json.dumps({**head, "object": "chat.completion", "choices": [message], "usage": USAGE}).encode() + ) + chunks: Final = ( + {**head, "object": "chat.completion.chunk", "choices": [{"index": 0, "delta": {"content": "ok"}}]}, + { + **head, + "object": "chat.completion.chunk", + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + "usage": USAGE, + }, + ) + return Reply( + content_type="text/event-stream", + chunks=(*(f"data: {json.dumps(chunk)}\n\n".encode() for chunk in chunks), b"data: [DONE]\n\n"), + ) + + +@pytest.fixture(scope="module") +def gateway(audit_sinks: SpanSinks) -> Iterator[Gateway]: + with gateway_from_environment() as base: + yield base + + +@pytest.fixture(scope="module", params=[False, True], ids=["otel-v1", "otel-v2"]) +def rig( + request: pytest.FixtureRequest, + gateway: Gateway, + audit_sinks: SpanSinks, + tmp_path_factory: pytest.TempPathFactory, +) -> Iterator[AutoRouterRig]: + otel_v2: Final = request.param is True + router: Final = "auto-router-" + uuid.uuid4().hex[:8] + alias: Final = "smart-" + uuid.uuid4().hex[:8] + routed: Final = "tier-model-" + uuid.uuid4().hex[:8] + directory: Final = tmp_path_factory.mktemp("auto-router-otel") + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + with wire_server(_upstream) as provider: + config["litellm_settings"] = {**config.get("litellm_settings", {}), "callbacks": ["otel"]} + config["router_settings"] = {**config.get("router_settings", {}), "model_group_alias": {alias: router}} + config["callback_settings"] = { + "otel": {"exporter": "http/json", "endpoint": audit_sinks.operator, "use_simple_processor": True} + } + config["model_list"] = [ + { + "model_name": routed, + "litellm_params": {"model": "openai/gpt-4o-mini", "api_base": provider.url, "api_key": "synthetic"}, + }, + { + "model_name": router, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": {"tiers": dict.fromkeys(TIERS, routed)}, + }, + }, + ] + path: Final = directory / "auto_router_otel.yaml" + path.write_text(yaml.safe_dump(config)) + overrides: Final = { + "OTEL_EXPORTER": "http/json", + "OTEL_ENDPOINT": audit_sinks.operator, + **({"LITELLM_OTEL_V2": "1"} if otel_v2 else {}), + } + with owned_proxy( + gateway, + directory, + overrides, + config=path, + remove_environment=() if otel_v2 else ("LITELLM_OTEL_V2",), + ) as proxy: + yield AutoRouterRig(proxy=proxy, sink=audit_sinks.operator, router=router, alias=alias, routed=routed) + + +@pytest.mark.parametrize( + ("route", "stream", "via_alias"), + [ + ("/v1/chat/completions", False, False), + ("/v1/chat/completions", True, False), + ("/v1/messages", False, False), + ("/v1/responses", False, False), + ("/v1/chat/completions", False, True), + ], + ids=["chat", "chat-stream", "messages", "responses", "chat-via-alias"], +) +def test_auto_router_llm_span_names_the_routed_model_and_carries_the_decision( + rig: AutoRouterRig, route: str, stream: bool, via_alias: bool +) -> None: + prompt: Final = "hi autorouter-" + uuid.uuid4().hex + body: Final[dict[str, JsonValue]] = { + "model": rig.alias if via_alias else rig.router, + "stream": stream, + **({"input": prompt} if route == "/v1/responses" else {"messages": [{"role": "user", "content": prompt}]}), + **({"max_tokens": 16} if route == "/v1/messages" else {}), + } + start, _ = recorded_spans(rig.sink) + response: Final = rig.proxy.request("POST", route, body) + assert response.status_code == 200, response.text + call_id: Final = response.headers["x-litellm-call-id"] + + def llm_span() -> Mapping[str, JsonValue] | None: + return next( + ( + span["attributes"] + for span in recorded_spans(rig.sink, start)[1] + if span["attributes"].get("litellm.call_id") == call_id + and "litellm.routing.router_model_name" in span["attributes"] + ), + None, + ) + + attributes: Final = eventually(llm_span, lambda found: found is not None, seconds=60) + assert attributes is not None + shown: Final = json.dumps(attributes, sort_keys=True, default=str) + underlying: Final = frozenset({rig.routed, "gpt-4o-mini"}) + assert attributes["gen_ai.request.model"] in underlying, shown + assert attributes["gen_ai.response.model"] in underlying, shown + assert attributes["litellm.routing.router_model_name"] == rig.router, shown + assert attributes["litellm.routing.router_type"] == "complexity", shown + assert attributes["litellm.routing.routed_model"] == rig.routed, shown + assert attributes["litellm.routing.tier"] in TIERS, shown diff --git a/tests/integration/observability/test_presidio_masking_endpoints.py b/tests/integration/observability/test_presidio_masking_endpoints.py index 33a436be8d5..ee1775eeae0 100644 --- a/tests/integration/observability/test_presidio_masking_endpoints.py +++ b/tests/integration/observability/test_presidio_masking_endpoints.py @@ -16,6 +16,7 @@ from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter CARD: Final = "4111 1111 1111 1111" +CARD_LEAK_PATTERN: Final = re.compile(r"(?" MASKED_EMAIL: Final = "" @@ -532,7 +533,7 @@ def test_default_scope_masks_a_card_number_the_model_generates(rig: Rig, model: answer: Final = _ask(rig, model, surface, rig.guardrails.default_scope, _card_prompt(_marker())) assert answer.text == MASKED_CARD, answer.wire - assert "4111" not in answer.wire, answer.wire + assert CARD_LEAK_PATTERN.search(answer.wire) is None, answer.wire @pytest.mark.parametrize( @@ -547,7 +548,7 @@ def test_spend_log_stores_the_masked_output_the_caller_received(rig: Rig, model: stored: Final = _spend_row(rig, model)["response"] assert "\n".join(_content_strings(stored)) == MASKED_CARD, stored - assert "4111" not in json.dumps(stored), stored + assert CARD_LEAK_PATTERN.search(json.dumps(stored)) is None, stored def test_masking_run_is_recorded_on_the_spend_log_with_the_detected_entities(rig: Rig, model: str) -> None: diff --git a/tests/integration/providers/test_bedrock_gpt_responses_native_wire.py b/tests/integration/providers/test_bedrock_gpt_responses_native_wire.py index 159097c80ee..a5cba28a48a 100644 --- a/tests/integration/providers/test_bedrock_gpt_responses_native_wire.py +++ b/tests/integration/providers/test_bedrock_gpt_responses_native_wire.py @@ -1,14 +1,15 @@ import base64 +import json import uuid from dataclasses import dataclass from typing import Final import openai -from integration._support.bedrock_runtime_peer import NATIVE_RESPONSES, answer, respond, target_of +from integration._support.bedrock_runtime_peer import NATIVE_RESPONSES, answer, marker_of, respond, target_of from integration._support.client import Gateway, Scenario, eventually from integration._support.database import read_rows -from integration._support.wire import Request, Wire, wire_server -from openai.types.responses import ResponseCompletedEvent, ResponseTextDeltaEvent +from integration._support.wire import Reply, Request, Wire, wire_server +from openai.types.responses import ResponseCompletedEvent, ResponseFailedEvent, ResponseTextDeltaEvent from pydantic import JsonValue, TypeAdapter from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with @@ -17,6 +18,7 @@ GPT: Final = "us.openai.gpt-5.6-sol" TOKEN: Final = "synthetic-bedrock-bearer" SALT: Final = "sk-integration-salt" _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +UPSTREAM_OUTAGE: Final = "The server had an error while processing your request. Sorry about that!" @dataclass(frozen=True, slots=True) @@ -123,6 +125,59 @@ async def test_async_openai_sdk_responses_stream_is_served_by_the_native_respons assert _spend_row(completed.response.id) == _success_row(model) +def _stream_then_bare_error(request: Request) -> Reply: + marker: Final = marker_of(request) + created: Final[dict[str, JsonValue]] = { + "type": "response.created", + "sequence_number": 0, + "response": { + "id": f"resp_upstream_{marker}", + "object": "response", + "created_at": 1, + "status": "in_progress", + "model": GPT, + "output": [], + }, + } + delta: Final[dict[str, JsonValue]] = { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": f"msg_{marker}", + "output_index": 0, + "content_index": 0, + "delta": answer(marker), + } + outage: Final[dict[str, JsonValue]] = { + "error": {"message": UPSTREAM_OUTAGE, "type": "server_error", "code": "service_unavailable"} + } + return Reply( + content_type="text/event-stream", + chunks=tuple(f"data: {json.dumps(event)}\n\n".encode() for event in (created, delta, outage)), + ) + + +async def test_mid_stream_upstream_error_without_an_event_type_reaches_the_client_as_response_failed( + gateway: Gateway, +) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_stream_then_bare_error) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire) + client: Final = openai.AsyncOpenAI(base_url=f"{gateway.client.base_url}/v1", api_key=gateway.key, max_retries=0) + stream: Final = await client.responses.create( + model=model, input=_prompt(marker), stream=True, extra_body={"cache": {"no-cache": True}} + ) + events: Final = [event async for event in stream] + assert [event.type for event in events] == [ + "response.created", + "response.output_text.delta", + "response.failed", + ], events + failed: Final = events[-1] + assert isinstance(failed, ResponseFailedEvent), failed + assert failed.response.error is not None and UPSTREAM_OUTAGE in failed.response.error.message, failed + _native_request(wire) + + def test_codex_agent_message_with_encrypted_content_reaches_the_native_responses_route_intact( gateway: Gateway, ) -> None: diff --git a/tests/integration/providers/test_hosted_tool_choice_wire.py b/tests/integration/providers/test_hosted_tool_choice_wire.py new file mode 100644 index 00000000000..8060b8f9fdd --- /dev/null +++ b/tests/integration/providers/test_hosted_tool_choice_wire.py @@ -0,0 +1,60 @@ +"""A chat completions ``tool_choice`` that forces an OpenAI hosted tool, such as ``{"type": "web_search"}``, reaches the +Responses API unchanged, and is a 400 before any provider call where the provider cannot enforce it.""" + +from typing import Final + +import pytest +from integration._support.client import JSON_OBJECT, Gateway, list_value +from integration._support.openai_wire import responses_reply +from integration._support.provider import SharedProvider +from pydantic import JsonValue + +_RESPONSES_MODEL: Final = "openai/responses/gpt-5.6-sol" +_ANTHROPIC_MODEL: Final = "anthropic/claude-haiku-4-5" +_WEATHER_TOOL: Final[dict[str, JsonValue]] = { + "type": "function", + "function": {"name": "get_weather", "parameters": {"type": "object", "properties": {}}}, +} + + +def _choice_id(tool_choice: dict[str, JsonValue]) -> str: + return str(tool_choice["type"]) + + +def _chat_request(model: str, tool_choice: dict[str, JsonValue]) -> dict[str, JsonValue]: + return { + "model": model, + "messages": [{"role": "user", "content": "What is the weather in Paris right now?"}], + "tools": [_WEATHER_TOOL, {"type": tool_choice["type"]}], + "tool_choice": tool_choice, + } + + +@pytest.mark.parametrize( + "tool_choice", + [{"type": "web_search"}, {"type": "web_search_preview"}, {"type": "file_search"}], + ids=_choice_id, +) +def test_hosted_tool_choice_reaches_the_responses_api_unchanged( + gateway: Gateway, provider: SharedProvider, tool_choice: dict[str, JsonValue] +) -> None: + provider.expect(responses_reply("resp_hosted_tool_choice", "gpt-5.6-sol", "It is 18C in Paris.", stream=False)) + response: Final = gateway.request("POST", "/v1/chat/completions", _chat_request(_RESPONSES_MODEL, tool_choice)) + assert response.status_code == 200, response.text + assert "It is 18C in Paris." in response.text, response.text + received: Final = provider.received() + assert [(request.method, request.target) for request in received] == [("POST", "/responses")] + sent: Final = JSON_OBJECT.validate_json(received[0].body) + assert sent["tool_choice"] == tool_choice + assert {"type": tool_choice["type"]} in list_value(sent["tools"]), sent["tools"] + + +def test_hosted_tool_choice_is_a_400_before_any_call_where_the_provider_cannot_enforce_it( + gateway: Gateway, provider: SharedProvider +) -> None: + response: Final = gateway.request( + "POST", "/v1/chat/completions", _chat_request(_ANTHROPIC_MODEL, {"type": "web_search"}) + ) + assert response.status_code == 400, response.text + assert "web_search" in response.text and "only supported for" in response.text, response.text + assert provider.received() == () diff --git a/tests/integration/spend/test_messages_stream_usage_cost.py b/tests/integration/spend/test_messages_stream_usage_cost.py index e28ab121de4..77ac95e2081 100644 --- a/tests/integration/spend/test_messages_stream_usage_cost.py +++ b/tests/integration/spend/test_messages_stream_usage_cost.py @@ -80,7 +80,9 @@ def _stream(message_id: str) -> bytes: ) -def _proxy_config(directory: Path, model: str, upstream_url: str) -> Path: +def _proxy_config( + directory: Path, model: str, upstream_url: str, litellm_settings: dict[str, JsonValue] | None = None +) -> Path: config: Final = directory / "streamed_usage_cost_config.yaml" config.write_text( json.dumps( @@ -109,7 +111,9 @@ def _proxy_config(directory: Path, model: str, upstream_url: str) -> Path: "proxy_batch_write_at": 1, "proxy_batch_polling_interval": 1, }, - "litellm_settings": {"include_cost_in_streaming_usage": True}, + "litellm_settings": {"include_cost_in_streaming_usage": True} + if litellm_settings is None + else litellm_settings, "router_settings": {"disable_cooldowns": True}, } ) @@ -176,3 +180,49 @@ def test_bedrock_messages_stream_usage_cost_matches_recorded_spend_with_custom_c "cost": pytest.approx(recorded_spend), }, (streamed_usage, rows, response.text) assert len(wire.drain()) == 1 + + +@pytest.mark.timeout(180) +def test_messages_stream_header_opts_into_usage_cost_while_global_flag_is_off(gateway: Gateway, tmp_path: Path) -> None: + model: Final = f"{BEDROCK_MODEL}-{uuid.uuid4().hex}" + message_ids: Final = {"plain": f"msg_{uuid.uuid4().hex}", "opted in": f"msg_{uuid.uuid4().hex}"} + + def respond(request: Request) -> Reply: + assert request.target == f"/model/{BEDROCK_MODEL}/invoke-with-response-stream", request.target + prompt: Final = json.loads(request.body)["messages"][0]["content"] + return Reply(content_type="application/vnd.amazon.eventstream", chunks=(_stream(message_ids[prompt]),)) + + def message_delta_usage(prompt: str, headers: dict[str, str]) -> dict[str, float]: + response: Final = candidate.request( + "POST", + "/v1/messages", + { + "model": model, + "messages": [{"role": "user", "content": prompt}], + "max_tokens": OUTPUT_TOKENS, + "stream": True, + }, + headers=headers, + ) + assert response.status_code == 200, response.text + message_delta: Final = next(event for event in _data_events(response.text) if event["type"] == "message_delta") + return STREAMED_USAGE.validate_python(message_delta["usage"]) + + with wire_server(respond) as wire: + config: Final = _proxy_config(tmp_path, model, wire.url, litellm_settings={}) + with owned_proxy(gateway, tmp_path, {}, config=config) as candidate: + plain_usage: Final = message_delta_usage("plain", {}) + opted_in_usage: Final = message_delta_usage( + "opted in", {"x-litellm-include-cost-in-streaming-usage": "true"} + ) + rows: Final = eventually( + lambda: read_rows( + 'SELECT spend FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (message_ids["opted in"],) + ), + lambda values: len(values) == 1, + seconds=70, + ) + assert "cost" not in plain_usage, plain_usage + assert opted_in_usage["cost"] == pytest.approx(float(str(rows[0]["spend"]))), (opted_in_usage, rows) + assert opted_in_usage["cost"] == pytest.approx(EXPECTED_SPEND), opted_in_usage + assert len(wire.drain()) == 2 diff --git a/tests/integration/spend/test_normalized_error_long_message.py b/tests/integration/spend/test_normalized_error_long_message.py index 23200c0450e..23769391427 100644 --- a/tests/integration/spend/test_normalized_error_long_message.py +++ b/tests/integration/spend/test_normalized_error_long_message.py @@ -358,16 +358,17 @@ def _burst( *, rejected_model: str = CRAFTED_MODEL, ) -> tuple[_BurstResult, ...]: + nonce: Final = uuid.uuid4().hex crafted: Final = tuple( (f"crafted-{path}-{index}", path, _body(path, rejected_model, uuid.uuid4().hex, stream=index % 2 == 1), open_key) for path in (CHAT, MESSAGES, RESPONSES) for index in range(4) ) happy: Final = tuple( - (f"happy-{index}", CHAT, _body(CHAT, happy_model, f"happy-{index}"), happy_key) for index in range(8) + (f"happy-{index}", CHAT, _body(CHAT, happy_model, f"happy-{index}-{nonce}"), happy_key) for index in range(8) ) late: Final = tuple( - (f"late-{index}", CHAT, _body(CHAT, happy_model, f"late-{index}"), happy_key) for index in range(8) + (f"late-{index}", CHAT, _body(CHAT, happy_model, f"late-{index}-{nonce}"), happy_key) for index in range(8) ) with ( ThreadPoolExecutor(max_workers=28) as pool, @@ -427,6 +428,7 @@ def test_killing_one_worker_mid_burst_leaves_the_other_serving_crafted_and_happy open_key: Final = scenario.key() control: Final = _burst(client, model, key, open_key, lambda: None, rejected_model=BENIGN_288KB_MODEL) assert all(result.status == (400 if result.label.startswith("crafted") else 200) for result in control), control + wire.drain() workers: Final = owned.worker_pids() assert len(workers) == 2, workers victim: Final = min(workers) @@ -438,6 +440,8 @@ def test_killing_one_worker_mid_burst_leaves_the_other_serving_crafted_and_happy _assert_burst_rejections_fast(results, control) late: Final = tuple(result for result in results if result.label.startswith("late")) assert all(result.status == 200 for result in late), late + seen: Final = tuple(request.body.decode() for request in wire.drain()) + assert all(any(f"normalized error audit {result.label}-" in body for body in seen) for result in late), seen _assert_rows_land_exactly_once(results, "late") survivor: Final = tuple(pid for pid in owned.worker_pids() if pid != victim) assert survivor, "no worker left serving" diff --git a/tests/integration/translation/messages/bases/azure.py b/tests/integration/translation/messages/bases/azure.py index b731f7b7760..cbe03515b0e 100644 --- a/tests/integration/translation/messages/bases/azure.py +++ b/tests/integration/translation/messages/bases/azure.py @@ -323,6 +323,24 @@ GPT_6_LUNA_CONTENT_FILTERS: Final[JsonValue] = [ }, ] +GPT_6_LUNA_ENCRYPTED_CONTENT: Final = ( + "gAAAAABq0CO7N77UmN8oIxFPRrP2rvE9LuqnrGGHsFzckseHLpNB24KIlg0OKnfemaCgYarrVksxeJBKgoiQ79xkF1NR9RI8-82w" + "4OGNB0ViGNYNA7nf5kOqOTrMh1_gHecIN0Evxn8I0meiqOds7h1biKfkiZlm19b-9r_Y2K--KXONVKFEbe8kP54_ka9KCxQ3_5Xn" + "L4IRfRS7lsLfoMjORKnapkVSnMjueRVVOprw_897JJypeT61mzzBYM_18T6k68OE70JndMsr4GNRq4c_Lf_rSRoHg7_iFrdHsBdv" + "MfN79dfp_PeLu5MZVRgcag-qM9aKUAvzhTD2-EutMw1C1ZNabU2S5Qa4RXq2OiU_iZljJI6JhpnJ9jT9g_M3W_0mklmFU2HBhQ6O" + "csIFmPePVcSaQlmofKgvq2ZThSHyFpLsT15ftBL2ZMZ3OWeodDnkfRr2fY5NDM7v7U1TLcAx7xCnP88W03kyb3SuJGSsbPnYTGbN" + "ZN7ye3SBJRkCMpRgUfguSYdC46mmMjPdvdkLnTWHZ23GOoDG5A2lffUJZ4WXNgPrMdG2FNThtHT0otHeDC75ahpa6i5xtqIsCGiF" + "RVH_iSxhItArxFXoc7e6oBU9YxsgkUkyXW_zz6h5h5E6DoJz4QckOAdL8E3SwL1nu6ovpS3_QgeiHa5RewKs6WclFCxeiwiyXO6b" + "fjhl_VtOA8KyWqkO-rt4T0RneV_bZHav6KoVanyCX-ZsAYj1WIhn4xEOXh6RdBjZJJB-yDw5cq6zdXkmN1hxs_nmA7hUemHkPclb" + "2lGF41N28IwMpcXq1rx2jrw2XleLnqFVKwgCLysinS9erU2vRN1CmmDQrGl7C5lIOrfSVCpB2W_p3BZWNDKwAfF4i8uXCxy1_5-E" + "jJ08CJXXYVchceOreHYNYxAXhVTHhJArQqVWQgTwJl9e_HhyA0OFZHojB54BjlkMr7QJMm9VE9LIkOJgRgIrtUM2ULIy1tZ-Ezuz" + "Ey9HYx5NuUKe_12ExKuUVrQmaJ2HRHvj_pPcmSk6lOVrN5VvVBMj_KHd3MfkbUWVeVfzIsgWc149TexK6FPtvHGPxf_VQBdyA4Jl" + "77dPp671wxH7puzy184LV1IMdxq5eY1JG4QdCOluIDkkqbw_qYEiKGNXImGnp9LsXCcYa6JivYmSAw_vgZ5YhO98L49aLEI6nT3H" + "WRWuzY3H39aCkFdA8Y0DttzwXYvhz8UuenFgoGvIk0W7WSOLBxi6pSefdczpRizWLMm10zlrYVFvdMzUPEx7OyDUhSXag61gYcsm" + "6L5mcXGb2cojThdHEXcFZg36bQYjQVJzu_ASlvN-Z9AFNajIIyw8qKx3Q8pRVg3HCpQcSyRiDhicvT_Gp8d0X7VhoAhG_-LPg4AO" + "z7U985JqP0wB2dB4sYhRgMZpQ29x7wgvgjNtsyqbCA==" +) + GPT_6_LUNA_TEST_CASE: Final = TranslationTestCase( scenario="basic", litellm_endpoint="/v1/messages", @@ -347,27 +365,7 @@ GPT_6_LUNA_TEST_CASE: Final = TranslationTestCase( "id": "rs_0534b0c3898d6d1b006ad023bb09b8819590793fe299bffdcf", "type": "reasoning", "content": [], - "encrypted_content": ( - "gAAAAABq0CO7N77UmN8oIxFPRrP2rvE9LuqnrGGHsFzckseHLpNB24KIlg0OKnfemaCgYarrVksxeJBK" - "goiQ79xkF1NR9RI8-82w4OGNB0ViGNYNA7nf5kOqOTrMh1_gHecIN0Evxn8I0meiqOds7h1biKfkiZlm" - "19b-9r_Y2K--KXONVKFEbe8kP54_ka9KCxQ3_5XnL4IRfRS7lsLfoMjORKnapkVSnMjueRVVOprw_897" - "JJypeT61mzzBYM_18T6k68OE70JndMsr4GNRq4c_Lf_rSRoHg7_iFrdHsBdvMfN79dfp_PeLu5MZVRgc" - "ag-qM9aKUAvzhTD2-EutMw1C1ZNabU2S5Qa4RXq2OiU_iZljJI6JhpnJ9jT9g_M3W_0mklmFU2HBhQ6O" - "csIFmPePVcSaQlmofKgvq2ZThSHyFpLsT15ftBL2ZMZ3OWeodDnkfRr2fY5NDM7v7U1TLcAx7xCnP88W" - "03kyb3SuJGSsbPnYTGbNZN7ye3SBJRkCMpRgUfguSYdC46mmMjPdvdkLnTWHZ23GOoDG5A2lffUJZ4WX" - "NgPrMdG2FNThtHT0otHeDC75ahpa6i5xtqIsCGiFRVH_iSxhItArxFXoc7e6oBU9YxsgkUkyXW_zz6h5" - "h5E6DoJz4QckOAdL8E3SwL1nu6ovpS3_QgeiHa5RewKs6WclFCxeiwiyXO6bfjhl_VtOA8KyWqkO-rt4" - "T0RneV_bZHav6KoVanyCX-ZsAYj1WIhn4xEOXh6RdBjZJJB-yDw5cq6zdXkmN1hxs_nmA7hUemHkPclb" - "2lGF41N28IwMpcXq1rx2jrw2XleLnqFVKwgCLysinS9erU2vRN1CmmDQrGl7C5lIOrfSVCpB2W_p3BZW" - "NDKwAfF4i8uXCxy1_5-EjJ08CJXXYVchceOreHYNYxAXhVTHhJArQqVWQgTwJl9e_HhyA0OFZHojB54B" - "jlkMr7QJMm9VE9LIkOJgRgIrtUM2ULIy1tZ-EzuzEy9HYx5NuUKe_12ExKuUVrQmaJ2HRHvj_pPcmSk6" - "lOVrN5VvVBMj_KHd3MfkbUWVeVfzIsgWc149TexK6FPtvHGPxf_VQBdyA4Jl77dPp671wxH7puzy184L" - "V1IMdxq5eY1JG4QdCOluIDkkqbw_qYEiKGNXImGnp9LsXCcYa6JivYmSAw_vgZ5YhO98L49aLEI6nT3H" - "WRWuzY3H39aCkFdA8Y0DttzwXYvhz8UuenFgoGvIk0W7WSOLBxi6pSefdczpRizWLMm10zlrYVFvdMzU" - "PEx7OyDUhSXag61gYcsm6L5mcXGb2cojThdHEXcFZg36bQYjQVJzu_ASlvN-Z9AFNajIIyw8qKx3Q8pR" - "Vg3HCpQcSyRiDhicvT_Gp8d0X7VhoAhG_-LPg4AOz7U985JqP0wB2dB4sYhRgMZpQ29x7wgvgjNtsyqb" - "CA==" - ), + "encrypted_content": GPT_6_LUNA_ENCRYPTED_CONTENT, "summary": [], }, { @@ -411,6 +409,10 @@ GPT_6_LUNA_TEST_CASE: Final = TranslationTestCase( "model": "azure/gpt-6-luna", "usage": {"input_tokens": 309, "output_tokens": 114}, "content": [ + { + "type": "redacted_thinking", + "data": "litellm_encrypted_reasoning:" + GPT_6_LUNA_ENCRYPTED_CONTENT, + }, { "type": "text", "text": ( @@ -457,6 +459,25 @@ GPT_6_1_SOL_CONTENT_FILTERS: Final[JsonValue] = [ }, ] +GPT_6_1_SOL_ENCRYPTED_CONTENT: Final = ( + "gAAAAABq0CO0QoINZ4oOVkvDvinFAJAMUBmzqqpXLoppdlyP4fFQsNqPOh-L08mQt8yf4OUtOwzk3E6Q7CLHieco0s_tfuYHRlTU" + "0Blr62KF_C77RScOjKYa6p4E2fjdwO44Y2HIiGfIo_GnUpkfwW1EiVZmj6L8PO0biPqihvK3KJ-GNKB_R4-lLyKXrrvVN-985glS" + "dHg-_fRrmyVGBvVxPFyERlD-WPm8fPujqMmXnUcAv3N-wUoIEccd7PfJGqBs52UzllpH27mNCI71C_qbc2nZ97jZgSabTvAZKF8Q" + "9zSJzWA_Gs72iQqd0X55F5vD0-xrMuVdXFTYLV-y-S_1rN8tY1bEbdJWr7Dl84NvNWCYqA5byZgsMxAVyjADWXcVaagrZhNtQ1ol" + "II2oePJL2WyrKrOs4d7ZuGFW5BcFCDS5zxk6MwckKbLkSi3HM34z-AhIOjr6SNc5qcR8yur4Phm9pSJQqFscqQwtmEgbAzXqC9U5" + "rwQ_c11N7o6Xc2bnoRXxeXQvKyGCPZWOqQr8AffHBCCOmeKjbPMV0ITek7HyRiOWg_Pjq-_xlVCoUGM7gIq5nq3HG1SqH7Yr6qev" + "5Arq24UALiBbr0NAtmq3p-YGRgia3oDnSrnygIX8AbqhnGrHz1kjBbwGr6fXF2XruSKuL96bQghfWGx6nhsbIDH4FI3akERb6Rh1" + "B9XIWsWw4Zzo6-tg6DnjzfaJteZqIY4qS5mHSD9stXEVqu_CfnGh7EdOlZlnkG6kIQ3m6oAo2ZG9E58m32X-KEqZDjD4yUgWss9y" + "7cnFbYeEzntQneytoBpaTNXtSEsM1fUoz73Ol7FzqPI0CKIsp3UGje2K7LqAwkV54JrhqCfxkD5HmPJnyEOnP0dODvyqD3CK-cbk" + "oPoszI2Ba0dk11re5kF985he-Fjm7dXDNhQX93__nj4awDCtnX8_Ke5ld7O3sA1ou7WHhgfGo-Jr4ZcGL5JWflPbI9oPpGOOvLsS" + "-e8CUjTcTRb_eA2jApeGU2nAdW42U-CnXeScaR5eJE1RoXTFXIP6MMSpxqpTRWYBfhIDyPEzwTveki-WOWgJtyNJ5W3GJxaEZla2" + "uJ5_cdiD4-HgRy1UwxIJODUh0F2hKC6niDVIZqoZ6Z4nb4OUvJP00_nGBA5-IL31Ij8I8U740RAZvnOsMOIsKelsJioM8Ii4MdPi" + "O2FYSdKjijEJmXLDdbYkQEXeqwfKGC7trXBOyRwO5AF6mLR9Yf8Bu_eiKbuf5DT9sw88P4JtfPi5LLIso2Ki44UnxePraO7vinqg" + "7H7oOi-1gw7BbDaKH7tJs3vUOaxl6mA0BNTg5cDt6TtFZMpFgKNRyanWzSwGWQGYlYbfUGVUnM_tjtgoBNFJAQz1s8fd-qOl1NQD" + "ggjlryl6TJk3RSiNYKyJ3fmpE10LM3LUtN3AoFMDFbEEBstDuVg0zzvwe64zaKosn7PIh_Ym-y7QlMdcvUbsqO-T6pfsq8iqEz2a" + "nv3yujHkAXIaa3292lbAdFIq80ycQSJbWNd2MtkYQHdWtWmQQoqMX1PdyJGp6N3GG45A4A==" +) + GPT_6_1_SOL_TEST_CASE: Final = TranslationTestCase( scenario="basic", litellm_endpoint="/v1/messages", @@ -480,28 +501,7 @@ GPT_6_1_SOL_TEST_CASE: Final = TranslationTestCase( "id": "rs_0a4a896afe94ce65006ad023b19f8481938d51de793f83feed", "type": "reasoning", "content": [], - "encrypted_content": ( - "gAAAAABq0CO0QoINZ4oOVkvDvinFAJAMUBmzqqpXLoppdlyP4fFQsNqPOh-L08mQt8yf4OUtOwzk3E6Q" - "7CLHieco0s_tfuYHRlTU0Blr62KF_C77RScOjKYa6p4E2fjdwO44Y2HIiGfIo_GnUpkfwW1EiVZmj6L8" - "PO0biPqihvK3KJ-GNKB_R4-lLyKXrrvVN-985glSdHg-_fRrmyVGBvVxPFyERlD-WPm8fPujqMmXnUcA" - "v3N-wUoIEccd7PfJGqBs52UzllpH27mNCI71C_qbc2nZ97jZgSabTvAZKF8Q9zSJzWA_Gs72iQqd0X55" - "F5vD0-xrMuVdXFTYLV-y-S_1rN8tY1bEbdJWr7Dl84NvNWCYqA5byZgsMxAVyjADWXcVaagrZhNtQ1ol" - "II2oePJL2WyrKrOs4d7ZuGFW5BcFCDS5zxk6MwckKbLkSi3HM34z-AhIOjr6SNc5qcR8yur4Phm9pSJQ" - "qFscqQwtmEgbAzXqC9U5rwQ_c11N7o6Xc2bnoRXxeXQvKyGCPZWOqQr8AffHBCCOmeKjbPMV0ITek7Hy" - "RiOWg_Pjq-_xlVCoUGM7gIq5nq3HG1SqH7Yr6qev5Arq24UALiBbr0NAtmq3p-YGRgia3oDnSrnygIX8" - "AbqhnGrHz1kjBbwGr6fXF2XruSKuL96bQghfWGx6nhsbIDH4FI3akERb6Rh1B9XIWsWw4Zzo6-tg6Dnj" - "zfaJteZqIY4qS5mHSD9stXEVqu_CfnGh7EdOlZlnkG6kIQ3m6oAo2ZG9E58m32X-KEqZDjD4yUgWss9y" - "7cnFbYeEzntQneytoBpaTNXtSEsM1fUoz73Ol7FzqPI0CKIsp3UGje2K7LqAwkV54JrhqCfxkD5HmPJn" - "yEOnP0dODvyqD3CK-cbkoPoszI2Ba0dk11re5kF985he-Fjm7dXDNhQX93__nj4awDCtnX8_Ke5ld7O3" - "sA1ou7WHhgfGo-Jr4ZcGL5JWflPbI9oPpGOOvLsS-e8CUjTcTRb_eA2jApeGU2nAdW42U-CnXeScaR5e" - "JE1RoXTFXIP6MMSpxqpTRWYBfhIDyPEzwTveki-WOWgJtyNJ5W3GJxaEZla2uJ5_cdiD4-HgRy1UwxIJ" - "ODUh0F2hKC6niDVIZqoZ6Z4nb4OUvJP00_nGBA5-IL31Ij8I8U740RAZvnOsMOIsKelsJioM8Ii4MdPi" - "O2FYSdKjijEJmXLDdbYkQEXeqwfKGC7trXBOyRwO5AF6mLR9Yf8Bu_eiKbuf5DT9sw88P4JtfPi5LLIs" - "o2Ki44UnxePraO7vinqg7H7oOi-1gw7BbDaKH7tJs3vUOaxl6mA0BNTg5cDt6TtFZMpFgKNRyanWzSwG" - "WQGYlYbfUGVUnM_tjtgoBNFJAQz1s8fd-qOl1NQDggjlryl6TJk3RSiNYKyJ3fmpE10LM3LUtN3AoFMD" - "FbEEBstDuVg0zzvwe64zaKosn7PIh_Ym-y7QlMdcvUbsqO-T6pfsq8iqEz2anv3yujHkAXIaa3292lbA" - "dFIq80ycQSJbWNd2MtkYQHdWtWmQQoqMX1PdyJGp6N3GG45A4A==" - ), + "encrypted_content": GPT_6_1_SOL_ENCRYPTED_CONTENT, "summary": [], }, { @@ -545,6 +545,10 @@ GPT_6_1_SOL_TEST_CASE: Final = TranslationTestCase( "model": "azure/gpt-6.1-sol", "usage": {"input_tokens": 309, "output_tokens": 137}, "content": [ + { + "type": "redacted_thinking", + "data": "litellm_encrypted_reasoning:" + GPT_6_1_SOL_ENCRYPTED_CONTENT, + }, { "type": "text", "text": ( @@ -592,6 +596,25 @@ GPT_6_ASTRA_CONTENT_FILTERS: Final[JsonValue] = [ }, ] +GPT_6_ASTRA_ENCRYPTED_CONTENT: Final = ( + "gAAAAABq0CPHdct5jJ3BolLC7QmXutcBsSZ0YzHiI8z2uK07jqwLifVhEROdF12ygdCVRtcrLrHPyHShpVbh41ULxn2L3SyVYLHb" + "bkFMcwbzfky_kqZdUe_1D-PLxsVyI1KFMPLg0Fb9K-N9u_ElvkFcXXsmAhRQCcB4SE9TMe_5tKBHiZbqvOI112wig7cc57rDbvM0" + "mEoKyFn3WgprP_ewHeCeP6DSEvVGDEMjUd_TIXAHH7y2Ae3y0WRW7WUag_wAiO7xzAH13mh4S6eKeBTKh62krmYeUH3wETO8BF1W" + "nkuZ_86lS0zWLBzeLJJnJiXKXVquKatJ6HvO3m6FkPF0tBXhdc8YYtTF1B_WBMECARKcw4n3mPL3R6J2jpHe05Eo7W7uZkcY5E-q" + "mE_cqUK8YQXTXx-mR339iuAvqSe0N2Rgl9Ey1l7z0MBATIV0c6J4wRav6Di5aAxEEXkgDVk1jF9LZSHomcGjh2B2UWFiAZEELRDD" + "7ovpP3c7B05h0QvwazXlZnskJf-xYZEHAl2CNZk0iqagFs6EtNrFguM5W2AH6dsuaGjNGmdWM_kt4JqIUeScZLNgMty8qTOHFzd7" + "tiBpuLkQgcnHzOigWb2vY_8ahy_zfzM2lVL-JQxP9SB7Gq1m49Ln-byC1O_kPyqEFXELhhQ8peXQY5do2R2T-3I0WLOhDhP5rruL" + "sL66vNHAAMjka6WmjI2goiy-JF_nVCplJd791Wqa3yf7Q7ksfasnSKtjR7-cgTtMcXX7n_5joMnqf4_LsqSq6gHdeKw4VTeT8wGB" + "Uo3MHLkgi8PnAIphdb3PKgifSjstE_oH4_ZbmlyCsHaFYq2aoy-S85ndlrR_PAHBtU8II-Ud_42EPzS-Etdt5d1Qjjri3KTyPjPO" + "7fAm8aPFr1k3GbjwASPMQ7cUs7yLhPmI-TCm5U-zgK1pIALviiU2o4IM7kdH0-Dg1eyyZN8hsFGSHaUB7G_tGB0fj0H0GSztH8lt" + "NjTK_276IOh8Jv5csxjjmlcMOHVBmDCkkE_DlVjvNQtWQS9QBt9jd1gFSzaOhXVkKZJK3vTIaq8JxTvYtgLc8WM5pRHbnQr71TfW" + "B-eLEYgN-aX5PueKkXGrGauDpKkzVJyNZ6mZ_dVXwl1p5IiNNJ1xzhmZK2GVDFWt0nfOBYyxmCDfIxvQL8YAhXXAi1Hrjal82c68" + "lw7UBauyMojDOcbH7EDG9KNbVV-kmkhCW3aH0QCIALzHqO_cukXDqnYX-iVNuxWpW5KUAJdg9cfgwX35O246HxDoSpoIYQ1GooqI" + "WSImYrlNQajPHpJcpxnWeSgKOSYIAYCAr70MouYhUg-lsId7n2pr4gspPf1ZUrBhAG3h4n9qSguKTlpnAYlln_2MzgVVN1cH5sop" + "vgJJLr20HKS1cfu9lB-CHDF5h74z7q3ydZJljn8dQLBqL3csr0ovOZkTjP21rOjPhdsedg4CLKnylgk5nTewbQz99lKxKPhfYK0-" + "XxiDdE6aKvobQUoRupGNHv-gMLqJjehMUynaC6XzBwJ63pPNhH3AlCj3zSte0Mj1FkPQQIY2PjhMB49CmXBV8Z9YJv0=" +) + GPT_6_ASTRA_TEST_CASE: Final = TranslationTestCase( scenario="basic", litellm_endpoint="/v1/messages", @@ -616,28 +639,7 @@ GPT_6_ASTRA_TEST_CASE: Final = TranslationTestCase( "id": "rs_019af305501e4822006ad023c49ce4819590da413b9879e06b", "type": "reasoning", "content": [], - "encrypted_content": ( - "gAAAAABq0CPHdct5jJ3BolLC7QmXutcBsSZ0YzHiI8z2uK07jqwLifVhEROdF12ygdCVRtcrLrHPyHSh" - "pVbh41ULxn2L3SyVYLHbbkFMcwbzfky_kqZdUe_1D-PLxsVyI1KFMPLg0Fb9K-N9u_ElvkFcXXsmAhRQ" - "CcB4SE9TMe_5tKBHiZbqvOI112wig7cc57rDbvM0mEoKyFn3WgprP_ewHeCeP6DSEvVGDEMjUd_TIXAH" - "H7y2Ae3y0WRW7WUag_wAiO7xzAH13mh4S6eKeBTKh62krmYeUH3wETO8BF1WnkuZ_86lS0zWLBzeLJJn" - "JiXKXVquKatJ6HvO3m6FkPF0tBXhdc8YYtTF1B_WBMECARKcw4n3mPL3R6J2jpHe05Eo7W7uZkcY5E-q" - "mE_cqUK8YQXTXx-mR339iuAvqSe0N2Rgl9Ey1l7z0MBATIV0c6J4wRav6Di5aAxEEXkgDVk1jF9LZSHo" - "mcGjh2B2UWFiAZEELRDD7ovpP3c7B05h0QvwazXlZnskJf-xYZEHAl2CNZk0iqagFs6EtNrFguM5W2AH" - "6dsuaGjNGmdWM_kt4JqIUeScZLNgMty8qTOHFzd7tiBpuLkQgcnHzOigWb2vY_8ahy_zfzM2lVL-JQxP" - "9SB7Gq1m49Ln-byC1O_kPyqEFXELhhQ8peXQY5do2R2T-3I0WLOhDhP5rruLsL66vNHAAMjka6WmjI2g" - "oiy-JF_nVCplJd791Wqa3yf7Q7ksfasnSKtjR7-cgTtMcXX7n_5joMnqf4_LsqSq6gHdeKw4VTeT8wGB" - "Uo3MHLkgi8PnAIphdb3PKgifSjstE_oH4_ZbmlyCsHaFYq2aoy-S85ndlrR_PAHBtU8II-Ud_42EPzS-" - "Etdt5d1Qjjri3KTyPjPO7fAm8aPFr1k3GbjwASPMQ7cUs7yLhPmI-TCm5U-zgK1pIALviiU2o4IM7kdH" - "0-Dg1eyyZN8hsFGSHaUB7G_tGB0fj0H0GSztH8ltNjTK_276IOh8Jv5csxjjmlcMOHVBmDCkkE_DlVjv" - "NQtWQS9QBt9jd1gFSzaOhXVkKZJK3vTIaq8JxTvYtgLc8WM5pRHbnQr71TfWB-eLEYgN-aX5PueKkXGr" - "GauDpKkzVJyNZ6mZ_dVXwl1p5IiNNJ1xzhmZK2GVDFWt0nfOBYyxmCDfIxvQL8YAhXXAi1Hrjal82c68" - "lw7UBauyMojDOcbH7EDG9KNbVV-kmkhCW3aH0QCIALzHqO_cukXDqnYX-iVNuxWpW5KUAJdg9cfgwX35" - "O246HxDoSpoIYQ1GooqIWSImYrlNQajPHpJcpxnWeSgKOSYIAYCAr70MouYhUg-lsId7n2pr4gspPf1Z" - "UrBhAG3h4n9qSguKTlpnAYlln_2MzgVVN1cH5sopvgJJLr20HKS1cfu9lB-CHDF5h74z7q3ydZJljn8d" - "QLBqL3csr0ovOZkTjP21rOjPhdsedg4CLKnylgk5nTewbQz99lKxKPhfYK0-XxiDdE6aKvobQUoRupGN" - "Hv-gMLqJjehMUynaC6XzBwJ63pPNhH3AlCj3zSte0Mj1FkPQQIY2PjhMB49CmXBV8Z9YJv0=" - ), + "encrypted_content": GPT_6_ASTRA_ENCRYPTED_CONTENT, "summary": [], }, { @@ -681,6 +683,10 @@ GPT_6_ASTRA_TEST_CASE: Final = TranslationTestCase( "model": "azure/gpt-6-astra", "usage": {"input_tokens": 309, "output_tokens": 138}, "content": [ + { + "type": "redacted_thinking", + "data": "litellm_encrypted_reasoning:" + GPT_6_ASTRA_ENCRYPTED_CONTENT, + }, { "type": "text", "text": ( @@ -728,6 +734,22 @@ GPT_5_6_TERRA_CONTENT_FILTERS: Final[JsonValue] = [ }, ] +GPT_5_6_TERRA_ENCRYPTED_CONTENT: Final = ( + "gAAAAABq0CPQDJzIdOV-2rO0aGYCUhKcJifPRv4ZkBaIrKZEARo0rk8CqAe8IV0c2kyBwOFe4M70GMM-HrR1L3hIUxan9dcDB5_S" + "7e45yDaxFMiKpbr20izrXNktdwBQfvYOKoMMd07fslVRR-GHOEtz8JvU98v581sxXgVTxNAWGSl_Kb22g2Fds02t1-1g2hT304vs" + "i1sBxbA6P2b-NFNQ8N3NEfX6FoWi7UF0nH69-Dg3sCMb4dpZ0gGZaZSHwvzXXHESEuiMv1nJSX18tfrFIt29gmw3Mq3sUxZ3eALk" + "PZ54yDUlfYXl6sm2MHXT2xn3l5x6jlG4X8hEqgdfHI76otxmfsXvgtbXCADCes87f00bkWatqWgq27oluRYsY02xDMhtrT4h4C4s" + "PIpV7PA9tdhS6Lo3znsQiOEQvMqnSS3dH1NKnZsAiyjNwVZVZa0opIzI5fpNCGjBmMv_4hB3jBRgenG-Ku75OcwPp6HXMej9Mcvp" + "LIKtXb8FMKS5wqiZzYBLp-F25hzOHI0QVl0pGFL5ya9mWU4hkQYeu6HuH6ZjSujzyy-8ZmlBjnRZSQp8mN4t34u2ofk4RtyuZKJk" + "gjJfN2m5MxBTT3AccMalmbiEa7lJv7GV9kLcqR0GT2Rzd6UMX6qPG92VdfjIgCj1Rzh1IWrJvMloKkMj5qB3nN2dABqfCtjBleSt" + "NoTjO2VDensXf11ceG9-KKWeCjfrG0ONgTVXPjTTUPHR0BxtoixnRxIYBZFb8sKdYAPI_QINjby2SVhhQH7UtookBf7N7OqmIJKb" + "zR4CLshKnvthH4fJQAZdsfoEPR9I943uBpZi4HnF8mQPF61Aiw9B3bh3USBc9_Oh2rJL044B6qYVD3QKKGydzcplRxQWA31m-4Zn" + "9eVxn-xOtx4Tvpdh6BhlFPKX9p2W0mSnb7kjpzX51Ec0J2AsTNUOonaaMFt82MJDELTGcAb3tsAkPzzNK0M_zjQ43CCH5Cu1yV2W" + "hbUzxCu4UM0spAIuzpAZv6JGmW1wF4RtAgLfJjc2LxoxcS9MJbLOUPAYdRlwBTj_IL9PIDeNlGVRgVQYL-9xnNKiWD8Iyu2vj8Ao" + "slM41wVTEXvOAYyEs25Eyo9ia2km8909d42DjFi2_ngA5ofZI89KQHtCzGwfCf_VTLbpMmrOB4NDFVTUYQlAMBvuSsXLXJJp0QPl" + "n2UBHhZ3erIcQ1D990SFRdx1uY47-E2FnCCPljyhJrWHnFxTlGbuEY4i5vPOweqYzj-uBbLVeDbcfjCLjSLFgM4Fa5Zg" +) + GPT_5_6_TERRA_TEST_CASE: Final = TranslationTestCase( scenario="basic", litellm_endpoint="/v1/messages", @@ -752,25 +774,7 @@ GPT_5_6_TERRA_TEST_CASE: Final = TranslationTestCase( "id": "rs_0d178c5d2401a624006ad023cfe304819486bdf040ec55eab7", "type": "reasoning", "content": [], - "encrypted_content": ( - "gAAAAABq0CPQDJzIdOV-2rO0aGYCUhKcJifPRv4ZkBaIrKZEARo0rk8CqAe8IV0c2kyBwOFe4M70GMM-" - "HrR1L3hIUxan9dcDB5_S7e45yDaxFMiKpbr20izrXNktdwBQfvYOKoMMd07fslVRR-GHOEtz8JvU98v5" - "81sxXgVTxNAWGSl_Kb22g2Fds02t1-1g2hT304vsi1sBxbA6P2b-NFNQ8N3NEfX6FoWi7UF0nH69-Dg3" - "sCMb4dpZ0gGZaZSHwvzXXHESEuiMv1nJSX18tfrFIt29gmw3Mq3sUxZ3eALkPZ54yDUlfYXl6sm2MHXT" - "2xn3l5x6jlG4X8hEqgdfHI76otxmfsXvgtbXCADCes87f00bkWatqWgq27oluRYsY02xDMhtrT4h4C4s" - "PIpV7PA9tdhS6Lo3znsQiOEQvMqnSS3dH1NKnZsAiyjNwVZVZa0opIzI5fpNCGjBmMv_4hB3jBRgenG-" - "Ku75OcwPp6HXMej9McvpLIKtXb8FMKS5wqiZzYBLp-F25hzOHI0QVl0pGFL5ya9mWU4hkQYeu6HuH6Zj" - "Sujzyy-8ZmlBjnRZSQp8mN4t34u2ofk4RtyuZKJkgjJfN2m5MxBTT3AccMalmbiEa7lJv7GV9kLcqR0G" - "T2Rzd6UMX6qPG92VdfjIgCj1Rzh1IWrJvMloKkMj5qB3nN2dABqfCtjBleStNoTjO2VDensXf11ceG9-" - "KKWeCjfrG0ONgTVXPjTTUPHR0BxtoixnRxIYBZFb8sKdYAPI_QINjby2SVhhQH7UtookBf7N7OqmIJKb" - "zR4CLshKnvthH4fJQAZdsfoEPR9I943uBpZi4HnF8mQPF61Aiw9B3bh3USBc9_Oh2rJL044B6qYVD3QK" - "KGydzcplRxQWA31m-4Zn9eVxn-xOtx4Tvpdh6BhlFPKX9p2W0mSnb7kjpzX51Ec0J2AsTNUOonaaMFt8" - "2MJDELTGcAb3tsAkPzzNK0M_zjQ43CCH5Cu1yV2WhbUzxCu4UM0spAIuzpAZv6JGmW1wF4RtAgLfJjc2" - "LxoxcS9MJbLOUPAYdRlwBTj_IL9PIDeNlGVRgVQYL-9xnNKiWD8Iyu2vj8AoslM41wVTEXvOAYyEs25E" - "yo9ia2km8909d42DjFi2_ngA5ofZI89KQHtCzGwfCf_VTLbpMmrOB4NDFVTUYQlAMBvuSsXLXJJp0QPl" - "n2UBHhZ3erIcQ1D990SFRdx1uY47-E2FnCCPljyhJrWHnFxTlGbuEY4i5vPOweqYzj-uBbLVeDbcfjCL" - "jSLFgM4Fa5Zg" - ), + "encrypted_content": GPT_5_6_TERRA_ENCRYPTED_CONTENT, "summary": [], }, { @@ -814,6 +818,10 @@ GPT_5_6_TERRA_TEST_CASE: Final = TranslationTestCase( "model": "azure/gpt-5.6-terra", "usage": {"input_tokens": 309, "output_tokens": 92}, "content": [ + { + "type": "redacted_thinking", + "data": "litellm_encrypted_reasoning:" + GPT_5_6_TERRA_ENCRYPTED_CONTENT, + }, { "type": "text", "text": ( diff --git a/tests/integration/translation/messages/bases/bedrock_mantle.py b/tests/integration/translation/messages/bases/bedrock_mantle.py index 7524bf8f7ea..065c8553231 100644 --- a/tests/integration/translation/messages/bases/bedrock_mantle.py +++ b/tests/integration/translation/messages/bases/bedrock_mantle.py @@ -467,6 +467,22 @@ BEDROCK_MANTLE_OPENAI_RESPONSES_MESSAGES_RESPONSE: Final[Mapping[str, JsonValue] } +GPT_6_1_SOL_ENCRYPTED_CONTENT: Final = ( + "rsn_rru6BQHgUc/VWS5vRIAlHwIFeB+Lw/sASAOaRTS5iOdUyDiwGqyLb6EJDqxBCdFOQvGfAF8AAQAVYXdzLWNyeXB0by1wdWJs" + "aWMta2V5AERBaTJyR044dmNxUGt3TndhZkNmeTBpOEFUWlBHSGN4L2s3U1E0UGpyUEp0bUdlZVhXVDREVy9tS2FPdStQUFd6OGc9" + "PQABABFhd3Mta21zLWhpZXJhcmNoeQAKYmVybS1zdG9yZQBcJ5O6OH7BzhGiCKEzUwANq0bcBtwE7aXPWBjYzgTyexPQO0PepR8j" + "RQ3zVYRhZpm6K2TOk2lvrp72PZ5N5jG3KGRkWT2Y2jsL9pnkWwE5kN16+aY2OKOY9ZlzQukCAAAQACDHXld1blcGog/Y0exPb0n+" + "8Vo0dT1X9pL401/Nv3/uQlT6BC0wwixOAufgfaY+ev////8AAAABAAAAAAAAAAAAAAABAAABrRdgMGsGAKubmJ7dCj1A9UQgYVFb" + "1O6+NB+Aom8IGUy3Y6E3c63eWNkWXxN8APSjWFwq012VWZ9Q11DzJmGSwM8gtV25ckYezPfNn46VmSj2idcLs2TxTSxP89NAOfSM" + "zWRnW6Dr3lSlx72uQKBrHWYPyAskoHTw4E6CfnFhXotCshAmT5dgImjniqdLFu4CWh2xpQcmYkFbxXakTBsK80TEI8fpJVjGYQbo" + "mlrZTwAb8bdyEp0/05xVeMVYC47F0tBY5zFOeZ0Nk0ps2PFteD1HQ5GEt7MRMb+VTGGSbq0wmjY7pvdC0f31cs+0GhEZ1N/3707c" + "GjuMevEMn17YYxJ0oelPOatqdAeSSxBGkEKLmzos0tDcisqwWU/94DL2hmTJWacHamboEN74xZ+KwOjKEoyEr0HOL7n5+B9Oojl7" + "MPdWgBnyGUwCsB3BH8t+h5MwMtMJY451PvO4hgl3qK2cN4WfgoYHRvWM2+rj2iJqYBXFPLo5yqWd3XIIimuTGltinK/0sr1uC/AM" + "Q2Qa0qEB0J4Zxla83H7HCEBg4e8wde2qlXZCG9AxVpPQVd1LijTvn7sPIJgYsmSMwKYAZzBlAjBvywiP8Kq5WqQ+kknoQv+bn4Ot" + "wISL+sXhXE+9JK6hhjn02Xd9iSkC/h6IGxAc22UCMQCBBpc3UEvXQSZAxGDy1xUr3JoG4X3sj2Y4yB3SHvBVcxmAWoVo+QAUVMNu" + "lgqXbwU=" +) + GPT_6_1_SOL_TEST_CASE: Final = TranslationTestCase( scenario="basic", litellm_endpoint="/v1/messages", @@ -492,7 +508,7 @@ GPT_6_1_SOL_TEST_CASE: Final = TranslationTestCase( "output": [ { "content": [], - "encrypted_content": "rsn_rru6BQHgUc/VWS5vRIAlHwIFeB+Lw/sASAOaRTS5iOdUyDiwGqyLb6EJDqxBCdFOQvGfAF8AAQAVYXdzLWNyeXB0by1wdWJsaWMta2V5AERBaTJyR044dmNxUGt3TndhZkNmeTBpOEFUWlBHSGN4L2s3U1E0UGpyUEp0bUdlZVhXVDREVy9tS2FPdStQUFd6OGc9PQABABFhd3Mta21zLWhpZXJhcmNoeQAKYmVybS1zdG9yZQBcJ5O6OH7BzhGiCKEzUwANq0bcBtwE7aXPWBjYzgTyexPQO0PepR8jRQ3zVYRhZpm6K2TOk2lvrp72PZ5N5jG3KGRkWT2Y2jsL9pnkWwE5kN16+aY2OKOY9ZlzQukCAAAQACDHXld1blcGog/Y0exPb0n+8Vo0dT1X9pL401/Nv3/uQlT6BC0wwixOAufgfaY+ev////8AAAABAAAAAAAAAAAAAAABAAABrRdgMGsGAKubmJ7dCj1A9UQgYVFb1O6+NB+Aom8IGUy3Y6E3c63eWNkWXxN8APSjWFwq012VWZ9Q11DzJmGSwM8gtV25ckYezPfNn46VmSj2idcLs2TxTSxP89NAOfSMzWRnW6Dr3lSlx72uQKBrHWYPyAskoHTw4E6CfnFhXotCshAmT5dgImjniqdLFu4CWh2xpQcmYkFbxXakTBsK80TEI8fpJVjGYQbomlrZTwAb8bdyEp0/05xVeMVYC47F0tBY5zFOeZ0Nk0ps2PFteD1HQ5GEt7MRMb+VTGGSbq0wmjY7pvdC0f31cs+0GhEZ1N/3707cGjuMevEMn17YYxJ0oelPOatqdAeSSxBGkEKLmzos0tDcisqwWU/94DL2hmTJWacHamboEN74xZ+KwOjKEoyEr0HOL7n5+B9Oojl7MPdWgBnyGUwCsB3BH8t+h5MwMtMJY451PvO4hgl3qK2cN4WfgoYHRvWM2+rj2iJqYBXFPLo5yqWd3XIIimuTGltinK/0sr1uC/AMQ2Qa0qEB0J4Zxla83H7HCEBg4e8wde2qlXZCG9AxVpPQVd1LijTvn7sPIJgYsmSMwKYAZzBlAjBvywiP8Kq5WqQ+kknoQv+bn4OtwISL+sXhXE+9JK6hhjn02Xd9iSkC/h6IGxAc22UCMQCBBpc3UEvXQSZAxGDy1xUr3JoG4X3sj2Y4yB3SHvBVcxmAWoVo+QAUVMNulgqXbwU=", + "encrypted_content": GPT_6_1_SOL_ENCRYPTED_CONTENT, "id": "rs_187e1b0bc4995fc9b7fd2796a1889637", "summary": [], "type": "reasoning", @@ -536,6 +552,10 @@ GPT_6_1_SOL_TEST_CASE: Final = TranslationTestCase( "model": "bedrock_mantle/openai.gpt-6.1-sol", "usage": {"input_tokens": 309, "output_tokens": 141}, "content": [ + { + "type": "redacted_thinking", + "data": "litellm_encrypted_reasoning:" + GPT_6_1_SOL_ENCRYPTED_CONTENT, + }, { "type": "text", "text": "I’ll read `shop/validation.py` first to check how it handles zero, negative, and " @@ -551,6 +571,21 @@ GPT_6_1_SOL_TEST_CASE: Final = TranslationTestCase( }, ) +GPT_6_LUNA_ENCRYPTED_CONTENT: Final = ( + "rsn_rru6BQHgUc/VWS5vRIAlHwIFeJWGzSslRfiCKfPUzG9PfxThM9nB8cV14wc1iKvcbztDAF8AAQAVYXdzLWNyeXB0by1wdWJs" + "aWMta2V5AERBb2ZiS1dnQ24waEpudTVSZGtlLyt3emJKY2VIUXlCeU9rWGd6enJDYU5ydXZYMEFmV1NwK01vVEUwZzhHWnpWb0E9" + "PQABABFhd3Mta21zLWhpZXJhcmNoeQAKYmVybS1zdG9yZQBcJhapxvTAKRUzTqbKGUoTlXJqSZjSljHkimXkxwTyexPQO0PepR8j" + "RQ3zVYQcb1ucdG8HiPjq86SsgccbqbIEmYllOUwDws5V0DtxxRnqoVEhkF3h+EG4I5uuB2kCAAAQABDKdaE231TqbLUInCwHGBD9" + "+U9x7lXzuoVQNhj0KATZ6U5+Ki+ZDt6KLjrNb/LM2v////8AAAABAAAAAAAAAAAAAAABAAABdglH8KRxssIDHD4kcMUX0rX0yogE" + "2V0AOj4vmUzgN+tpKEUhA0/BNphpi42nYRZK61EgLzZk42sWn6sJhR5/WKHHn/7N7FppZkt6kFCDTyhRfPimhQT0S7NiACE4511F" + "HXE/WbRGd2/3uuLOx6PfMFlxuKgJHYQuMEttkYRZz1lwzMrtQTyTtghKKz0QaNtp7umhKHlaFVetqVsKqLD7JpKyBxjbe5VojoKm" + "MQEY3dExNYgvCnOBPb8uLhZCcBemLq1ZQzinbMRyRO1bhQmfGsx4GgFbuA4cyJlyzHe+Af4VqcRM94bELBMXiU64UveyqRkWBdoc" + "sPfPFpVWa+Hnvl37fgn4jBTnhG1xqNuoVCXX2nV/GAMYHdPwsKweIhxg9I8EcrsXOcTWv26hwDRpmdoHHtJkSqKSbudxEs1F5kr5" + "ELytEFS2OWJwwsMvDfBAlC4boaYtxUBG5sREazNwDNRkhFWHb8MwsXBpOqwOBtu1L2n953WUP2I2tHE6WcZ3DpoD2dP3pwBnMGUC" + "MQDonP31z1HPHjm2mGf7SINbrdFZ5xl5ztMBj3g8hJmIpdWtN/W5dJkc7kfNSErvVrQCMDGAL4pNuVt19O4BfjAmLKnNyz5x/z1F" + "PLGJDz2j+9b0Y0FYBIqafypg88wJ6XSc0Q==" +) + GPT_6_LUNA_TEST_CASE: Final = TranslationTestCase( scenario="basic", litellm_endpoint="/v1/messages", @@ -578,7 +613,7 @@ GPT_6_LUNA_TEST_CASE: Final = TranslationTestCase( "output": [ { "content": [], - "encrypted_content": "rsn_rru6BQHgUc/VWS5vRIAlHwIFeJWGzSslRfiCKfPUzG9PfxThM9nB8cV14wc1iKvcbztDAF8AAQAVYXdzLWNyeXB0by1wdWJsaWMta2V5AERBb2ZiS1dnQ24waEpudTVSZGtlLyt3emJKY2VIUXlCeU9rWGd6enJDYU5ydXZYMEFmV1NwK01vVEUwZzhHWnpWb0E9PQABABFhd3Mta21zLWhpZXJhcmNoeQAKYmVybS1zdG9yZQBcJhapxvTAKRUzTqbKGUoTlXJqSZjSljHkimXkxwTyexPQO0PepR8jRQ3zVYQcb1ucdG8HiPjq86SsgccbqbIEmYllOUwDws5V0DtxxRnqoVEhkF3h+EG4I5uuB2kCAAAQABDKdaE231TqbLUInCwHGBD9+U9x7lXzuoVQNhj0KATZ6U5+Ki+ZDt6KLjrNb/LM2v////8AAAABAAAAAAAAAAAAAAABAAABdglH8KRxssIDHD4kcMUX0rX0yogE2V0AOj4vmUzgN+tpKEUhA0/BNphpi42nYRZK61EgLzZk42sWn6sJhR5/WKHHn/7N7FppZkt6kFCDTyhRfPimhQT0S7NiACE4511FHXE/WbRGd2/3uuLOx6PfMFlxuKgJHYQuMEttkYRZz1lwzMrtQTyTtghKKz0QaNtp7umhKHlaFVetqVsKqLD7JpKyBxjbe5VojoKmMQEY3dExNYgvCnOBPb8uLhZCcBemLq1ZQzinbMRyRO1bhQmfGsx4GgFbuA4cyJlyzHe+Af4VqcRM94bELBMXiU64UveyqRkWBdocsPfPFpVWa+Hnvl37fgn4jBTnhG1xqNuoVCXX2nV/GAMYHdPwsKweIhxg9I8EcrsXOcTWv26hwDRpmdoHHtJkSqKSbudxEs1F5kr5ELytEFS2OWJwwsMvDfBAlC4boaYtxUBG5sREazNwDNRkhFWHb8MwsXBpOqwOBtu1L2n953WUP2I2tHE6WcZ3DpoD2dP3pwBnMGUCMQDonP31z1HPHjm2mGf7SINbrdFZ5xl5ztMBj3g8hJmIpdWtN/W5dJkc7kfNSErvVrQCMDGAL4pNuVt19O4BfjAmLKnNyz5x/z1FPLGJDz2j+9b0Y0FYBIqafypg88wJ6XSc0Q==", + "encrypted_content": GPT_6_LUNA_ENCRYPTED_CONTENT, "id": "rs_8de42923e05d5521b0dbd93c1adcea91", "summary": [], "type": "reasoning", @@ -622,6 +657,10 @@ GPT_6_LUNA_TEST_CASE: Final = TranslationTestCase( "model": "bedrock_mantle/openai.gpt-6-luna", "usage": {"input_tokens": 309, "output_tokens": 119}, "content": [ + { + "type": "redacted_thinking", + "data": "litellm_encrypted_reasoning:" + GPT_6_LUNA_ENCRYPTED_CONTENT, + }, { "type": "text", "text": "I’ll read `shop/validation.py` first to check the behavior for zero, negative, and " @@ -637,6 +676,21 @@ GPT_6_LUNA_TEST_CASE: Final = TranslationTestCase( }, ) +GPT_6_ASTRA_ENCRYPTED_CONTENT: Final = ( + "rsn_rru6BQHgUc/VWS5vRIAlHwIFeK6rtxWPF9TjmgBuzPaPZhcqwiqUtYeqNCEyD/fgEZAfAF8AAQAVYXdzLWNyeXB0by1wdWJs" + "aWMta2V5AERBaUE5eUozT3MzTE1rRFhLWGRuWUFYRjdZU2Q3OStjT3V5LzBIRUNnZXBUWDloQXZncVppTC9qVUc3Y2h3dlQ5QVE9" + "PQABABFhd3Mta21zLWhpZXJhcmNoeQAKYmVybS1zdG9yZQBcuTxMV80kCzRfLH6n87syGDT8XBcjqEaE8tDuSQTyexPQO0PepR8j" + "RQ3zVYTdphRKCdq3CJr491CVE1gdVwCTcbqsK/1k7shfoHg93E2fCZXYKq0tnNLGBKhnl90CAAAQAHNHMusCy4BS4gUSPEf8T3Dw" + "PC4nfHsYIw/8I2M8N4Z4dxH0a/6ULM8Zaf5qgZInLv////8AAAABAAAAAAAAAAAAAAABAAABmFno3HYlCh8skcb8d7DvpK045XzP" + "/rPTF9J1TXQEj+Oam1ybdDFFN5bbUrMUhvhaAKoUdhp0TPWFBQnQD7rM1QSBtZY60vbMXdscVeDkmAPh7MsZ6OLAY1olCWPsid7O" + "qATdHKm+S6UfnqnvNLJQ0Hp12YYilLOhW5+BLnmqdm04wkU2kvH4G4vxgRK0hFGMGaoKlCIKRvAmwre2DGXdWx85E9jsq2j1ta1X" + "T6/2eCr7zyXn8+tRI69tNwFVRTcaHu8DvSS1LBzdO7ZwNwUr6PbksO44r/sXeNstFU4u3ZQnUhWz8JNym0ydaia57zXQG3/pl8rJ" + "aA3IshV+Gm+Ie7N3qySlcIEonqN+cYhQN669G4oBuvCKARc2LFdRB1OXhErskZ8c2c3NA9Xt/dXBbclKY9i/I3Usp46t46mALqz7" + "rA2ft/cGLtX3gNoPPEKL2+i1ZDa0EB1Y7agcESa5pYLKgajhUDbaxupsVQ93wZiCcDTr7SlEfMrBkNFGijc3kpdn1avRFEzVC3H/" + "Ys1cPPYJ93hOUYO42b/RibFpBlkLnKtNf//v81kAZzBlAjAdmUn/WQ5cOGyMLQF2CQxoUrAO9HhKULmwznaURUbjukpBRw8v6Pwf" + "lSA+mmU8iv0CMQCASBvNjVVmEd4/s4jH27TFub4sikcdi7+LR5HeitoRhyEJIfB41mRciDxhY80hp6s=" +) + GPT_6_ASTRA_TEST_CASE: Final = TranslationTestCase( scenario="basic", litellm_endpoint="/v1/messages", @@ -664,7 +718,7 @@ GPT_6_ASTRA_TEST_CASE: Final = TranslationTestCase( "output": [ { "content": [], - "encrypted_content": "rsn_rru6BQHgUc/VWS5vRIAlHwIFeK6rtxWPF9TjmgBuzPaPZhcqwiqUtYeqNCEyD/fgEZAfAF8AAQAVYXdzLWNyeXB0by1wdWJsaWMta2V5AERBaUE5eUozT3MzTE1rRFhLWGRuWUFYRjdZU2Q3OStjT3V5LzBIRUNnZXBUWDloQXZncVppTC9qVUc3Y2h3dlQ5QVE9PQABABFhd3Mta21zLWhpZXJhcmNoeQAKYmVybS1zdG9yZQBcuTxMV80kCzRfLH6n87syGDT8XBcjqEaE8tDuSQTyexPQO0PepR8jRQ3zVYTdphRKCdq3CJr491CVE1gdVwCTcbqsK/1k7shfoHg93E2fCZXYKq0tnNLGBKhnl90CAAAQAHNHMusCy4BS4gUSPEf8T3DwPC4nfHsYIw/8I2M8N4Z4dxH0a/6ULM8Zaf5qgZInLv////8AAAABAAAAAAAAAAAAAAABAAABmFno3HYlCh8skcb8d7DvpK045XzP/rPTF9J1TXQEj+Oam1ybdDFFN5bbUrMUhvhaAKoUdhp0TPWFBQnQD7rM1QSBtZY60vbMXdscVeDkmAPh7MsZ6OLAY1olCWPsid7OqATdHKm+S6UfnqnvNLJQ0Hp12YYilLOhW5+BLnmqdm04wkU2kvH4G4vxgRK0hFGMGaoKlCIKRvAmwre2DGXdWx85E9jsq2j1ta1XT6/2eCr7zyXn8+tRI69tNwFVRTcaHu8DvSS1LBzdO7ZwNwUr6PbksO44r/sXeNstFU4u3ZQnUhWz8JNym0ydaia57zXQG3/pl8rJaA3IshV+Gm+Ie7N3qySlcIEonqN+cYhQN669G4oBuvCKARc2LFdRB1OXhErskZ8c2c3NA9Xt/dXBbclKY9i/I3Usp46t46mALqz7rA2ft/cGLtX3gNoPPEKL2+i1ZDa0EB1Y7agcESa5pYLKgajhUDbaxupsVQ93wZiCcDTr7SlEfMrBkNFGijc3kpdn1avRFEzVC3H/Ys1cPPYJ93hOUYO42b/RibFpBlkLnKtNf//v81kAZzBlAjAdmUn/WQ5cOGyMLQF2CQxoUrAO9HhKULmwznaURUbjukpBRw8v6PwflSA+mmU8iv0CMQCASBvNjVVmEd4/s4jH27TFub4sikcdi7+LR5HeitoRhyEJIfB41mRciDxhY80hp6s=", + "encrypted_content": GPT_6_ASTRA_ENCRYPTED_CONTENT, "id": "rs_e6237911ea3e5775acee304296484fab", "summary": [], "type": "reasoning", @@ -708,6 +762,10 @@ GPT_6_ASTRA_TEST_CASE: Final = TranslationTestCase( "model": "bedrock_mantle/openai.gpt-6-astra", "usage": {"input_tokens": 309, "output_tokens": 138}, "content": [ + { + "type": "redacted_thinking", + "data": "litellm_encrypted_reasoning:" + GPT_6_ASTRA_ENCRYPTED_CONTENT, + }, { "type": "text", "text": "I’ll read `shop/validation.py` first to check how it handles zero, negative, and " @@ -723,6 +781,20 @@ GPT_6_ASTRA_TEST_CASE: Final = TranslationTestCase( }, ) +GPT_5_6_TERRA_ENCRYPTED_CONTENT: Final = ( + "rsn_rru6BQHgUc/VWS5vRIAlHwIFeCIDnqsrWiO+hx6u0oBst+OnJIA0RIm1eTWC5ORgkFFNAF8AAQAVYXdzLWNyeXB0by1wdWJs" + "aWMta2V5AERBMEZqYXRRMkhMSEw1clFlcythQm0wc1Bid2hBWHJhN3ZZOFdGenRwOXlsd0NqSkxBOGxWd2loU2dleHhmTFN0Snc9" + "PQABABFhd3Mta21zLWhpZXJhcmNoeQAKYmVybS1zdG9yZQBcWxvahbwUWb5zEZ4mn/Ri1V5Z4KCuYAzvlyiI2ATyexPQO0PepR8j" + "RQ3zVYTwOcBo7nb2ZZfI77C2nwtaEjD275c1LF19dmrCUZ7VPKIYQCMlismkx7f2X1LE4xUCAAAQAPc8v4U/p6fjgu2ZpZgPQUfp" + "PsrIs2TWrckQcNPOhJmhKYDB5gthEp93h1Hf2+ad3P////8AAAABAAAAAAAAAAAAAAABAAABHw7y8mi4YLItPYN81mZNbuQhK0FT" + "G04aQuvQ0966UQxFIMfbeJe+qtyO7uHmyKoz58NNb0ud67KhOPoip6wfSxVcw2RTBCAljj+gkLciZdZlGDiLyutdaHfCRh6StNZD" + "tXG0xPgbn0mN93xnE0hDa05mjcNsNckDaS2KkTMzoXF0tqWoi+Vl0fMSK3CF/49QKyGXYsDeOPB0b8XghpfF1PvMqMI/Ua9Sd98l" + "999mYu71TaBKvgRS9yLWBXg9YcmfZx2hulOUkK2a9aCQ5aN661B0MTdykIm+CokaKogZCv6e/MXncHPPtBeqtuQtsggedinoW/qD" + "09HU+rLUysC+vbHT65YqoVL9wW42XZa79V5QBFYIQL787eVy54qbPxrryWKuTOsN9xRFqaSOv1L10ABnMGUCMFnJHyhVr8ujzWTw" + "nSQCuS+zKrvPgXQfEePxJOR/KFHa7CUXZ8kBuEoW0kVxx3biLgIxAK+8PazVOuepKMy5cuKKsJChOnvQ1Npz7Gpmrx0CXO7X2sIN" + "1WJTlhdNHEclLuwoSQ==" +) + GPT_5_6_TERRA_TEST_CASE: Final = TranslationTestCase( scenario="basic", litellm_endpoint="/v1/messages", @@ -750,7 +822,7 @@ GPT_5_6_TERRA_TEST_CASE: Final = TranslationTestCase( "output": [ { "content": [], - "encrypted_content": "rsn_rru6BQHgUc/VWS5vRIAlHwIFeCIDnqsrWiO+hx6u0oBst+OnJIA0RIm1eTWC5ORgkFFNAF8AAQAVYXdzLWNyeXB0by1wdWJsaWMta2V5AERBMEZqYXRRMkhMSEw1clFlcythQm0wc1Bid2hBWHJhN3ZZOFdGenRwOXlsd0NqSkxBOGxWd2loU2dleHhmTFN0Snc9PQABABFhd3Mta21zLWhpZXJhcmNoeQAKYmVybS1zdG9yZQBcWxvahbwUWb5zEZ4mn/Ri1V5Z4KCuYAzvlyiI2ATyexPQO0PepR8jRQ3zVYTwOcBo7nb2ZZfI77C2nwtaEjD275c1LF19dmrCUZ7VPKIYQCMlismkx7f2X1LE4xUCAAAQAPc8v4U/p6fjgu2ZpZgPQUfpPsrIs2TWrckQcNPOhJmhKYDB5gthEp93h1Hf2+ad3P////8AAAABAAAAAAAAAAAAAAABAAABHw7y8mi4YLItPYN81mZNbuQhK0FTG04aQuvQ0966UQxFIMfbeJe+qtyO7uHmyKoz58NNb0ud67KhOPoip6wfSxVcw2RTBCAljj+gkLciZdZlGDiLyutdaHfCRh6StNZDtXG0xPgbn0mN93xnE0hDa05mjcNsNckDaS2KkTMzoXF0tqWoi+Vl0fMSK3CF/49QKyGXYsDeOPB0b8XghpfF1PvMqMI/Ua9Sd98l999mYu71TaBKvgRS9yLWBXg9YcmfZx2hulOUkK2a9aCQ5aN661B0MTdykIm+CokaKogZCv6e/MXncHPPtBeqtuQtsggedinoW/qD09HU+rLUysC+vbHT65YqoVL9wW42XZa79V5QBFYIQL787eVy54qbPxrryWKuTOsN9xRFqaSOv1L10ABnMGUCMFnJHyhVr8ujzWTwnSQCuS+zKrvPgXQfEePxJOR/KFHa7CUXZ8kBuEoW0kVxx3biLgIxAK+8PazVOuepKMy5cuKKsJChOnvQ1Npz7Gpmrx0CXO7X2sIN1WJTlhdNHEclLuwoSQ==", + "encrypted_content": GPT_5_6_TERRA_ENCRYPTED_CONTENT, "id": "rs_a4e6e9898f225c88b4d948b6488bfacc", "summary": [], "type": "reasoning", @@ -795,6 +867,10 @@ GPT_5_6_TERRA_TEST_CASE: Final = TranslationTestCase( "model": "bedrock_mantle/openai.gpt-5.6-terra", "usage": {"input_tokens": 309, "output_tokens": 109}, "content": [ + { + "type": "redacted_thinking", + "data": "litellm_encrypted_reasoning:" + GPT_5_6_TERRA_ENCRYPTED_CONTENT, + }, { "type": "text", "text": "I’ll read `shop/validation.py` first to verify the current handling of zero, " diff --git a/tests/unit/accounting/test_metrics.py b/tests/unit/accounting/test_metrics.py index 054ceaba967..d05d860bd66 100644 --- a/tests/unit/accounting/test_metrics.py +++ b/tests/unit/accounting/test_metrics.py @@ -5,12 +5,11 @@ import pytest from prometheus_client.metrics import MetricWrapperBase from litellm.accounting import metrics -from litellm.accounting.metrics import record_deliveries -from litellm.repositories.accounting_delivery_repository import DeliveryMetrics +from litellm.accounting.metrics import Cadence, record_backlog, record_storage +from litellm.repositories.accounting_delivery_repository import DeliveryBacklog, DeliveryStorage -_DELIVERY_METRICS = DeliveryMetrics( - pending_count=4, parked_count=1, oldest_pending_seconds=-3.0, batch_count=7, table_bytes=8192 -) +_BACKLOG = DeliveryBacklog(pending_count=4, parked_count=1, oldest_pending_seconds=-3.0) +_STORAGE = DeliveryStorage(batch_count=7, table_bytes=8192) def _sample(metric: MetricWrapperBase, suffix: str = "") -> float: @@ -23,7 +22,8 @@ def test_delivery_metrics_publish_backlog_and_clamp_clock_skew() -> None: assert instruments is not None batches_before = _sample(instruments.batch_time, "_count") - record_deliveries(_DELIVERY_METRICS, 0.5) + record_backlog(_BACKLOG, 0.5) + record_storage(_STORAGE) assert ( _sample(instruments.pending), @@ -46,4 +46,17 @@ def test_delivery_metrics_are_skipped_when_prometheus_client_is_not_installed( spec.loader.exec_module(module) - assert module.record_deliveries(_DELIVERY_METRICS, 0.5) is None + assert module.record_backlog(_BACKLOG, 0.5) is None + assert module.record_storage(_STORAGE) is None + + +def test_cadence_is_due_once_per_interval() -> None: + now = [100.0] + cadence = Cadence(lambda: now[0]) + + observed = [] + for moment in (100.0, 101.0, 999.9, 1000.0, 1000.0, 1899.0, 1900.0): + now[0] = moment + observed.append(cadence.due(900)) + + assert observed == [True, False, False, True, False, False, True] diff --git a/tests/unit/caching/test_request_redis_batch_post_call.py b/tests/unit/caching/test_request_redis_batch_post_call.py index 41d508a0eea..01feee0be54 100644 --- a/tests/unit/caching/test_request_redis_batch_post_call.py +++ b/tests/unit/caching/test_request_redis_batch_post_call.py @@ -424,7 +424,7 @@ async def test_spend_counter_increments_ride_the_native_post_call_batch(monkeypa with request_redis_batch_scope(native_factory=natives, native_send_factory=fake_send): # pyright: ignore[reportArgumentType] # test double await proxy_server._apply_spend_counter_increments(pending) - declared = list(spend.post_call) + declared = [(request, increments) for request, increments, _ in spend.post_call] assert natives.built[0].events == [] await flush_post_call_redis_batches() @@ -433,6 +433,27 @@ async def test_spend_counter_increments_ride_the_native_post_call_batch(monkeypa assert spend.commands == [] +@pytest.mark.asyncio +async def test_the_request_sees_its_committed_spend_charge_settle_whoever_sends_the_batch(monkeypatch): + from litellm.proxy import proxy_server + + spend = FakeSpendCounters(redis={}) + monkeypatch.setattr(proxy_server, "spend_counter_store", spend.store()) + + with request_redis_batch_scope(native_factory=FakeNatives(), native_send_factory=fake_send) as request: # pyright: ignore[reportArgumentType] # test double + await proxy_server._apply_spend_counter_increments([PendingSpendIncrement("spend:key:k1", 0.5)]) + without_committed_spend = request.post_call_settlements + await proxy_server._apply_spend_counter_increments( + [PendingSpendIncrement("spend:key:k1", 0.5), PendingSpendIncrement("spend:committed:key:k1", 0.5)] + ) + declared = [settlement.settled for settlement in request.post_call_settlements] + spend.settle_post_call() + sent = [settlement.settled for settlement in request.post_call_settlements] + + assert without_committed_spend == () + assert (declared, sent) == ([False], [True]) + + @pytest.mark.asyncio async def test_native_post_call_ops_go_out_with_the_post_call_flush_not_the_request_boundary(): natives = FakeNatives() diff --git a/tests/unit/integration_support/test_database_relay.py b/tests/unit/integration_support/test_database_relay.py index 98501f84058..de2bf1f5d2b 100644 --- a/tests/unit/integration_support/test_database_relay.py +++ b/tests/unit/integration_support/test_database_relay.py @@ -4,7 +4,7 @@ from typing import Final import pytest -from tests.integration._support.database_relay import TriggerScanner +from tests.integration._support.database_relay import StatementMarkerScanner, TriggerScanner TRIGGER: Final = b'SELECT "startTime" FROM "LiteLLM_SpendLogs"' @@ -26,3 +26,55 @@ def test_trigger_scanner_reports_a_match_once() -> None: scanner: Final = TriggerScanner(TRIGGER) assert scanner.feed(b"x" + TRIGGER + b"y") assert not scanner.feed(b"z") + + +STARTUP: Final = b"\x00\x00\x00\x08\x00\x03\x00\x00" + + +def _message(kind: bytes, payload: bytes) -> bytes: + return kind + (len(payload) + 4).to_bytes(4, "big") + payload + + +PARSE: Final = _message(b"P", b"s1\x00" + TRIGGER + b"\x00\x00\x00") +BIND: Final = _message(b"B", b"\x00s1\x00\x00\x00\x00\x00\x00\x00") + + +@pytest.mark.parametrize("split_at", range(1, len(STARTUP + PARSE + BIND))) +def test_statement_markers_survive_every_protocol_frame_split(split_at: int) -> None: + scanner: Final = StatementMarkerScanner((TRIGGER,)) + stream: Final = STARTUP + PARSE + BIND + assert (*scanner.feed(stream[:split_at]), *scanner.feed(stream[split_at:])) == ( + frozenset({TRIGGER}), + frozenset({TRIGGER}), + ) + + +def test_statement_markers_are_local_to_a_connection_and_forget_closed_statements() -> None: + scanner: Final = StatementMarkerScanner((TRIGGER,)) + other: Final = StatementMarkerScanner((TRIGGER,)) + assert tuple(scanner.feed(STARTUP + PARSE)) == (frozenset({TRIGGER}),) + assert tuple(other.feed(STARTUP + BIND)) == (frozenset(),) + assert tuple(scanner.feed(_message(b"C", b"Ss1\x00") + BIND)) == (frozenset(),) + + +def test_replaced_and_unnamed_statements_do_not_reuse_old_markers() -> None: + scanner: Final = StatementMarkerScanner((TRIGGER,)) + assert tuple(scanner.feed(STARTUP + PARSE)) == (frozenset({TRIGGER}),) + replacement: Final = _message(b"P", b"s1\x00SELECT 1\x00\x00\x00") + unnamed: Final = _message(b"P", b"\x00" + TRIGGER + b"\x00\x00\x00") + assert tuple(scanner.feed(replacement + BIND + unnamed)) == (frozenset(), frozenset(), frozenset({TRIGGER})) + + +def test_statement_scanner_rejects_invalid_frame_lengths() -> None: + scanner: Final = StatementMarkerScanner((TRIGGER,)) + assert tuple(scanner.feed(STARTUP)) == () + with pytest.raises(ValueError, match="Invalid PostgreSQL message length"): + tuple(scanner.feed(b"P\x00\x00\x00\x03")) + + +def test_statement_scanner_matches_bound_parameter_markers() -> None: + scanner: Final = StatementMarkerScanner((TRIGGER,)) + bind: Final = _message( + b"B", b"\x00s2\x00\x00\x00\x00\x01" + len(TRIGGER).to_bytes(4, "big") + TRIGGER + b"\x00\x00" + ) + assert tuple(scanner.feed(STARTUP + bind)) == (frozenset({TRIGGER}),) diff --git a/tests/unit/integrations/generic_api/test_generic_api_callback.py b/tests/unit/integrations/generic_api/test_generic_api_callback.py index afaae02321d..8f1e79562db 100644 --- a/tests/unit/integrations/generic_api/test_generic_api_callback.py +++ b/tests/unit/integrations/generic_api/test_generic_api_callback.py @@ -20,6 +20,8 @@ from litellm.integrations.gcs_pubsub.pub_sub import * from litellm.integrations.generic_api.generic_api_callback import LOG_FORMAT_TYPES, GenericAPILogger from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys from litellm.types.utils import StandardLoggingPayload from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -795,3 +797,27 @@ async def test_generic_api_concurrent_flushes_deliver_each_event_once(log_format task.cancel() await asyncio.gather(*owned_tasks, return_exceptions=True) await handler.close() + + +@pytest.mark.asyncio +async def test_v1_payload_of_a_chat_row_never_shows_the_committed_spend_auth_read() -> None: + counter_keys: Final = EntitySpendCounterKeys.of( + token="hashed-token", team_id="t1", user_id=None, org_id=None, end_user_id=None + ) + auth: Final = UserAPIKeyAuth( + api_key="hashed-token", + team_id="t1", + committed_budget_spend=CommittedBudgetSpend(counter_keys=counter_keys, key=0.123456, team=0.654321), + ) + generic_logger: Final = GenericAPILogger(endpoint="https://example.com/api/logs", flush_interval=3600) + + payload: Final = generic_logger._get_v1_logging_payload( + kwargs={"litellm_params": {"metadata": {"user_api_key_auth": auth}}}, + response_obj={"usage": {}}, + start_time=None, + end_time=None, + ) + + assert "team_id='t1'" in payload["metadata"] + assert "0.123456" not in payload["metadata"] + assert str(counter_keys.team) not in payload["metadata"] diff --git a/tests/unit/integrations/otel/test_otel_v2_sources_of_truth.py b/tests/unit/integrations/otel/test_otel_v2_sources_of_truth.py index c29da4607eb..3c95a85e9a8 100644 --- a/tests/unit/integrations/otel/test_otel_v2_sources_of_truth.py +++ b/tests/unit/integrations/otel/test_otel_v2_sources_of_truth.py @@ -539,6 +539,27 @@ def test_llm_call_adapter_extracts_all_fields(): assert data.identity.key_hash == "hsh" +@pytest.mark.parametrize( + ("model_group", "response_model", "expected"), + [ + ("auto-router", "auto-router", ("gpt-4o-group", "gpt-4o-group")), + ("smart", "smart", ("gpt-4o-group", "gpt-4o-group")), + ("auto-router", "gpt-4o-2024", ("gpt-4o-group", "gpt-4o-2024")), + (None, None, ("gpt-4o", None)), + ], + ids=["restamped-to-router", "restamped-to-alias", "provider-model", "no-group-no-response"], +) +def test_llm_call_adapter_names_the_routed_model_for_an_auto_router_call(model_group, response_model, expected): + data = LLMCallSpanData.from_standard_logging_payload( + _sample_payload( + model_group=model_group, + response={"id": "resp_1", "model": response_model, "choices": [{"finish_reason": "stop"}]}, + metadata={"routing_decision": {"routed_model": "gpt-4o-group"}}, + ) + ) + assert (data.request_model, data.response_model) == expected + + def test_llm_call_adapter_extracts_cache_tokens_from_usage_object(): payload = _sample_payload() payload["metadata"] = { diff --git a/tests/unit/integrations/test_opentelemetry.py b/tests/unit/integrations/test_opentelemetry.py index c4fd7ed51c7..4ef95d3a79d 100644 --- a/tests/unit/integrations/test_opentelemetry.py +++ b/tests/unit/integrations/test_opentelemetry.py @@ -372,6 +372,35 @@ def test_v1_routing_diagnostics_add_scalar_attributes_and_preserve_legacy_metada assert finished.attributes["metadata.routing_decision"] == str(decision) +@pytest.mark.parametrize( + ("decision", "expected"), + [({"router_model_name": "auto-router", "routed_model": "answer"}, "answer"), ("not-a-mapping", "auto-router")], + ids=["routed", "malformed-decision"], +) +def test_v1_auto_router_response_model_names_the_routed_model(decision: object, expected: str) -> None: + exporter: Final = InMemorySpanExporter() + provider: Final = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + otel: Final = OpenTelemetry(tracer_provider=provider) + with provider.get_tracer(__name__).start_as_current_span("inference") as span: + otel.set_attributes( + span, + { + "model": "gpt-4o-mini", + "standard_logging_object": { + "id": "routing-call", + "call_type": "acompletion", + "model_group": "auto-router", + "metadata": {"routing_decision": decision}, + }, + }, + {"model": "auto-router", "choices": [], "usage": {"prompt_tokens": 3}}, + ) + (finished,) = exporter.get_finished_spans() + attributes: Final = finished.attributes + assert (attributes["gen_ai.response.model"], attributes["gen_ai.usage.input_tokens"]) == (expected, 3) + + def test_v1_server_tool_calls_lead_the_completion_function_call_attributes() -> None: exporter: Final = InMemorySpanExporter() provider: Final = TracerProvider() diff --git a/tests/unit/integrations/test_prometheus_carried_budget_state.py b/tests/unit/integrations/test_prometheus_carried_budget_state.py index e49e6baa8b0..9e006c20079 100644 --- a/tests/unit/integrations/test_prometheus_carried_budget_state.py +++ b/tests/unit/integrations/test_prometheus_carried_budget_state.py @@ -4,13 +4,22 @@ from request metadata. get_*_object only runs when that state is missing (custom auth, SDK callers, failure paths) """ +import asyncio +from collections.abc import Sequence from datetime import datetime, timezone +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest from prometheus_client import REGISTRY +from litellm.caching.redis_batch import ( + active_post_call_redis_batch, + flush_post_call_redis_batches, + request_redis_batch_scope, +) from litellm.integrations.prometheus import PrometheusLogger +from litellm.proxy import proxy_server from litellm.proxy._types import ( LiteLLM_BudgetTable, LiteLLM_OrganizationTable, @@ -24,11 +33,16 @@ from litellm.proxy.spend_tracking.carried_budget_state import ( carry_organization_budget_state, carry_team_and_user_budget_state, ) +from litellm.proxy.spend_tracking.spend_counter_batch import PendingSpendIncrement from litellm.types.proxy.carried_budget_state import ( + CommittedBudgetSpend, + EntitySpendCounterKeys, KeyBudgetSnapshot, TeamBudgetSnapshot, UserBudgetSnapshot, ) +from tests.unit.caching.test_redis_batch import FakeNatives, fake_send, make +from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters TEAM_RESET_AT = datetime(2026, 10, 1, tzinfo=timezone.utc) USER_RESET_AT = datetime(2026, 11, 1, tzinfo=timezone.utc) @@ -55,6 +69,15 @@ def prometheus_logger(): return PrometheusLogger() +@pytest.fixture(autouse=True) +def spend_counters(): + spend: Final = FakeSpendCounters(redis=None) + with patch( # test-quality-ok: prometheus reads the proxy_server spend counters at emit time, no injection seam + "litellm.proxy.proxy_server.spend_counter_store", spend.store() + ): + yield spend + + @pytest.fixture def getters(): """Every response-path object getter, patched where prometheus imports them from.""" @@ -137,6 +160,7 @@ def _request_metadata(token: UserAPIKeyAuth) -> dict: return { **LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(token), **carried_budget_metadata(token), + "user_api_key_auth": token, } @@ -198,6 +222,499 @@ async def test_authed_request_sets_every_gauge_without_any_object_getter(prometh prometheus_logger.litellm_user_budget_remaining_hours_metric.labels().set.assert_called_once() +_COUNTER_KEYS: Final = ("spend:key:hashed", "spend:team:t1", "spend:user:u1", "spend:org:o1") +_TWIN_KEYS: Final = ( + "spend:committed:key:hashed", + "spend:committed:team:t1", + "spend:committed:user:u1", + "spend:committed:org:o1", +) + + +def _store(spend: FakeSpendCounters, counter_keys: Sequence[str], values: Sequence[float]) -> None: + tier: Final = spend.memory if spend.redis is None else spend.redis + tier.update(zip(counter_keys, values, strict=True)) + + +def _store_counters(spend: FakeSpendCounters, key: float, team: float, user: float, org: float) -> None: + _store(spend, _COUNTER_KEYS, (key, team, user, org)) + + +def _store_committed(spend: FakeSpendCounters, key: float, team: float, user: float, org: float) -> None: + _store(spend, _TWIN_KEYS, (key, team, user, org)) + + +def _remaining(prometheus_logger: PrometheusLogger) -> dict[str, float]: + return { + "key": prometheus_logger.litellm_remaining_api_key_budget_metric.labels().set.call_args[0][0], + "team": prometheus_logger.litellm_remaining_team_budget_metric.labels().set.call_args[0][0], + "user": prometheus_logger.litellm_remaining_user_budget_metric.labels().set.call_args[0][0], + "org": prometheus_logger.litellm_remaining_org_budget_metric.labels().set.call_args[0][0], + } + + +def _token_committed_at_auth() -> UserAPIKeyAuth: + token = _authed_token() + token.committed_budget_spend = CommittedBudgetSpend( + counter_keys=EntitySpendCounterKeys.of( + token="hashed", team_id="t1", user_id="u1", org_id="o1", end_user_id=None + ), + key=3.0, + team=30.0, + user=8.0, + org=60.0, + ) + return token + + +@pytest.mark.asyncio +async def test_gauges_use_the_committed_spend_at_emit_when_other_requests_finished_after_auth( + prometheus_logger, getters, spend_counters +): + """This request's 2.0 plus 1.0 from another request that finished first, so the gauge is never set back.""" + _stub_gauges(prometheus_logger) + _store_committed(spend_counters, key=6.0, team=33.0, user=11.0, org=63.0) + + await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth())) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(4.0), + "team": pytest.approx(267.0), + "user": pytest.approx(39.0), + "org": pytest.approx(437.0), + } + + +@pytest.mark.asyncio +async def test_gauges_add_the_cost_to_another_requests_charge_while_the_own_charge_is_not_in_yet( + prometheus_logger, getters, spend_counters +): + """1.0 from another request reached the counters during this long request, this request's 2.0 has not.""" + _stub_gauges(prometheus_logger) + _store_committed(spend_counters, key=4.0, team=31.0, user=9.0, org=61.0) + + await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth())) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(4.0), + "team": pytest.approx(267.0), + "user": pytest.approx(39.0), + "org": pytest.approx(437.0), + } + + +@pytest.mark.asyncio +async def test_gauges_read_no_committed_counter_auth_did_not_name_as_budgeted( + prometheus_logger, getters, spend_counters +): + _stub_gauges(prometheus_logger) + _store_committed(spend_counters, key=9.0, team=99.0, user=19.0, org=99.0) + + await _emit(prometheus_logger, _request_metadata(_authed_token())) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(7.0), + "team": pytest.approx(278.0), + "user": pytest.approx(43.0), + "org": 458.0, + } + + +def _billed_row_metadata() -> dict: + return {key: value for key, value in _request_metadata(_authed_token()).items() if key != "user_api_key_auth"} + + +@pytest.mark.asyncio +async def test_row_with_no_auth_adds_its_cost_to_the_committed_spend_of_the_entities_it_is_charged_to( + prometheus_logger, getters, spend_counters +): + """Counter increments flush after the success callbacks, so the row's own charge is not in the counters yet.""" + _stub_gauges(prometheus_logger) + _store_committed(spend_counters, key=6.0, team=33.0, user=11.0, org=63.0) + + await _emit(prometheus_logger, _billed_row_metadata()) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(2.0), + "team": pytest.approx(265.0), + "user": pytest.approx(37.0), + "org": 435.0, + } + + +@pytest.mark.asyncio +async def test_row_with_no_auth_and_no_committed_counters_keeps_the_stored_spend( + prometheus_logger, getters, spend_counters +): + _stub_gauges(prometheus_logger) + + await _emit(prometheus_logger, _billed_row_metadata()) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(7.0), + "team": pytest.approx(278.0), + "user": pytest.approx(43.0), + "org": 458.0, + } + + +@pytest.mark.asyncio +async def test_slow_committed_read_on_a_failed_request_is_bounded_and_keeps_the_auth_time_spend( + prometheus_logger, getters, monkeypatch +): + async def never_answers(counter: str) -> None: + await asyncio.sleep(30) + + monkeypatch.setattr( + "litellm.proxy.proxy_server.spend_counter_store", + FakeSpendCounters(redis={}, batch_reads=False, read_hook=never_answers).store(), + ) + monkeypatch.setenv("PROMETHEUS_BUDGET_METRICS_PER_REQUEST_TIMEOUT", "0.01") + _stub_gauges(prometheus_logger) + + await asyncio.wait_for( + prometheus_logger.async_log_failure_event( + _failed_request(_request_metadata(_token_committed_at_auth())), None, KEY_RESET_AT, KEY_RESET_AT + ), + timeout=5, + ) + + prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(pytest.approx(440.0)) + + +@pytest.mark.asyncio +async def test_slow_committed_read_on_a_successful_request_still_sets_the_gauges_from_the_auth_time_spend( + prometheus_logger, getters, monkeypatch +): + async def never_answers(counter: str) -> None: + await asyncio.sleep(30) + + monkeypatch.setattr( + "litellm.proxy.proxy_server.spend_counter_store", + FakeSpendCounters(redis={}, batch_reads=False, read_hook=never_answers).store(), + ) + monkeypatch.setenv("PROMETHEUS_BUDGET_METRICS_PER_REQUEST_TIMEOUT", "0.5") + _stub_gauges(prometheus_logger) + + await asyncio.wait_for(_emit(prometheus_logger, _request_metadata(_token_committed_at_auth())), timeout=5) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(5.0), + "team": pytest.approx(268.0), + "user": pytest.approx(40.0), + "org": pytest.approx(438.0), + } + + +@pytest.fixture +def redis_spend_counters(): + spend: Final = FakeSpendCounters(redis={}) + with patch( # test-quality-ok: prometheus reads the proxy_server spend counters at emit time, no injection seam + "litellm.proxy.proxy_server.spend_counter_store", spend.store() + ): + yield spend + + +def _request_scope() -> request_redis_batch_scope: + return request_redis_batch_scope( + native_factory=FakeNatives(), # pyright: ignore[reportArgumentType] # test double + native_send_factory=fake_send, # pyright: ignore[reportArgumentType] # test double + ) + + +async def _declare_the_own_two_dollar_charge() -> None: + await proxy_server._apply_spend_counter_increments( + tuple(PendingSpendIncrement(counter_key, 2.0) for counter_key in (*_COUNTER_KEYS, *_TWIN_KEYS)) + ) + + +_BOTH_TWO_DOLLAR_CHARGES_COUNTED_ONCE: Final = { + "key": pytest.approx(3.0), + "team": pytest.approx(266.0), + "user": pytest.approx(38.0), + "org": pytest.approx(436.0), +} + + +@pytest.mark.asyncio +async def test_gauges_add_the_own_charge_to_a_same_size_charge_of_another_request_while_the_own_waits_in_the_pipeline( + prometheus_logger, getters, redis_spend_counters +): + """Another request's 2.0 reached the counters and this request's 2.0 is queued, so both count.""" + _stub_gauges(prometheus_logger) + _store_committed(redis_spend_counters, key=5.0, team=32.0, user=10.0, org=62.0) + + with _request_scope(): + await _declare_the_own_two_dollar_charge() + await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth())) + await flush_post_call_redis_batches() + + assert _remaining(prometheus_logger) == _BOTH_TWO_DOLLAR_CHARGES_COUNTED_ONCE + + +@pytest.mark.asyncio +async def test_gauges_do_not_add_the_own_charge_twice_once_the_pipeline_went_out( + prometheus_logger, getters, redis_spend_counters +): + """The 2.0 over the auth time spend is this request's own charge, settled before the gauges read it.""" + _stub_gauges(prometheus_logger) + _store_committed(redis_spend_counters, key=3.0, team=30.0, user=8.0, org=60.0) + + with _request_scope(): + await _declare_the_own_two_dollar_charge() + redis_spend_counters.settle_post_call() + await flush_post_call_redis_batches() + await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth())) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(5.0), + "team": pytest.approx(268.0), + "user": pytest.approx(40.0), + "org": pytest.approx(438.0), + } + + +@pytest.mark.asyncio +async def test_an_unrelated_post_call_flush_does_not_make_the_gauge_drop_its_own_unsent_charge( + prometheus_logger, getters, redis_spend_counters +): + """An unrelated write went out first, then another request's 2.0 reached the counters. This request's 2.0 counts + once whether it is undeclared, declared but unsent, or settled.""" + _stub_gauges(prometheus_logger) + _store_committed(redis_spend_counters, key=5.0, team=32.0, user=10.0, org=62.0) + unrelated_cache, _ = make() + metadata: Final = _request_metadata(_token_committed_at_auth()) + + with _request_scope(): + active_post_call_redis_batch(unrelated_cache).increment("unrelated-write", 1.0) + await flush_post_call_redis_batches() + await _emit(prometheus_logger, metadata) + undeclared: Final = _remaining(prometheus_logger) + await _declare_the_own_two_dollar_charge() + await _emit(prometheus_logger, metadata) + declared_unsent: Final = _remaining(prometheus_logger) + redis_spend_counters.settle_post_call() + await _emit(prometheus_logger, metadata) + settled: Final = _remaining(prometheus_logger) + await flush_post_call_redis_batches() + + assert {"undeclared": undeclared, "declared_unsent": declared_unsent, "settled": settled} == { + "undeclared": _BOTH_TWO_DOLLAR_CHARGES_COUNTED_ONCE, + "declared_unsent": _BOTH_TWO_DOLLAR_CHARGES_COUNTED_ONCE, + "settled": _BOTH_TWO_DOLLAR_CHARGES_COUNTED_ONCE, + } + + +@pytest.mark.asyncio +async def test_gauges_count_the_own_charge_once_when_it_settles_while_the_gauges_read_the_counters( + prometheus_logger, getters, redis_spend_counters +): + """Another request's 2.0 reached the counters, and this request's 2.0 goes out during the gauges' read.""" + _stub_gauges(prometheus_logger) + _store_committed(redis_spend_counters, key=5.0, team=32.0, user=10.0, org=62.0) + + async def settle_during_read(_counter: str) -> None: + redis_spend_counters.settle_post_call() + + redis_spend_counters.batch_reads = False + redis_spend_counters.read_hook = settle_during_read + + with _request_scope(): + await _declare_the_own_two_dollar_charge() + await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth())) + await flush_post_call_redis_batches() + + assert _remaining(prometheus_logger) == _BOTH_TWO_DOLLAR_CHARGES_COUNTED_ONCE + + +@pytest.mark.asyncio +async def test_gauges_without_redis_treat_a_charge_equal_to_the_own_as_the_own( + prometheus_logger, getters, spend_counters +): + """Without Redis the charge increments in memory at once, so a request scope says nothing about it.""" + _stub_gauges(prometheus_logger) + _store_committed(spend_counters, key=5.0, team=32.0, user=10.0, org=62.0) + + with request_redis_batch_scope(): + await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth())) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(5.0), + "team": pytest.approx(268.0), + "user": pytest.approx(40.0), + "org": pytest.approx(438.0), + } + + +@pytest.mark.asyncio +async def test_gauges_add_the_cost_to_the_auth_time_spend_before_the_charge_reaches_the_counters( + prometheus_logger, getters, spend_counters +): + _stub_gauges(prometheus_logger) + _store_committed(spend_counters, key=3.0, team=30.0, user=8.0, org=60.0) + + await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth())) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(5.0), + "team": pytest.approx(268.0), + "user": pytest.approx(40.0), + "org": pytest.approx(438.0), + } + + +@pytest.mark.asyncio +async def test_gauges_keep_the_auth_time_spend_when_the_proxy_spend_counters_cannot_be_read( + prometheus_logger, getters, monkeypatch +): + _stub_gauges(prometheus_logger) + monkeypatch.delattr("litellm.proxy.proxy_server.spend_counter_store") + + await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth())) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(5.0), + "team": pytest.approx(268.0), + "user": pytest.approx(40.0), + "org": pytest.approx(438.0), + } + + +@pytest.mark.asyncio +async def test_a_budget_hold_on_the_enforcement_counters_never_shows_in_the_gauges( + prometheus_logger, getters, spend_counters +): + _stub_gauges(prometheus_logger) + _store_committed(spend_counters, key=5.0, team=32.0, user=10.0, org=62.0) + _store_counters(spend_counters, key=9.0, team=36.0, user=14.0, org=66.0) + + await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth())) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(5.0), + "team": pytest.approx(268.0), + "user": pytest.approx(40.0), + "org": pytest.approx(438.0), + } + + +def _failed_request(metadata: dict, response_cost: float = 0.0) -> dict: + return { + "model": "openai/gpt-5.4-nano", + "litellm_params": {"metadata": metadata}, + "standard_logging_object": { + "model": "openai/gpt-5.4-nano", + "model_id": "dep-1", + "model_group": "chat", + "api_base": "https://api.openai.com", + "custom_llm_provider": "openai", + "status": "failure", + "response_cost": response_cost, + "request_tags": [], + "hidden_params": {"litellm_overhead_time_ms": None, "additional_headers": None}, + "metadata": { + "user_api_key_hash": "hashed", + "user_api_key_alias": "key-alias", + "user_api_key_team_id": "t1", + "user_api_key_team_alias": "team-alias", + "user_api_key_user_id": "u1", + "user_api_key_org_id": "o1", + "requester_metadata": None, + "user_api_key_end_user_id": None, + }, + }, + "exception": Exception("upstream failed"), + } + + +@pytest.mark.asyncio +async def test_failed_request_sets_the_org_gauge_from_the_carried_org_and_committed_spend( + prometheus_logger, getters, spend_counters +): + _stub_gauges(prometheus_logger) + _store_committed(spend_counters, key=3.0, team=30.0, user=8.0, org=65.0) + _store_counters(spend_counters, key=9.0, team=36.0, user=14.0, org=90.0) + + await prometheus_logger.async_log_failure_event( + _failed_request(_request_metadata(_token_committed_at_auth())), None, KEY_RESET_AT, KEY_RESET_AT + ) + + assert getters["get_org_object"].await_count == 0 + prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(pytest.approx(435.0)) + + +@pytest.mark.asyncio +async def test_failed_request_without_the_auth_object_still_reads_the_committed_org_spend( + prometheus_logger, getters, spend_counters +): + _stub_gauges(prometheus_logger) + _store_committed(spend_counters, key=3.0, team=30.0, user=8.0, org=65.0) + metadata = _request_metadata(_token_committed_at_auth()) + del metadata["user_api_key_auth"] + + await prometheus_logger.async_log_failure_event(_failed_request(metadata), None, KEY_RESET_AT, KEY_RESET_AT) + + prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(pytest.approx(435.0)) + + +@pytest.mark.asyncio +async def test_failed_request_without_a_carried_org_reads_no_committed_org_spend( + prometheus_logger, getters, spend_counters +): + _stub_gauges(prometheus_logger) + _store_committed(spend_counters, key=3.0, team=30.0, user=8.0, org=65.0) + token = _token_committed_at_auth().model_copy(update={"org_budget_snapshot": None}) + metadata = {key: value for key, value in _request_metadata(token).items() if key != "user_api_key_auth"} + + await prometheus_logger.async_log_failure_event(_failed_request(metadata), None, KEY_RESET_AT, KEY_RESET_AT) + + prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(pytest.approx(460.0)) + + +@pytest.mark.asyncio +async def test_failed_request_that_was_charged_counts_its_own_charge_while_it_waits_in_the_pipeline( + prometheus_logger, getters, redis_spend_counters +): + """A stream cut after 40 events still costs 2.0, and that charge is queued when the failure logs.""" + _stub_gauges(prometheus_logger) + _store_committed(redis_spend_counters, key=5.0, team=32.0, user=10.0, org=62.0) + + with _request_scope(): + await _declare_the_own_two_dollar_charge() + await prometheus_logger.async_log_failure_event( + _failed_request(_request_metadata(_token_committed_at_auth()), response_cost=2.0), + None, + KEY_RESET_AT, + KEY_RESET_AT, + ) + await flush_post_call_redis_batches() + + prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(pytest.approx(436.0)) + + +@pytest.mark.asyncio +async def test_ordinary_metadata_keys_named_like_snapshot_fields_change_no_gauge(prometheus_logger, getters): + _stub_gauges(prometheus_logger) + metadata = { + "user_api_key_team_spend": 20.0, + "user_api_key_team_max_budget": 300.0, + "spend": 1000.0, + "max_budget": 1.0, + "budget_reset_at": None, + "user_alias": "not-the-user", + } + + await _emit(prometheus_logger, metadata) + + assert {name: getter.await_count for name, getter in getters.items()} == { + "get_key_object": 1, + "get_team_object": 1, + "get_user_object": 1, + "get_org_object": 1, + } + prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(458.0) + + @pytest.mark.asyncio async def test_metadata_without_carried_state_still_fetches_each_object_once(prometheus_logger, getters): _stub_gauges(prometheus_logger) diff --git a/tests/unit/integrations/test_prometheus_user_team_metrics.py b/tests/unit/integrations/test_prometheus_user_team_metrics.py index ab0fb67b52f..acec1c882ef 100644 --- a/tests/unit/integrations/test_prometheus_user_team_metrics.py +++ b/tests/unit/integrations/test_prometheus_user_team_metrics.py @@ -1065,6 +1065,79 @@ async def test_set_customer_budget_metrics_after_api_request_uses_cached_default assert _customer_sample("litellm_customer_max_budget_metric", "cust-default") == pytest.approx(3.0) +@pytest.mark.asyncio +async def test_customer_gauge_subtracts_the_committed_spend_instead_of_the_cached_spend( + prometheus_logger, customer_metrics_enabled +): + """Auth saw 1.5 committed, and by emit the counter holds 2.5 with this request's 0.5 and another request's.""" + import sys + + from litellm.models.budget import LiteLLM_BudgetTable + from litellm.models.end_user import LiteLLM_EndUserTable + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys + from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters + + mock_proxy_server = MagicMock() + mock_proxy_server.user_api_key_cache.async_get_cache = AsyncMock( + return_value=LiteLLM_EndUserTable( + user_id="cust-live", + blocked=False, + spend=0.5, + litellm_budget_table=LiteLLM_BudgetTable(budget_id="b", max_budget=3.0), + ) + ) + mock_proxy_server.spend_counter_store = FakeSpendCounters( + redis=None, memory={"spend:committed:end_user:cust-live": 2.5} + ).store() + token = UserAPIKeyAuth(token="hashed") + token.committed_budget_spend = CommittedBudgetSpend( + counter_keys=EntitySpendCounterKeys(end_user="spend:end_user:cust-live"), end_user=1.5 + ) + + with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}): + await prometheus_logger._increment_remaining_budget_metrics( + user_api_team=None, + user_api_team_alias=None, + user_api_key=None, + user_api_key_alias=None, + litellm_params={"metadata": {"user_api_key_auth": token}}, + response_cost=0.5, + end_user_id="cust-live", + ) + + assert _customer_sample("litellm_remaining_customer_budget_metric", "cust-live") == pytest.approx(0.5) + + +@pytest.mark.asyncio +async def test_customer_gauge_counts_a_charge_the_cached_customer_row_already_holds_once( + prometheus_logger, customer_metrics_enabled +): + import sys + + from litellm.models.budget import LiteLLM_BudgetTable + from litellm.models.end_user import LiteLLM_EndUserTable + + mock_proxy_server = MagicMock() + mock_proxy_server.user_api_key_cache.async_get_cache = AsyncMock( + return_value=LiteLLM_EndUserTable( + user_id="cust-charged", + blocked=False, + spend=0.5, + litellm_budget_table=LiteLLM_BudgetTable(budget_id="b", max_budget=3.0), + ) + ) + + with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}): + await prometheus_logger._set_customer_budget_metrics_after_api_request( + end_user_id="cust-charged", + response_cost=0.5, + committed_spend=0.0, + ) + + assert _customer_sample("litellm_remaining_customer_budget_metric", "cust-charged") == pytest.approx(2.5) + + @pytest.mark.asyncio async def test_set_customer_budget_metrics_after_api_request_without_budget_only_emits_remaining( prometheus_logger, customer_metrics_enabled diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_utils.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_utils.py index 1d2a5e4cc9b..f007d3740aa 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_utils.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_utils.py @@ -1,4 +1,5 @@ import asyncio +import math import uuid from collections.abc import Mapping from types import MappingProxyType @@ -483,3 +484,42 @@ async def test_router_prices_balanced_columns_by_the_same_rules_as_flex( assert balanced_cost == pytest.approx(flex_cost) if not overridden_columns: assert balanced_cost == pytest.approx(_cost_at(PARITY_ROW, "_balanced")) + + +def _vertex_cost_with_cache_writes(model: str, prompt_tokens: int, cache_creation_tokens: int) -> float: + response: Final = ModelResponse( + model=model, + usage=Usage( + prompt_tokens=prompt_tokens, + completion_tokens=100, + total_tokens=prompt_tokens + 100, + cache_creation_input_tokens=cache_creation_tokens, + ), + ) + return litellm.completion_cost(completion_response=response, model=model, custom_llm_provider="vertex_ai") + + +@pytest.mark.parametrize( + "model", + ( + "gemini-2.5-pro", + "gemini-3.1-pro-preview", + "gemini-3.1-pro-preview-customtools", + "vertex_ai/gemini-3-pro-preview", + "vertex_ai/gemini-3.1-pro-preview", + "vertex_ai/gemini-3.1-pro-preview-customtools", + ), +) +def test_vertex_gemini_pro_bills_cache_writes_above_200k_at_the_above_200k_input_rate( + local_model_cost_map: None, model: str +) -> None: + rates: Final = litellm.get_model_info(model=model, custom_llm_provider="vertex_ai") + input_rate_above_200k: Final = rates.get("input_cost_per_token_above_200k_tokens") + output_rate_above_200k: Final = rates.get("output_cost_per_token_above_200k_tokens") + cache_creation_tokens: Final = 150_000 + below: Final = _vertex_cost_with_cache_writes(model, 190_000, cache_creation_tokens) + above: Final = _vertex_cost_with_cache_writes(model, 210_000, cache_creation_tokens) + + assert input_rate_above_200k is not None and output_rate_above_200k is not None + assert above > below + assert math.isclose(above, 210_000 * input_rate_above_200k + 100 * output_rate_above_200k) diff --git a/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py b/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py index 6ebbdcf73bf..d020c7fcd85 100644 --- a/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py +++ b/tests/unit/litellm_core_utils/test_allowed_tools_tool_choice.py @@ -4,10 +4,13 @@ from typing import Final import pytest +from litellm.exceptions import BadRequestError from litellm.litellm_core_utils.allowed_tools_tool_choice import ( ChatAllowedToolsToolChoice, + hosted_tool_choice_type, is_allowed_tools_tool_choice, parse_chat_allowed_tools_tool_choice, + raise_if_hosted_tool_choice_unsupported, responses_allowed_tools_tool_choice, ) @@ -98,3 +101,42 @@ def test_parse_rejects_anything_but_the_chat_layout(tool_choice: dict[str, objec ) def test_is_allowed_tools_tool_choice(tool_choice: Mapping[str, object] | str | None, expected: bool): assert is_allowed_tools_tool_choice(tool_choice) is expected + + +@pytest.mark.parametrize( + ("tool_choice", "expected"), + [ + ({"type": "web_search"}, "web_search"), + ({"type": "mcp", "server_label": "docs"}, "mcp"), + ({"type": "function", "function": {"name": "get_time"}}, None), + ({"type": "allowed_tools"}, None), + ("required", None), + (None, None), + ], +) +def test_hosted_tool_choice_type(tool_choice: Mapping[str, object] | str | None, expected: str | None): + assert hosted_tool_choice_type(tool_choice) == expected + + +@pytest.mark.parametrize("provider", ["openai", "azure", "litellm_proxy"]) +def test_hosted_tool_choice_passes_where_the_responses_api_runs_the_tool(provider: str): + assert ( + raise_if_hosted_tool_choice_unsupported({"type": "web_search"}, custom_llm_provider=provider, model="m") is None + ) + + +@pytest.mark.parametrize("provider", ["anthropic", "bedrock", "gemini", "vertex_ai"]) +def test_hosted_tool_choice_is_a_400_where_the_provider_cannot_run_it(provider: str): + with pytest.raises(BadRequestError, match=r"tool_choice type 'file_search'.*only supported for") as exc: + raise_if_hosted_tool_choice_unsupported({"type": "file_search"}, custom_llm_provider=provider, model="m") + assert exc.value.status_code == 400 + assert f"custom_llm_provider={provider}" in str(exc.value) + + +def test_a_function_tool_choice_is_not_a_hosted_tool_anywhere(): + assert ( + raise_if_hosted_tool_choice_unsupported( + {"type": "function", "function": {"name": "get_time"}}, custom_llm_provider="anthropic", model="m" + ) + is None + ) diff --git a/tests/unit/litellm_core_utils/test_cli_token_utils.py b/tests/unit/litellm_core_utils/test_cli_token_utils.py index 7e7eee5373f..68d3098cc91 100644 --- a/tests/unit/litellm_core_utils/test_cli_token_utils.py +++ b/tests/unit/litellm_core_utils/test_cli_token_utils.py @@ -31,11 +31,13 @@ from litellm.litellm_core_utils.cli_token_utils import ( CredentialNotCleared, CredentialNotRecorded, CredentialNotSaved, + LoginRecordNotMoved, clear_cli_token, get_cli_token_file_path, get_litellm_gateway_api_key, is_cli_token_fresh, load_cli_token, + load_cli_token_left_in_keychain, save_cli_token, ) @@ -54,6 +56,10 @@ def _token_file(home): return home / ".litellm" / "token.json" +def _everything_on_disk(home): + return "".join(path.read_text() for path in (home / ".litellm").rglob("*") if path.is_file()) + + def _write_legacy_file(home, **overrides): payload = { "base_url": SERVER, @@ -123,6 +129,10 @@ class _MkstempThatNeedsTheOldFileGone: return _REAL_MKSTEMP(*args, **kwargs) +def _full_disk_mkstemp(*args, **kwargs): + raise OSError(errno.ENOSPC, "No space left on device") + + _REAL_REPLACE = os.replace @@ -777,6 +787,17 @@ class TestClearCliToken: assert clear_cli_token(vault=vault) == SecretStranded() assert load_cli_token(vault=vault) is None + def test_a_keychain_that_kept_the_secret_is_still_warned_about_once_it_is_out_of_reach( + self, isolated_home, secret_vault_factory + ): + _write_metadata_only_file(isolated_home) + vault = secret_vault_factory(blob=_blob(), erasable=False) + assert clear_cli_token(vault=vault) == SecretStranded() + + vault.available = False + + assert clear_cli_token(vault=vault) == KeyringUnreachable() + @pytest.mark.parametrize( "failure", [KeyringDisabled(), KeyringUnreachable(), KeyringNotInstalled()] ) @@ -794,7 +815,8 @@ class TestClearCliToken: vault = secret_vault_factory(available=False, failure=failure) assert clear_cli_token(vault=vault) == failure - assert json.loads(_token_file(isolated_home).read_text()).get("key") is None + assert not _token_file(isolated_home).exists() + assert "sk-legacy" not in _everything_on_disk(isolated_home) def test_a_second_logout_still_reports_the_keychain_it_could_not_clear( self, isolated_home, secret_vault_factory @@ -817,20 +839,20 @@ class TestClearCliToken: vault = secret_vault_factory(available=False, failure=KeyringNotInstalled()) assert clear_cli_token(vault=vault) == KeyringNotInstalled() - assert json.loads(_token_file(isolated_home).read_text()).get("key") is None + assert load_cli_token(vault=vault) is None - def test_a_logout_that_cannot_clear_the_keychain_keeps_the_record_that_it_has_to( - self, isolated_home, secret_vault_factory + @pytest.mark.parametrize("failure", [KeyringDisabled(), KeyringUnreachable(), KeyringNotInstalled()]) + def test_a_logout_that_cannot_check_the_keychain_still_ends_the_local_login( + self, isolated_home, secret_vault_factory, failure ): - """The file left behind holds no secret. It is what a later run reads to tell a machine with - a credential it cannot reach apart from one that never had a login, which is the difference - between warning the user and inventing a credential for them to worry about.""" - _write_metadata_only_file(isolated_home) - vault = secret_vault_factory(available=False, failure=KeyringUnreachable()) + vault = secret_vault_factory(available=False, failure=failure) + save_cli_token(CliTokenRecord(base_url=SERVER, key="sk-in-file", user_id="u-1"), vault=vault) - clear_cli_token(vault=vault) + assert clear_cli_token(vault=vault) == failure - assert json.loads(_token_file(isolated_home).read_text()).get("key") is None + assert load_cli_token(vault=vault) is None + assert get_litellm_gateway_api_key(vault=vault) is None + assert not _token_file(isolated_home).exists() def test_a_logout_that_cannot_clear_the_keychain_still_takes_the_file_secret_away( self, isolated_home, secret_vault_factory @@ -842,7 +864,7 @@ class TestClearCliToken: clear_cli_token(vault=vault) - left_on_disk = _token_file(isolated_home).read_text() + left_on_disk = _everything_on_disk(isolated_home) assert "sk-legacy" not in left_on_disk assert "rt-legacy" not in left_on_disk @@ -862,7 +884,91 @@ class TestClearCliToken: assert clear_cli_token(vault=vault) == KeyringUnreachable() assert clear_cli_token(vault=vault) == KeyringUnreachable() assert vault.blob is not None - assert "sk-second" not in _token_file(isolated_home).read_text() + assert load_cli_token(vault=vault) is None + assert "sk-second" not in _everything_on_disk(isolated_home) + + def test_a_keychain_that_answers_again_never_signs_an_ended_login_back_in( + self, isolated_home, secret_vault_factory + ): + vault = secret_vault_factory() + save_cli_token(CliTokenRecord(base_url=SERVER, key="sk-first"), vault=vault) + vault.available = False + save_cli_token(CliTokenRecord(base_url=SERVER, key="sk-second"), vault=vault) + assert clear_cli_token(vault=vault) == KeyringUnreachable() + + vault.available = True + + assert load_cli_token(vault=vault) is None + assert get_litellm_gateway_api_key(vault=vault) is None + assert clear_cli_token(vault=vault) == SecretErased() + assert vault.blob is None + + def test_a_keychain_confirmed_clean_is_not_warned_about_again(self, isolated_home, secret_vault_factory): + vault = secret_vault_factory() + save_cli_token(CliTokenRecord(base_url=SERVER, key="sk-first"), vault=vault) + vault.available = False + save_cli_token(CliTokenRecord(base_url=SERVER, key="sk-second"), vault=vault) + assert clear_cli_token(vault=vault) == KeyringUnreachable() + vault.available = True + assert clear_cli_token(vault=vault) == SecretErased() + + vault.available = False + + assert clear_cli_token(vault=vault) == SecretErased() + + def test_a_second_logout_that_misses_the_keychain_keeps_the_first_login_revocable( + self, isolated_home, secret_vault_factory + ): + vault = secret_vault_factory() + save_cli_token( + CliTokenRecord(base_url=SERVER, key="sk-first", refresh_token="rt-first", timestamp=1000.0), vault=vault + ) + vault.available = False + assert clear_cli_token(vault=vault) == KeyringUnreachable() + save_cli_token( + CliTokenRecord(base_url=SERVER, key="sk-second", refresh_token="rt-second", timestamp=2000.0), vault=vault + ) + assert clear_cli_token(vault=vault) == KeyringUnreachable() + + vault.available = True + left = load_cli_token_left_in_keychain(vault=vault) + + assert left is not None + assert (left.key, left.refresh_token) == ("sk-first", "rt-first") + + def test_a_locked_login_over_a_keychain_login_keeps_the_keychain_login_revocable( + self, isolated_home, secret_vault_factory + ): + vault = secret_vault_factory() + save_cli_token( + CliTokenRecord(base_url=SERVER, key="sk-first", refresh_token="rt-first", timestamp=1000.0), vault=vault + ) + vault.available = False + save_cli_token( + CliTokenRecord(base_url=SERVER, key="sk-second", refresh_token="rt-second", timestamp=2000.0), vault=vault + ) + + vault.available = True + left = load_cli_token_left_in_keychain(vault=vault) + + assert left is not None + assert (left.key, left.refresh_token) == ("sk-first", "rt-first") + + def test_the_note_never_pairs_with_a_keychain_entry_a_later_login_wrote( + self, isolated_home, secret_vault_factory + ): + vault = secret_vault_factory() + save_cli_token( + CliTokenRecord(base_url=SERVER, key="sk-first", refresh_token="rt-first", timestamp=1000.0), vault=vault + ) + vault.available = False + assert clear_cli_token(vault=vault) == KeyringUnreachable() + vault.available = True + save_cli_token( + CliTokenRecord(base_url=SERVER, key="sk-second", refresh_token="rt-second", timestamp=2000.0), vault=vault + ) + + assert load_cli_token_left_in_keychain(vault=vault) is None @pytest.mark.parametrize("failure", [KeyringNotInstalled(), KeyringDisabled(), KeyringUnreachable()]) def test_logging_out_of_a_machine_that_never_logged_in_invents_nothing_to_warn_about( @@ -874,6 +980,7 @@ class TestClearCliToken: vault = secret_vault_factory(available=False, failure=failure) assert clear_cli_token(vault=vault) == SecretErased() + assert clear_cli_token(vault=vault) == SecretErased() def test_a_file_backed_login_cannot_vouch_for_a_keychain_no_package_can_reach( self, isolated_home, secret_vault_factory @@ -890,7 +997,8 @@ class TestClearCliToken: assert clear_cli_token(vault=vault) == KeyringNotInstalled() assert vault.blob is not None - assert "sk-in-file" not in _token_file(isolated_home).read_text() + assert not _token_file(isolated_home).exists() + assert "sk-in-file" not in _everything_on_disk(isolated_home) @pytest.mark.skipif(os.geteuid() == 0, reason="root ignores file permissions") def test_a_file_that_gives_up_neither_its_secret_nor_itself_is_reported_not_raised( @@ -917,10 +1025,10 @@ class TestClearCliToken: def test_a_directory_that_takes_no_new_file_still_gives_up_the_secret_in_the_old_one( self, isolated_home, secret_vault_factory ): - """A read-only `~/.litellm` accepts no replacement token file and no removal of the one it - has, and still lets that one be shortened. The secret goes, the file stays as the note that - the keychain went unchecked, and the logout after it warns again instead of reading the gap - the removal would have left as a clean keychain. + """A read-only `~/.litellm` accepts no removal of the token file and no new file beside it, + and still lets the one it has be shortened. The secret goes, the file stays because nothing + else can, and the logout after it says the file is still there, since a keychain that answers + again would sign it back in, rather than calling the machine logged out. The key is a realistic length so the file genuinely shrinks: a rewrite in place that leaves the tail of the old contents behind hands the next run a file it cannot parse.""" @@ -929,25 +1037,16 @@ class TestClearCliToken: path.parent.chmod(0o500) try: assert clear_cli_token(vault=vault) == KeyringUnreachable() - assert clear_cli_token(vault=vault) == KeyringUnreachable() + assert isinstance(clear_cli_token(vault=vault), LoginRecordNotMoved) finally: path.parent.chmod(0o700) assert json.loads(path.read_text()).get("key") is None - @pytest.mark.skipif(os.geteuid() == 0, reason="root ignores file permissions") - def test_a_note_the_logout_had_to_remove_is_written_again_for_the_next_one( + def test_a_full_disk_still_keeps_the_unchecked_keychain_on_record( self, isolated_home, secret_vault_factory, monkeypatch ): - """A full disk refuses the replacement file and a read-only token file refuses the rewrite - in place, so the only way left to get the secret off disk is to remove the file carrying it. - That file was also the note saying the keychain went unchecked, and its absence is what the - next logout would read as a keychain already known to be clean. - - Removing it is what frees the room the replacement was refused for, so the note is written - again on the way out and the logout after this one still warns.""" path = _write_legacy_file(isolated_home) - path.chmod(0o400) monkeypatch.setattr( "litellm.litellm_core_utils.private_json.tempfile.mkstemp", _MkstempThatNeedsTheOldFileGone(path), @@ -957,7 +1056,26 @@ class TestClearCliToken: assert clear_cli_token(vault=vault) == KeyringUnreachable() assert clear_cli_token(vault=vault) == KeyringUnreachable() - assert json.loads(path.read_text()).get("key") is None + assert not path.exists() + assert "sk-legacy" not in _everything_on_disk(isolated_home) + + def test_a_full_disk_keeps_a_locked_keychain_login_revocable_without_signing_it_back_in( + self, isolated_home, secret_vault_factory, monkeypatch + ): + vault = secret_vault_factory() + save_cli_token( + CliTokenRecord(base_url=SERVER, key="sk-first", refresh_token="rt-first", timestamp=1000.0), vault=vault + ) + monkeypatch.setattr("litellm.litellm_core_utils.private_json.tempfile.mkstemp", _full_disk_mkstemp) + vault.available = False + + assert clear_cli_token(vault=vault) == KeyringUnreachable() + + vault.available = True + left = load_cli_token_left_in_keychain(vault=vault) + assert left is not None + assert left.refresh_token == "rt-first" + assert load_cli_token(vault=vault) is None @pytest.mark.skipif(os.geteuid() == 0, reason="root ignores directory permissions") def test_a_metadata_file_that_will_not_go_is_not_worth_alarming_the_user_over( diff --git a/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py b/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py index de252b17f00..b4e3ef2c566 100644 --- a/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py +++ b/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py @@ -30,9 +30,9 @@ CUSTOM_PROVIDER: Final = "test-onprem-llm" @pytest.mark.parametrize( ("model", "provider"), - (("openai/gpt-luna-latest", None), ("gpt-luna-latest", "openai")), + (("openai/gpt-luna-latest", None), ("gpt-luna-latest", "openai"), ("gpt-luna-latest", None)), ) -def test_latest_alias_resolves_to_the_bundled_target(model: str, provider: str | None) -> None: +def test_latest_alias_resolves_to_the_cost_map_target(model: str, provider: str | None) -> None: target: Final = get_model_latest_aliases()["openai"]["gpt-luna-latest"] resolved, resolved_provider, _, api_base = get_llm_provider( model, custom_llm_provider=provider, api_base="https://api.openai.com/v1" @@ -42,7 +42,12 @@ def test_latest_alias_resolves_to_the_bundled_target(model: str, provider: str | @pytest.mark.parametrize( "model", - ("gemini/gemini-pro-latest", "gemini/gemini-flash-latest", "dashscope/qwen-plus-latest", "mistral/mistral-large-latest"), + ( + "gemini/gemini-pro-latest", + "gemini/gemini-flash-latest", + "dashscope/qwen-plus-latest", + "mistral/mistral-large-latest", + ), ) def test_provider_native_latest_alias_passes_through(model: str) -> None: provider, name = model.split("/", 1) @@ -153,6 +158,7 @@ def _vcr_outcome_gate(request, vcr): yield record_vcr_outcome(request, vcr) + @pytest.fixture(scope="function") def isolate_litellm_state(): """ @@ -205,6 +211,7 @@ def isolate_litellm_state(): setattr(litellm, attr, original_value) _invalidate_model_cost_lowercase_map() + _SCALAR_DEFAULTS = { "num_retries": getattr(litellm, "num_retries", None), "num_retries_per_request": getattr(litellm, "num_retries_per_request", None), @@ -225,6 +232,7 @@ _SCALAR_DEFAULTS = { "api_key": getattr(litellm, "api_key", None), } + @pytest.fixture(scope="module") def setup_and_teardown(): """ @@ -247,12 +255,14 @@ def setup_and_teardown(): litellm.in_memory_llm_clients_cache.flush_cache() yield + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_get_llm_provider(): _, response, _, _ = litellm.get_llm_provider(model="anthropic.claude-v2:1") assert response == "bedrock" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_get_llm_provider_fireworks(): # tests finetuned fireworks models - https://github.com/BerriAI/litellm/issues/4923 model, custom_llm_provider, _, _ = litellm.get_llm_provider(model="fireworks_ai/accounts/my-test-1234") @@ -260,17 +270,20 @@ def test_get_llm_provider_fireworks(): # tests finetuned fireworks models - htt assert custom_llm_provider == "fireworks_ai" assert model == "accounts/my-test-1234" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_get_llm_provider_catch_all(): _, response, _, _ = litellm.get_llm_provider(model="*") assert response == "openai" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_get_llm_provider_gpt_instruct(): _, response, _, _ = litellm.get_llm_provider(model="gpt-3.5-turbo-instruct-0914") assert response == "text-completion-openai" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_get_llm_provider_mistral_custom_api_base(): model, custom_llm_provider, dynamic_api_key, api_base = litellm.get_llm_provider( @@ -281,6 +294,7 @@ def test_get_llm_provider_mistral_custom_api_base(): assert model == "mistral-large-fr" assert api_base == "https://mistral-large-fr-ishaan.francecentral.inference.ai.azure.com/v1" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_get_llm_provider_deepseek_custom_api_base(): os.environ["DEEPSEEK_API_BASE"] = "MY-FAKE-BASE" @@ -293,6 +307,7 @@ def test_get_llm_provider_deepseek_custom_api_base(): os.environ.pop("DEEPSEEK_API_BASE") + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_get_llm_provider_vertex_ai_image_models(monkeypatch): monkeypatch.setattr(litellm, "vertex_ai_image_models", set()) @@ -310,6 +325,7 @@ def test_get_llm_provider_vertex_ai_image_models(monkeypatch): ) assert custom_llm_provider == "vertex_ai" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_get_llm_provider_ai21_chat(): model, custom_llm_provider, dynamic_api_key, api_base = litellm.get_llm_provider( @@ -319,6 +335,7 @@ def test_get_llm_provider_ai21_chat(): assert model == "jamba-1.5-large" assert api_base == "https://api.ai21.com/studio/v1" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_get_llm_provider_ai21_chat_test2(): """ @@ -335,6 +352,7 @@ def test_get_llm_provider_ai21_chat_test2(): assert model == "jamba-1.5-large" assert api_base == "https://api.ai21.com/studio/v1" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_get_llm_provider_cohere_chat_test2(): """ @@ -350,6 +368,7 @@ def test_get_llm_provider_cohere_chat_test2(): assert custom_llm_provider == "cohere_chat" assert model == "command-r-plus-08-2024" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_get_llm_provider_azure_o1(): @@ -359,6 +378,7 @@ def test_get_llm_provider_azure_o1(): assert custom_llm_provider == "azure" assert model == "o1-mini" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_hosted_vllm_default_api_key(): from litellm.litellm_core_utils.get_llm_provider_logic import ( @@ -373,6 +393,7 @@ def test_hosted_vllm_default_api_key(): ) assert dynamic_api_key == "fake-api-key" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_get_llm_provider_jina_ai(): model, custom_llm_provider, dynamic_api_key, api_base = litellm.get_llm_provider( @@ -382,6 +403,7 @@ def test_get_llm_provider_jina_ai(): assert api_base == "https://api.jina.ai/v1" assert model == "jina-embeddings-v3" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_get_llm_provider_hosted_vllm(): model, custom_llm_provider, dynamic_api_key, api_base = litellm.get_llm_provider( @@ -391,6 +413,7 @@ def test_get_llm_provider_hosted_vllm(): assert model == "llama-3.1-70b-instruct" assert dynamic_api_key == "fake-api-key" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_get_llm_provider_llamafile(): model, custom_llm_provider, dynamic_api_key, api_base = litellm.get_llm_provider( @@ -401,6 +424,7 @@ def test_get_llm_provider_llamafile(): assert dynamic_api_key == "fake-api-key" assert api_base == "http://127.0.0.1:8080/v1" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_get_llm_provider_watson_text(): model, custom_llm_provider, dynamic_api_key, api_base = litellm.get_llm_provider( @@ -409,6 +433,7 @@ def test_get_llm_provider_watson_text(): assert custom_llm_provider == "watsonx_text" assert model == "watson-text-to-speech" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_azure_global_standard_get_llm_provider(): model, custom_llm_provider, dynamic_api_key, api_base = litellm.get_llm_provider( @@ -418,6 +443,7 @@ def test_azure_global_standard_get_llm_provider(): ) assert custom_llm_provider == "azure_ai" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_nova_bedrock_converse(): model, custom_llm_provider, dynamic_api_key, api_base = litellm.get_llm_provider( @@ -426,6 +452,7 @@ def test_nova_bedrock_converse(): assert custom_llm_provider == "bedrock" assert model == "amazon.nova-micro-v1:0" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_bedrock_invoke_anthropic(): model, custom_llm_provider, dynamic_api_key, api_base = litellm.get_llm_provider( @@ -434,6 +461,7 @@ def test_bedrock_invoke_anthropic(): assert custom_llm_provider == "bedrock" assert model == "invoke/anthropic.claude-haiku-4-5-20251001-v1:0" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.parametrize("model", ["xai/grok-2-vision-latest", "grok-2-vision-latest"]) def test_xai_api_base(model): @@ -450,6 +478,7 @@ def test_xai_api_base(model): assert api_base == "https://api.x.ai/v1" assert dynamic_api_key == "xai-my-specialkey" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_get_litellm_proxy_custom_llm_provider(): """ @@ -479,6 +508,7 @@ def test_get_litellm_proxy_custom_llm_provider(): assert key == expected_api_key assert base == expected_api_base + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_get_litellm_proxy_with_args_override_env_vars(): """ @@ -510,6 +540,7 @@ def test_get_litellm_proxy_with_args_override_env_vars(): assert key == arg_api_key assert base == arg_api_base + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_get_litellm_proxy_model_prefix_stripping(): """ @@ -540,6 +571,7 @@ def test_get_litellm_proxy_model_prefix_stripping(): assert key == expected_api_key assert base == expected_api_base + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_get_llm_provider_LITELLM_PROXY_ALWAYS_true(): """ @@ -568,6 +600,7 @@ def test_get_llm_provider_LITELLM_PROXY_ALWAYS_true(): assert key == proxy_api_key assert base == proxy_api_base + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_get_llm_provider_LITELLM_PROXY_ALWAYS_true_model_prefix(): """ @@ -594,6 +627,7 @@ def test_get_llm_provider_LITELLM_PROXY_ALWAYS_true_model_prefix(): assert key == proxy_api_key assert base == proxy_api_base + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_get_llm_provider_use_proxy_arg_true(): """ @@ -623,6 +657,7 @@ def test_get_llm_provider_use_proxy_arg_true(): assert key == proxy_api_key assert base == proxy_api_base + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_get_llm_provider_use_proxy_arg_true_with_direct_args(): """ @@ -659,6 +694,7 @@ def test_get_llm_provider_use_proxy_arg_true_with_direct_args(): assert key == arg_api_key # Should use the argument key assert base == arg_api_base # Should use the argument base + @pytest.fixture def shipped_generalizations(): """Install the rules shipped in the bundled backup, then restore. @@ -682,6 +718,7 @@ def shipped_generalizations(): finally: set_fallback_generalizations(previous) + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") class TestClaudeModelPatternMatching: """ @@ -749,9 +786,7 @@ class TestClaudeModelPatternMatching: @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") class TestTypesafeJevPatternMatching: - def test_every_priced_jev_id_routes_to_typesafe( - self, shipped_generalizations: None - ) -> None: + def test_every_priced_jev_id_routes_to_typesafe(self, shipped_generalizations: None) -> None: from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap priced_ids: Final = tuple( @@ -760,9 +795,9 @@ class TestTypesafeJevPatternMatching: if key.startswith("typesafe/") ) assert "jev-latest" in priced_ids - assert { - model: litellm.get_llm_provider(model=model)[1] for model in priced_ids - } == dict.fromkeys(priced_ids, "typesafe") + assert {model: litellm.get_llm_provider(model=model)[1] for model in priced_ids} == dict.fromkeys( + priced_ids, "typesafe" + ) @pytest.mark.parametrize( "model", @@ -837,7 +872,9 @@ def test_metadata_callers_never_log_in_for_a_declared_chatgpt_deployment( assert get_api_base("chatgpt/gpt-5.4", {"model": "chatgpt/gpt-5.4"}) == "https://chatgpt.com/backend-api/codex" identity = _resolve_model("chatgpt/gpt-5.4", None) assert identity is not None and identity.provider == "chatgpt" - group = Router(model_list=[{"model_name": "codex", "litellm_params": {"model": "chatgpt/gpt-5.4"}}]).registry().group( - "codex" + group = ( + Router(model_list=[{"model_name": "codex", "litellm_params": {"model": "chatgpt/gpt-5.4"}}]) + .registry() + .group("codex") ) assert group is not None and group.providers == ["chatgpt"] diff --git a/tests/unit/litellm_core_utils/test_model_latest_aliases.py b/tests/unit/litellm_core_utils/test_model_latest_aliases.py index c052a24b4dc..c6def915640 100644 --- a/tests/unit/litellm_core_utils/test_model_latest_aliases.py +++ b/tests/unit/litellm_core_utils/test_model_latest_aliases.py @@ -3,14 +3,17 @@ from typing import Final import pytest import litellm +from litellm.litellm_core_utils.get_llm_provider_logic import PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO +from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry from litellm.litellm_core_utils.model_latest_aliases import ( expand_model_latest_alias, get_model_latest_aliases, is_model_latest_alias, + resolve_model_latest_alias, ) -@pytest.mark.parametrize("model", ("", "gpt-luna-latest", "openai/unknown", "gemini/gemini-pro-latest")) +@pytest.mark.parametrize("model", ("", "kimi-latest", "openai/unknown", "gemini/gemini-pro-latest")) def test_non_alias_expansion_keeps_the_requested_name(model: str) -> None: assert expand_model_latest_alias(model) == (model,) @@ -21,15 +24,56 @@ def test_alias_expansion_keeps_requested_and_target_names() -> None: assert expand_model_latest_alias(f"{provider}/{alias}") == (f"{provider}/{alias}", f"{provider}/{target}") -def test_latest_alias_targets_have_chat_pricing_without_shadowing_native_models() -> None: - for provider, aliases in get_model_latest_aliases().items(): - for alias, target in aliases.items(): +def test_bare_alias_expands_to_its_cost_map_owner() -> None: + alias, target = next(iter(get_model_latest_aliases()["openai"].items())) + assert expand_model_latest_alias(alias) == (alias, target) + + +def test_latest_aliases_price_and_route_like_their_targets() -> None: + aliases: Final = get_model_latest_aliases() + assert aliases + for provider, provider_aliases in aliases.items(): + for alias, target in provider_aliases.items(): + prefixed: Final = f"{provider}/{alias}" + alias_key: Final = prefixed if prefixed in litellm.model_cost else alias target_key: Final = f"{provider}/{target}" if f"{provider}/{target}" in litellm.model_cost else target - entry: Final = litellm.model_cost[target_key] - assert (entry["litellm_provider"], entry["mode"]) == (provider, "chat") - assert f"{provider}/{alias}" not in litellm.model_cost - assert litellm.model_cost.get(alias, {}).get("litellm_provider") != provider - assert litellm.get_llm_provider(f"{provider}/{alias}")[:2] == (target, provider) + entry: Final = litellm.model_cost[alias_key] + assert entry["mode"] in ("chat", "responses") + assert entry == litellm.model_cost[target_key] + assert resolve_model_latest_alias(alias, provider) == target + if provider in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO: + continue + assert litellm.get_llm_provider(prefixed)[:2] == (target, provider) + assert litellm.get_llm_provider(alias, custom_llm_provider=provider)[:2] == (target, provider) + + +@pytest.mark.parametrize( + ("model", "provider", "target_provider"), + ( + ("bedrock/claude-haiku-latest", "bedrock", "bedrock_converse"), + ("vertex_ai/claude-haiku-latest", "vertex_ai", "vertex_ai-anthropic_models"), + ("azure/gpt-mini-latest", "azure", "azure"), + ("fireworks_ai/kimi-latest", "fireworks_ai", "fireworks_ai"), + ), +) +def test_hosted_alias_is_sent_as_the_hosts_model_id(model: str, provider: str, target_provider: str) -> None: + entry: Final = litellm.model_cost[model] + assert isinstance(entry, AliasedCostEntry) + name: Final = model.partition("/")[2] + target: Final = entry.names[0].removeprefix(f"{provider}/") + assert target != name + assert entry["litellm_provider"] == target_provider + assert litellm.get_llm_provider(model)[:2] == (target, provider) + assert is_model_latest_alias(model) + assert resolve_model_latest_alias(name, provider) == target + assert get_model_latest_aliases()[provider][name] == target + + +def test_alias_resolution_is_scoped_to_the_owning_provider() -> None: + alias, target = next(iter(get_model_latest_aliases()["openai"].items())) + assert resolve_model_latest_alias(alias, "openai") == target + assert resolve_model_latest_alias(alias, "anthropic") == alias + assert resolve_model_latest_alias(target, "openai") == target @pytest.mark.parametrize( @@ -37,11 +81,14 @@ def test_latest_alias_targets_have_chat_pricing_without_shadowing_native_models( ( ("openai/gpt-luna-latest", None, True), ("gpt-luna-latest", "openai", True), - ("gpt-luna-latest", None, False), + ("gpt-luna-latest", None, True), + ("gpt-luna-latest", "anthropic", False), + ("moonshot/kimi-latest", None, True), + ("kimi-latest", None, False), ("gemini/gemini-pro-latest", None, False), ("", None, False), ("openai/unknown-model", None, False), ), ) -def test_latest_alias_detection_requires_an_explicit_provider(model: str, provider: str | None, expected: bool) -> None: +def test_latest_alias_detection_follows_the_cost_map_owner(model: str, provider: str | None, expected: bool) -> None: assert is_model_latest_alias(model, provider) is expected diff --git a/tests/unit/litellm_core_utils/test_total_deadline.py b/tests/unit/litellm_core_utils/test_total_deadline.py new file mode 100644 index 00000000000..f3bc99354ea --- /dev/null +++ b/tests/unit/litellm_core_utils/test_total_deadline.py @@ -0,0 +1,36 @@ +from collections.abc import Mapping + +import pytest + +from litellm.litellm_core_utils.total_deadline import InvalidTotalTimeout, TotalTimeout, total_timeout_from_request + + +@pytest.mark.parametrize( + ("header", "body", "expected"), + [ + (None, {}, None), + (None, {"total_timeout": 10}, TotalTimeout(10.0, from_body=True)), + (None, {"total_timeout": "2.5"}, TotalTimeout(2.5, from_body=True)), + ("4", {}, TotalTimeout(4.0, from_body=False)), + ("4", {"total_timeout": 10}, TotalTimeout(4.0, from_body=False)), + ("4", {"total_timeout": "bad"}, TotalTimeout(4.0, from_body=False)), + ], + ids=["absent", "body", "body_numeric_string", "header", "header_wins", "header_wins_over_a_bad_body"], +) +def test_total_timeout_is_read_from_the_header_first_then_the_body( + header: str | None, body: Mapping[str, object], expected: TotalTimeout | None +) -> None: + assert total_timeout_from_request(header, body) == expected + + +@pytest.mark.parametrize( + "value", + [None, True, False, 0, -1, "0", "-2", "nan", "inf", "-inf", float("nan"), 10**400, "ten", "", [10], {"s": 10}], +) +def test_a_total_timeout_that_is_not_a_positive_finite_number_of_seconds_is_invalid(value: object) -> None: + assert isinstance(total_timeout_from_request(None, {"total_timeout": value}), InvalidTotalTimeout) + + +@pytest.mark.parametrize("value", ["0", "-1", "nan", "inf", "ten", ""]) +def test_an_invalid_total_timeout_header_is_invalid(value: str) -> None: + assert isinstance(total_timeout_from_request(value, {}), InvalidTotalTimeout) diff --git a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_transformation.py b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_transformation.py index 54b29ca6aa7..8f86888728f 100644 --- a/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_transformation.py +++ b/tests/unit/llms/anthropic/messages/native/test_anthropic_messages_transformation.py @@ -721,3 +721,17 @@ def test_a_first_party_messages_request_never_fetches_the_remote_beta_headers_co assert request["messages"][1] == {"role": "system", "content": "be terse"} assert not remote_config.called + + +@pytest.mark.parametrize( + "model, expected", + [ + ("claude-sonnet-5", {"type": "disabled"}), + ("claude-haiku-4-5", None), + ("claude-sonnet-5-5", None), + ], +) +def test_none_disables_only_optional_default_thinking(model: str, expected: dict[str, str] | None) -> None: + params: Final[dict[str, object]] = {"reasoning_effort": "none", "output_config": {"effort": "high"}} + AnthropicMessagesConfig._translate_reasoning_effort_to_anthropic(model, params, 1024, "anthropic") + assert params == ({"thinking": expected} if expected is not None else {}) diff --git a/tests/unit/llms/anthropic/test_anthropic_reasoning_effort.py b/tests/unit/llms/anthropic/test_anthropic_reasoning_effort.py index 573f586de1e..1b066e49e3b 100644 --- a/tests/unit/llms/anthropic/test_anthropic_reasoning_effort.py +++ b/tests/unit/llms/anthropic/test_anthropic_reasoning_effort.py @@ -5,11 +5,42 @@ Verifies that reasoning_effort=None returns None for all models, including Claude Opus 4.6. """ +from typing import Final + import pytest +from litellm.exceptions import BadRequestError from litellm.llms.anthropic.chat.transformation import AnthropicConfig +@pytest.mark.parametrize( + "model, expected", + [ + ("claude-sonnet-5", {"type": "disabled"}), + ("claude-opus-5", {"type": "disabled"}), + ("claude-haiku-5-5", {"type": "disabled"}), + ("claude-haiku-4-5", None), + ("claude-sonnet-5-5", None), + ("claude-opus-5-5", None), + ("claude-fable-5", None), + ("claude-mythos-5", None), + ], +) +def test_none_disables_only_optional_default_thinking(model: str, expected: dict[str, str] | None) -> None: + mapped: Final = AnthropicConfig().map_openai_params({"reasoning_effort": "none"}, {}, model, False) + assert mapped.get("thinking") == expected + assert "output_config" not in mapped + + +@pytest.mark.parametrize("effort", [None, "bogus"]) +def test_default_thinking_preserves_absent_and_invalid_effort(effort: str | None) -> None: + if effort is None: + assert AnthropicConfig.map_reasoning_effort(effort, "claude-sonnet-5", "anthropic") is None + else: + with pytest.raises(BadRequestError, match="Invalid reasoning_effort"): + AnthropicConfig().map_openai_params({"reasoning_effort": effort}, {}, "claude-sonnet-5", False) + + class TestMapReasoningEffort: def test_none_returns_none_for_opus_4_6(self): """reasoning_effort=None should return None for Opus 4.6, not adaptive.""" diff --git a/tests/unit/llms/anthropic/test_anthropic_wif.py b/tests/unit/llms/anthropic/test_anthropic_wif.py index 733e1472a8a..acaa7c7c8bf 100644 --- a/tests/unit/llms/anthropic/test_anthropic_wif.py +++ b/tests/unit/llms/anthropic/test_anthropic_wif.py @@ -452,6 +452,31 @@ class TestResolutionMatrix: assert params.service_account_id is None assert params.workspace_id is None + def test_env_inline_oidc_reference_resolves_like_the_param(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("ANTHROPIC_FEDERATION_RULE_ID", "fdrl_env") + monkeypatch.setenv("ANTHROPIC_ORGANIZATION_ID", "org-env") + monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN", "oidc/google/https://api.anthropic.com") + poster = ScriptedPoster([token_response()]) + engine = JwtBearerTokenExchangeEngine( + poster=poster, + assertion_reader={"oidc/google/https://api.anthropic.com": "google-id-token"}.get, + clock=FakeClock(), + refresh_executor=ManualExecutor(), + ) + + token = get_anthropic_wif_token(None, None, "claude-opus-4-6", engine) + + assert token == "sk-ant-oat01-minted" + assert [request.json_body()["assertion"] for request in poster.requests] == ["google-id-token"] + + def test_env_inline_env_path_reference_is_rejected(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("ANTHROPIC_FEDERATION_RULE_ID", "fdrl_env") + monkeypatch.setenv("ANTHROPIC_ORGANIZATION_ID", "org-env") + monkeypatch.setenv("ANTHROPIC_IDENTITY_TOKEN", "oidc/env_path/ANTHROPIC_TOKEN_PATH") + + with pytest.raises(litellm.AuthenticationError): + resolve_anthropic_wif_params(None) + def test_file_param_beats_inline_param(self): params = resolve_anthropic_wif_params( { diff --git a/tests/unit/llms/base_llm/auth/test_shared_token_store.py b/tests/unit/llms/base_llm/auth/test_shared_token_store.py index 298e61c4c79..5acea09d296 100644 --- a/tests/unit/llms/base_llm/auth/test_shared_token_store.py +++ b/tests/unit/llms/base_llm/auth/test_shared_token_store.py @@ -111,6 +111,44 @@ def test_a_minted_assertion_never_reaches_the_shared_store(tmp_path: Path): assert second.access_token.get_secret_value() != first.access_token.get_secret_value() +def test_a_google_metadata_assertion_is_exchanged_again_inside_the_mandatory_window(tmp_path: Path): + """A Google metadata-server token carries no ``jti``, so Anthropic exchanges it as often as asked. The + store used to answer a mandatory refresh with the stored token minted from that same assertion, which + was the 20-seconds-left token the refresh existed to replace.""" + poster = ScriptedPoster([token_response("sk-ant-oat01-first", expires_in=600), token_response("sk-ant-oat01-second")]) + clock = FakeClock() + google_ref = "oidc/google/https://api.anthropic.com" + engine = store_engine( + poster, FileTokenStore(tmp_path), reader={google_ref: "google-id-token"}, clock=clock, wall_clock=clock + ) + spec = make_spec(assertion_ref=google_ref) + + first = minted(engine.get_token(spec)) + clock.advance(580.0) + refreshed = minted(engine.get_token(spec)) + + assert refreshed.access_token.get_secret_value() == "sk-ant-oat01-second" + assert refreshed.expires_at - clock.now == pytest.approx(3600.0), (first.expires_at, refreshed.expires_at) + assert [request.json_body()["assertion"] for request in poster.requests] == ["google-id-token", "google-id-token"] + assert stored_files(tmp_path) == [], "a re-exchangeable assertion has no token to share across workers" + + +def test_a_token_file_assertion_still_serves_the_stored_token_inside_the_mandatory_window(tmp_path: Path): + """The projected token file is the case the store was built for: with the file unrotated, a second + exchange would only be refused as ``jti_reused``, so the stored token keeps serving until it dies.""" + poster = SingleUsePoster(expires_in=600) + clock = FakeClock() + store = FileTokenStore(tmp_path) + engine = store_engine(poster, store, clock=clock, wall_clock=clock) + + first = minted(engine.get_token(make_spec())) + clock.advance(580.0) + served = minted(store_engine(poster, store, clock=clock, wall_clock=clock).get_token(make_spec())) + + assert served.access_token.get_secret_value() == first.access_token.get_secret_value() + assert len(poster.requests) == 1 + + def test_a_failed_write_leaves_no_staging_file_holding_a_live_token(tmp_path: Path): """Nothing ever sweeps this directory, so a staging file a failed write leaves behind would keep a working token readable on disk for as long as the pod lives.""" diff --git a/tests/unit/llms/base_llm/batches/test_transformation.py b/tests/unit/llms/base_llm/batches/test_transformation.py index 0c360ce2ed9..87d1900b07c 100644 --- a/tests/unit/llms/base_llm/batches/test_transformation.py +++ b/tests/unit/llms/base_llm/batches/test_transformation.py @@ -9,13 +9,13 @@ one abstract property), so the only concrete behavior to regression-lock is: subclass missing any abstract member also cannot be instantiated; a subclass implementing all of them can. - get_config(): a classmethod that reflects over ``cls.__dict__`` and returns - the class-level config attributes, filtering out dunders, ``_abc`` internals, + the class-level config attributes, filtering out every underscore-prefixed name, callables (function/builtin/classmethod/staticmethod), and ``None`` values. These tests assert the exact dict get_config() produces for hand-built subclasses, so a change to the filter predicate (e.g. dropping the ``None`` -filter, dropping the staticmethod/classmethod filter, or widening the prefix -filter to all single-underscore names) makes a test fail. +filter, dropping the staticmethod/classmethod filter, or narrowing the prefix +filter so single-underscore names leak back in) makes a test fail. """ @@ -38,7 +38,7 @@ class _ConcreteBatchesConfig(BaseBatchesConfig): int_attr = 42 list_attr = [1, 2, 3] none_attr = None - _single_underscore = "kept" + _single_underscore = "private" @property def custom_llm_provider(self) -> LlmProviders: @@ -136,8 +136,7 @@ def test_subclass_missing_any_abstract_member_cannot_instantiate(missing_member) def test_get_config_returns_class_level_non_none_data_attrs(): """Exact contents: only class-level data attributes that are not None, - not dunders, not callables. Single-underscore names ARE kept (only ``__`` - and ``_abc`` prefixes are filtered). The ``custom_llm_provider`` property + not underscore-prefixed, not callables. The ``custom_llm_provider`` property object also survives the filter (a property is neither a function nor None), matching how real provider subclasses define it.""" config = _ConcreteBatchesConfig.get_config() @@ -147,7 +146,6 @@ def test_get_config_returns_class_level_non_none_data_attrs(): "string_attr": "hello", "int_attr": 42, "list_attr": [1, 2, 3], - "_single_underscore": "kept", } diff --git a/tests/unit/llms/bedrock/chat/test_converse_transformation.py b/tests/unit/llms/bedrock/chat/test_converse_transformation.py index 4dcf312737d..23ac5baf1ba 100644 --- a/tests/unit/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/unit/llms/bedrock/chat/test_converse_transformation.py @@ -12203,6 +12203,21 @@ def test_converse_claude_request_over_the_cache_control_cap_drops_the_oldest_hin assert messages == messages_before +@pytest.mark.parametrize( + "model, expected", + [ + ("us.anthropic.claude-sonnet-5", {"type": "disabled"}), + ("us.anthropic.claude-haiku-4-5", None), + ("us.anthropic.claude-sonnet-5-5", None), + ], +) +def test_none_disables_only_optional_default_thinking(model: str, expected: dict[str, str] | None) -> None: + params: Final[dict[str, object]] = {} + AmazonConverseConfig()._handle_reasoning_effort_parameter(model, "none", params) + assert params.get("thinking") == expected + assert "output_config" not in params + + def test_converse_claude_tool_message_with_several_marked_blocks_counts_as_one_cache_point_against_the_cap( monkeypatch: pytest.MonkeyPatch, local_model_cost_map: None ) -> None: diff --git a/tests/unit/llms/bedrock/test_bedrock_common_utils.py b/tests/unit/llms/bedrock/test_bedrock_common_utils.py index 7cc4629e497..31999ee0379 100644 --- a/tests/unit/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/unit/llms/bedrock/test_bedrock_common_utils.py @@ -846,6 +846,7 @@ def test_sign_aws_request_assumes_role_with_external_id(monkeypatch): import boto3 from botocore.exceptions import ClientError + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.common_utils import CommonBatchFilesUtils monkeypatch.delenv("AWS_EXTERNAL_ID", raising=False) @@ -879,7 +880,7 @@ def test_sign_aws_request_assumes_role_with_external_id(monkeypatch): } with patch.object(boto3, "client", return_value=FakeSTSClient()): - signed_headers, signed_data = CommonBatchFilesUtils().sign_aws_request( + signed_headers, signed_data = CommonBatchFilesUtils(BaseAWSLLM()).sign_aws_request( service_name="bedrock", data={"jobName": "litellm-batch-job"}, endpoint_url="https://bedrock.us-east-1.amazonaws.com/model-invocation-job", @@ -899,6 +900,7 @@ def test_sign_aws_request_assumes_role_with_session_tags(monkeypatch): import boto3 from botocore.exceptions import ClientError + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.common_utils import CommonBatchFilesUtils monkeypatch.delenv("AWS_WEB_IDENTITY_TOKEN_FILE", raising=False) @@ -934,7 +936,7 @@ def test_sign_aws_request_assumes_role_with_session_tags(monkeypatch): } with patch.object(boto3, "client", return_value=FakeSTSClient()): - signed_headers, _signed_data = CommonBatchFilesUtils().sign_aws_request( + signed_headers, _signed_data = CommonBatchFilesUtils(BaseAWSLLM()).sign_aws_request( service_name="bedrock", data={"jobName": "litellm-batch-job"}, endpoint_url="https://bedrock.us-east-1.amazonaws.com/model-invocation-job", diff --git a/tests/unit/llms/bedrock/test_web_identity_session_policy.py b/tests/unit/llms/bedrock/test_web_identity_session_policy.py index 332e1258e9c..a181087cba9 100644 --- a/tests/unit/llms/bedrock/test_web_identity_session_policy.py +++ b/tests/unit/llms/bedrock/test_web_identity_session_policy.py @@ -30,16 +30,54 @@ action. """ import base64 +import itertools import json -from datetime import datetime, timedelta, timezone +from collections.abc import Callable +from datetime import datetime, timezone +from fnmatch import fnmatchcase from types import MappingProxyType -from typing import Final -from unittest.mock import MagicMock, patch +from typing import Final, TypeAlias +from unittest.mock import AsyncMock, MagicMock, patch import pytest +from botocore.session import get_session from pydantic import TypeAdapter -from litellm.llms.bedrock.base_aws_llm import WebIdentitySessionPolicy, _SessionPolicyStatement +from litellm.integrations.s3_v2 import S3Logger +from litellm.integrations.sqs import SQSLogger +from litellm.llms.aws_polly.text_to_speech.transformation import AWSPollyTextToSpeechConfig +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig +from litellm.llms.bedrock.chat.converse_handler import BedrockConverseLLM +from litellm.llms.bedrock.embed.embedding import BedrockEmbedding +from litellm.llms.bedrock.files.handler import BedrockFilesHandler +from litellm.llms.bedrock.files.transformation import BedrockFilesConfig +from litellm.llms.bedrock.model_listing import BedrockModelLister +from litellm.llms.bedrock.passthrough.transformation import BedrockAgentRuntimeAWS +from litellm.llms.bedrock.realtime.handler import BedrockRealtime +from litellm.llms.bedrock.web_identity_session_policy import WebIdentitySessionPolicy, _SessionPolicyStatement +from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.opensearch.vector_stores.transformation import OpenSearchVectorStoreConfig +from litellm.llms.s3_vectors.vector_stores.transformation import S3VectorsVectorStoreConfig +from litellm.llms.sagemaker.chat.handler import SagemakerChatHandler +from litellm.llms.sagemaker.chat.transformation import SagemakerChatConfig +from litellm.llms.sagemaker.completion.handler import SagemakerLLM +from litellm.proxy.common_utils.load_config_utils import s3_object_reader +from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import _is_bedrock_agent_runtime_route +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.comprehend_medical_passthrough_logging_handler import ( + COMPREHEND_MEDICAL_SUPPORTED_OPERATIONS, + ComprehendMedicalAWS, +) +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.transcribe_passthrough_logging_handler import ( + TranscribeAWS, + transcribe_supported_operations, +) +from litellm.rag.ingestion.bedrock_ingestion import BedrockRAGIngestion +from litellm.rag.ingestion.s3_vectors_ingestion import S3VectorsRAGIngestion +from litellm.secret_managers.aws_secret_manager_v2 import AWSSecretsManagerV2 +from litellm.types.router import GenericLiteLLMParams + +_STS_SESSION_EXPIRATION: Final = datetime(2099, 1, 1, tzinfo=timezone.utc) # Actions the Claude Platform on AWS service is documented to call. # Source: AWS IAM action reference + the #27678 surface area. @@ -67,7 +105,7 @@ def _captured_policy_document() -> str: "AccessKeyId": "k", "SecretAccessKey": "s", "SessionToken": "t", - "Expiration": datetime.now(timezone.utc) + timedelta(hours=1), + "Expiration": _STS_SESSION_EXPIRATION, }, "PackedPolicySize": 0, } @@ -100,17 +138,20 @@ def _captured_policy() -> WebIdentitySessionPolicy: def _granted_actions(policy: WebIdentitySessionPolicy) -> frozenset[str]: - return frozenset(action for stmt in policy["Statement"] for action in stmt["Action"]) + return frozenset(itertools.chain.from_iterable(stmt["Action"] for stmt in policy["Statement"])) + + +def _invoke_model_grants_over_any_transport(policy: WebIdentitySessionPolicy) -> tuple[_SessionPolicyStatement, ...]: + return tuple( + stmt for stmt in policy["Statement"] if "bedrock:InvokeModel" in stmt["Action"] and "Condition" not in stmt + ) def _statement_by_sid(policy: WebIdentitySessionPolicy, sid: str) -> _SessionPolicyStatement: for stmt in policy["Statement"]: if stmt.get("Sid") == sid: return stmt - raise AssertionError( - f"Sid={sid!r} not found in session policy; " - f"saw {[s.get('Sid') for s in policy['Statement']]}" - ) + raise AssertionError(f"Sid={sid!r} not found in session policy; saw {[s.get('Sid') for s in policy['Statement']]}") class TestWebIdentitySessionPolicyShape: @@ -122,14 +163,12 @@ class TestWebIdentitySessionPolicyShape: def test_bedrock_statement_actions_preserved(self): """The original bedrock action set must still be granted — regression for the pre-existing bedrock/* routes.""" - policy = _captured_policy() - bedrock_stmt = _statement_by_sid(policy, "BedrockLiteLLM") - actions = set(bedrock_stmt["Action"]) + granted: Final = _granted_actions(_captured_policy()) for required in ( "bedrock:InvokeModel", "bedrock:InvokeModelWithResponseStream", ): - assert required in actions, f"{required} missing from BedrockLiteLLM" + assert required in granted, f"{required} missing from the Bedrock session policy" def test_bedrock_count_tokens_action_present(self): """Regression for #33142: the CountTokens handler authorizes @@ -140,8 +179,7 @@ class TestWebIdentitySessionPolicyShape: bedrock_stmt = _statement_by_sid(policy, "BedrockLiteLLM") actions = set(bedrock_stmt["Action"]) assert "bedrock:CountTokens" in actions, ( - "bedrock:CountTokens missing from BedrockLiteLLM — " - "count-tokens requests will 403 on OIDC auth" + "bedrock:CountTokens missing from BedrockLiteLLM — count-tokens requests will 403 on OIDC auth" ) @@ -152,8 +190,7 @@ class TestClaudePlatformActionsCovered: @pytest.mark.parametrize("action", sorted(_CLAUDE_PLATFORM_ACTIONS)) def test_claude_platform_action_present(self, action: str): assert action in _granted_actions(_captured_policy()), ( - f"{action} missing from session policy — " - f"bedrock/claude_platform/* requests will 403 on OIDC auth" + f"{action} missing from session policy — bedrock/claude_platform/* requests will 403 on OIDC auth" ) def test_claude_platform_statement_allows(self): @@ -240,8 +277,7 @@ class TestBedrockMantleActionsCovered: if isinstance(actions, str): actions = [actions] assert "bedrock-mantle:*" not in actions, ( - "session policy must not grant bedrock-mantle:* — " - "the ceiling should match the documented action set" + "session policy must not grant bedrock-mantle:* — the ceiling should match the documented action set" ) def test_bedrock_mantle_statement_carries_secure_transport_condition(self): @@ -249,8 +285,7 @@ class TestBedrockMantleActionsCovered: stmt = _statement_by_sid(policy, "BedrockMantleLiteLLM") cond = stmt.get("Condition") or {} assert cond.get("Bool", {}).get("aws:SecureTransport") == "true", ( - "BedrockMantleLiteLLM must require aws:SecureTransport=true " - "to keep parity with the bedrock statement" + "BedrockMantleLiteLLM must require aws:SecureTransport=true to keep parity with the bedrock statement" ) @@ -285,12 +320,8 @@ class TestInvalidIdentityTokenSurfacesAudience: class _InvalidIdentityTokenException(Exception): pass - mock_sts.exceptions.InvalidIdentityTokenException = ( - _InvalidIdentityTokenException - ) - mock_sts.assume_role_with_web_identity.side_effect = ( - _InvalidIdentityTokenException(self._STS_MESSAGE) - ) + mock_sts.exceptions.InvalidIdentityTokenException = _InvalidIdentityTokenException + mock_sts.assume_role_with_web_identity.side_effect = _InvalidIdentityTokenException(self._STS_MESSAGE) with ( patch("boto3.client", return_value=mock_sts), @@ -340,16 +371,20 @@ class TestPolicyTransportConditions: stmt = _statement_by_sid(policy, "ClaudePlatformLiteLLM") cond = stmt.get("Condition") or {} assert cond.get("Bool", {}).get("aws:SecureTransport") == "true", ( - "ClaudePlatformLiteLLM must require aws:SecureTransport=true " - "to keep parity with the bedrock statement" + "ClaudePlatformLiteLLM must require aws:SecureTransport=true to keep parity with the bedrock statement" ) - def test_only_the_outbound_identity_statement_requires_a_claude_platform_caller(self): + def test_statements_for_services_litellm_calls_directly_require_no_calling_service(self): + directly_called_services: Final = frozenset( + {"bedrock", "bedrock-agentcore", "bedrock-mantle", "aws-external-anthropic"} + ) directly_called: Final = tuple( - stmt for stmt in _captured_policy()["Statement"] if stmt["Sid"] != "ClaudePlatformOutboundIdentityLiteLLM" + stmt + for stmt in _captured_policy()["Statement"] + if all(_service(action) in directly_called_services for action in stmt["Action"]) ) assert directly_called - assert all("StringEquals" not in stmt["Condition"] for stmt in directly_called), ( + assert all("StringEquals" not in (stmt.get("Condition") or {}) for stmt in directly_called), ( "statements LiteLLM calls directly must not require aws:CalledViaLast, or OIDC auth 403s on those routes" ) @@ -363,6 +398,9 @@ _BEDROCK_ROUTE_ACTIONS: Final = MappingProxyType( "model/{model_id}/converse": "bedrock:InvokeModel", "model/{model_id}/converse-stream": "bedrock:InvokeModelWithResponseStream", "model/{model_id}/count-tokens": "bedrock:CountTokens", + "model/{model_id}/invoke-with-bidirectional-stream": "bedrock:InvokeModel", + "bedrock-mantle anthropic/v1/messages": "bedrock-mantle:CreateInference", + "bedrock-mantle anthropic/v1/messages/count_tokens": "bedrock-mantle:CountTokens", "guardrail/{guardrail_id}/version/{version}/apply": "bedrock:ApplyGuardrail", "rerank": "bedrock:Rerank", "knowledgebases/{knowledge_base_id}/retrieve": "bedrock:Retrieve", @@ -393,3 +431,551 @@ class TestSessionPolicyGrantsEveryBedrockRoute: def test_policy_document_fits_the_sts_plaintext_limit(self): assert len(_captured_policy_document()) <= _STS_SESSION_POLICY_PLAINTEXT_LIMIT + + +def _s3_logger() -> S3Logger: + with patch("asyncio.create_task"): + return S3Logger(s3_bucket_name="litellm-logs", s3_region_name="us-east-1") + + +def _sqs_logger() -> SQSLogger: + with patch("asyncio.create_task"), patch("litellm.aws_sqs_callback_params", {}): + return SQSLogger( + sqs_queue_url="https://sqs.us-east-1.amazonaws.com/123456789012/logs", sqs_region_name="us-east-1" + ) + + +def _bedrock_knowledge_base_ingestion() -> BedrockRAGIngestion: + return BedrockRAGIngestion( + ingest_options={ + "vector_store": {"custom_llm_provider": "bedrock", "aws_region_name": "us-east-1"}, + } + ) + + +def _bedrock_model_lister() -> BedrockModelLister: + return BedrockModelLister(deployment={"aws_region_name": "us-east-1"}, client=HTTPHandler()) + + +def _s3_vectors_ingestion() -> S3VectorsRAGIngestion: + return S3VectorsRAGIngestion( + ingest_options={ + "embedding": {"model": "bedrock/amazon.titan-embed-text-v2:0"}, + "vector_store": { + "custom_llm_provider": "s3_vectors", + "vector_bucket_name": "litellm-vectors", + "aws_region_name": "us-east-1", + }, + } + ) + + +_CallerAndActions: TypeAlias = tuple[Callable[[], BaseAWSLLM], tuple[str, ...]] + +_CALLER_ROUTE_ACTIONS: Final[MappingProxyType[str, _CallerAndActions]] = MappingProxyType( + { + "bedrock batches": ( + BedrockBatchesConfig, + ( + "bedrock:CreateModelInvocationJob", + "bedrock:GetModelInvocationJob", + "bedrock:StopModelInvocationJob", + "bedrock:TagResource", + "iam:PassRole", + ), + ), + "bedrock files": ( + BedrockFilesConfig, + ("s3:PutObject", "s3:GetObject", "s3:DeleteObject", "s3:ListBucket", "kms:GenerateDataKey"), + ), + "bedrock file download": (BedrockFilesHandler, ("s3:GetObject", "kms:Decrypt")), + "bedrock embeddings": ( + BedrockEmbedding, + ( + "bedrock:InvokeModel", + "bedrock:ApplyGuardrail", + "bedrock:GetAsyncInvoke", + "s3:GetObject", + "s3:PutObject", + "kms:GenerateDataKey", + "kms:Decrypt", + ), + ), + "bedrock realtime": (BedrockRealtime, ("bedrock:InvokeModel", "bedrock:ApplyGuardrail")), + "bedrock agent runtime passthrough": ( + BedrockAgentRuntimeAWS, + ( + "bedrock:InvokeAgent", + "bedrock:InvokeInlineAgent", + "bedrock:InvokeFlow", + "bedrock:StartFlowExecution", + "bedrock:StopFlowExecution", + "bedrock:Retrieve", + "bedrock:RetrieveAndGenerate", + "bedrock:GenerateQuery", + "bedrock:GetKnowledgeBase", + "bedrock:GetDocumentContent", + "bedrock:CheckIngestedDocumentAcl", + "bedrock:GetIngestedDocumentAcl", + "bedrock:Rerank", + "bedrock:OptimizePrompt", + "bedrock:InvokeModel", + "bedrock:ApplyGuardrail", + "bedrock:GetInferenceProfile", + "sqlworkbench:GetSqlRecommendations", + ), + ), + "bedrock model listing": ( + _bedrock_model_lister, + ("bedrock:ListFoundationModels", "bedrock:ListInferenceProfiles"), + ), + "bedrock knowledge base ingestion": ( + _bedrock_knowledge_base_ingestion, + ( + "bedrock:ListDataSources", + "bedrock:GetDataSource", + "bedrock:CreateKnowledgeBase", + "bedrock:GetKnowledgeBase", + "bedrock:CreateDataSource", + "bedrock:StartIngestionJob", + "bedrock:GetIngestionJob", + "s3:CreateBucket", + "s3:PutObject", + "kms:GenerateDataKey", + "aoss:CreateSecurityPolicy", + "aoss:CreateAccessPolicy", + "aoss:CreateCollection", + "aoss:BatchGetCollection", + "aoss:APIAccessAll", + "iam:CreateServiceLinkedRole", + "iam:CreateRole", + "iam:PutRolePolicy", + "iam:PassRole", + ), + ), + "sagemaker completion": ( + SagemakerLLM, + ("sagemaker:InvokeEndpoint", "sagemaker:InvokeEndpointWithResponseStream"), + ), + "sagemaker chat": ( + SagemakerChatConfig, + ("sagemaker:InvokeEndpoint", "sagemaker:InvokeEndpointWithResponseStream"), + ), + "sagemaker chat handler": ( + SagemakerChatHandler, + ("sagemaker:InvokeEndpoint", "sagemaker:InvokeEndpointWithResponseStream"), + ), + "polly speech": (AWSPollyTextToSpeechConfig, ("polly:SynthesizeSpeech",)), + "s3 vectors search": (S3VectorsVectorStoreConfig, ("s3vectors:QueryVectors",)), + "s3 vectors ingest": ( + _s3_vectors_ingestion, + ( + "s3vectors:GetVectorBucket", + "s3vectors:CreateVectorBucket", + "s3vectors:GetIndex", + "s3vectors:CreateIndex", + "s3vectors:PutVectors", + ), + ), + "s3 logging": (_s3_logger, ("s3:PutObject", "kms:GenerateDataKey")), + "sqs logging": (_sqs_logger, ("sqs:SendMessage", "kms:GenerateDataKey")), + "secrets manager": ( + AWSSecretsManagerV2, + ( + "secretsmanager:GetSecretValue", + "secretsmanager:DescribeSecret", + "secretsmanager:CreateSecret", + "secretsmanager:PutSecretValue", + "secretsmanager:UpdateSecret", + "secretsmanager:RestoreSecret", + "secretsmanager:TagResource", + "secretsmanager:ReplicateSecretToRegions", + "secretsmanager:DeleteSecret", + "kms:GenerateDataKey", + "kms:Decrypt", + ), + ), + "comprehend medical passthrough": ( + ComprehendMedicalAWS, + tuple(f"comprehendmedical:{operation}" for operation in sorted(COMPREHEND_MEDICAL_SUPPORTED_OPERATIONS)), + ), + "transcribe passthrough": ( + TranscribeAWS, + ( + "transcribe:StartTranscriptionJob", + "transcribe:GetTranscriptionJob", + "s3:GetObject", + "s3:PutObject", + "iam:PassRole", + "kms:GenerateDataKey", + "kms:Decrypt", + ), + ), + } +) + + +_KMS_VIA_SERVICE_BY_ROUTE: Final = MappingProxyType( + { + "bedrock files": "s3.*", + "bedrock file download": "s3.*", + "bedrock embeddings": "s3.*", + "s3 logging": "s3.*", + "sqs logging": "sqs.*", + "bedrock knowledge base ingestion": "s3.*", + "secrets manager": "secretsmanager.*", + "transcribe passthrough": "s3.*", + } +) + + +_PASS_ROLE_TARGET_BY_ROUTE: Final = MappingProxyType( + { + "bedrock batches": "bedrock.amazonaws.com", + "bedrock knowledge base ingestion": "bedrock.amazonaws.com", + "transcribe passthrough": "transcribe.amazonaws.com", + } +) + + +def _sts_mock() -> MagicMock: + sts: Final = MagicMock() + sts.assume_role_with_web_identity.return_value = { + "Credentials": { + "AccessKeyId": "k", + "SecretAccessKey": "s", + "SessionToken": "t", + "Expiration": _STS_SESSION_EXPIRATION, + }, + "PackedPolicySize": 0, + } + return sts + + +def _policies_sent_to_sts(callers: tuple[BaseAWSLLM, ...], session_name: str) -> tuple[WebIdentitySessionPolicy, ...]: + sts: Final = _sts_mock() + with ( + patch("boto3.client", return_value=sts), + patch("litellm.llms.bedrock.base_aws_llm.get_secret", return_value="oidc-jwt-token"), + ): + for caller in callers: + caller.get_credentials( + aws_web_identity_token="oidc/google/litellm", + aws_role_name="arn:aws:iam::123456789012:role/litellm-role", + aws_session_name=session_name, + aws_region_name="us-east-1", + ) + return tuple( + _SESSION_POLICY_ADAPTER.validate_python(json.loads(call.kwargs["Policy"])) + for call in sts.assume_role_with_web_identity.call_args_list + ) + + +_PROXY_ROLE_ARN: Final = "arn:aws:iam::123456789012:role/litellm-role" + +_KNOWLEDGE_BASE_ROLE_ACTIONS: Final = frozenset({"iam:CreateRole", "iam:PutRolePolicy", "iam:PassRole"}) + + +def _role_created_in_account(**create_role_request: str) -> dict[str, dict[str, str]]: + return {"Role": {"Arn": f"arn:aws:iam::123456789012:role/{create_role_request['RoleName']}"}} + + +def _service(action: str) -> str: + return action.split(":", 1)[0] + + +def _grants_kms(statement: _SessionPolicyStatement) -> bool: + return any(_service(action) == "kms" for action in statement["Action"]) + + +def _ungranted(actions: tuple[str, ...], policy: WebIdentitySessionPolicy) -> tuple[str, ...]: + granted: Final = _granted_actions(policy) + return tuple(action for action in actions if not any(fnmatchcase(action, pattern) for pattern in granted)) + + +class TestEachCallerGetsItsOwnCeiling: + """Every AWS caller besides Bedrock runtime (batches, files, model listing, knowledge base + ingestion, SageMaker, Polly, S3 Vectors, the loggers, Secrets Manager, the config bucket, and the + Comprehend Medical and Transcribe pass-throughs) authorizes against actions outside the Bedrock + runtime ceiling, so on web identity auth it 403d with "no session policy allows" even when the + role allowed it. Each caller's STS session must carry a ceiling that grants what its routes call + and nothing from another caller's service.""" + + @pytest.mark.parametrize(("route", "caller_and_actions"), sorted(_CALLER_ROUTE_ACTIONS.items())) + def test_route_actions_are_granted_by_the_callers_ceiling(self, route: str, caller_and_actions: _CallerAndActions): + make_caller, actions = caller_and_actions + (policy,) = _policies_sent_to_sts((make_caller(),), session_name=f"grant-{route}") + missing: Final = _ungranted(actions, policy) + assert not missing, f"{route} calls {sorted(missing)}, which its session policy does not grant" + + @pytest.mark.parametrize(("route", "caller_and_actions"), sorted(_CALLER_ROUTE_ACTIONS.items())) + def test_ceiling_grants_no_service_the_route_does_not_call(self, route: str, caller_and_actions: _CallerAndActions): + make_caller, actions = caller_and_actions + (policy,) = _policies_sent_to_sts((make_caller(),), session_name=f"deny-{route}") + called_services: Final = frozenset(_service(action) for action in actions) + extra: Final = sorted(a for a in _granted_actions(policy) if _service(a) not in called_services) + assert not extra, f"{route}'s session policy also grants {extra}" + + @pytest.mark.parametrize(("route", "caller_and_actions"), sorted(_CALLER_ROUTE_ACTIONS.items())) + def test_ceiling_fits_the_sts_plaintext_limit(self, route: str, caller_and_actions: _CallerAndActions): + make_caller, _actions = caller_and_actions + (policy,) = _policies_sent_to_sts((make_caller(),), session_name=f"size-{route}") + assert len(json.dumps(policy, separators=(",", ":"))) <= _STS_SESSION_POLICY_PLAINTEXT_LIMIT + + def test_bedrock_chat_ceiling_reaches_no_other_callers_service(self): + every_route_action: Final = itertools.chain.from_iterable( + actions for _caller, actions in _CALLER_ROUTE_ACTIONS.values() + ) + other_services: Final = frozenset( + _service(action) for action in every_route_action if _service(action) != "bedrock" + ) + (policy,) = _policies_sent_to_sts((BedrockConverseLLM(),), session_name="chat-reach") + reached: Final = sorted(action for action in _granted_actions(policy) if _service(action) in other_services) + assert not reached, f"a Bedrock chat session can also call {reached}" + + @pytest.mark.parametrize(("route", "target_service"), sorted(_PASS_ROLE_TARGET_BY_ROUTE.items())) + def test_service_role_can_only_be_passed_to_the_routes_service(self, route: str, target_service: str): + make_caller, _actions = _CALLER_ROUTE_ACTIONS[route] + (policy,) = _policies_sent_to_sts((make_caller(),), session_name=f"pass-role-{route}") + granting: Final = tuple(stmt for stmt in policy["Statement"] if "iam:PassRole" in stmt["Action"]) + assert granting + assert all( + stmt["Condition"].get("StringEquals") == {"iam:PassedToService": target_service} for stmt in granting + ) + + @pytest.mark.parametrize("action", ["s3:GetObject", "s3:PutObject"]) + def test_async_invoke_bucket_access_is_only_granted_through_bedrock(self, action: str): + (policy,) = _policies_sent_to_sts((BedrockEmbedding(),), session_name=f"async-{action}") + granting: Final = tuple(stmt for stmt in policy["Statement"] if action in stmt["Action"]) + assert granting + assert all( + stmt["Condition"].get("StringEquals") == {"aws:CalledViaLast": "bedrock.amazonaws.com"} for stmt in granting + ) + + def test_transcribe_ceiling_grants_every_operation_the_passthrough_serves(self): + (policy,) = _policies_sent_to_sts((TranscribeAWS(),), session_name="transcribe-operations") + operations: Final = tuple(f"transcribe:{operation}" for operation in sorted(transcribe_supported_operations())) + assert not _ungranted(operations, policy) + + def test_transcript_output_write_is_only_granted_through_transcribe(self): + (policy,) = _policies_sent_to_sts((TranscribeAWS(),), session_name="transcribe-output") + granting: Final = tuple(stmt for stmt in policy["Statement"] if "s3:PutObject" in stmt["Action"]) + assert granting + assert all( + stmt["Condition"].get("StringEquals") == {"aws:CalledViaLast": "transcribe.amazonaws.com"} + for stmt in granting + ) + + def test_knowledge_base_ingestion_can_only_create_the_opensearch_serverless_service_linked_role(self): + (policy,) = _policies_sent_to_sts((_bedrock_knowledge_base_ingestion(),), session_name="kb-slr") + granting: Final = tuple(stmt for stmt in policy["Statement"] if "iam:CreateServiceLinkedRole" in stmt["Action"]) + assert granting + assert all( + stmt["Condition"].get("StringEquals") == {"iam:AWSServiceName": "observability.aoss.amazonaws.com"} + for stmt in granting + ) + + async def test_knowledge_base_ingestion_can_only_create_and_pass_the_roles_it_names(self): + ingestion: Final = BedrockRAGIngestion( + ingest_options={ + "vector_store": { + "custom_llm_provider": "bedrock", + "aws_region_name": "us-east-1", + "aws_web_identity_token": "oidc/google/litellm", + "aws_role_name": _PROXY_ROLE_ARN, + "aws_session_name": "kb-role-scope", + } + } + ) + sts: Final = _sts_mock() + iam: Final = MagicMock() + iam.create_role.side_effect = _role_created_in_account + with ( + patch("boto3.client", return_value=sts), + patch("boto3.Session", return_value=MagicMock(client=MagicMock(return_value=iam))), + patch("litellm.llms.bedrock.base_aws_llm.get_secret", return_value="oidc-jwt-token"), + patch("asyncio.sleep", new=AsyncMock()), + ): + role_arn: Final = await ingestion._create_bedrock_role( + "abcd1234", "123456789012", "arn:aws:aoss:us-east-1:123456789012:collection/kb" + ) + policy: Final = _SESSION_POLICY_ADAPTER.validate_python( + json.loads(sts.assume_role_with_web_identity.call_args.kwargs["Policy"]) + ) + granting: Final = tuple( + stmt for stmt in policy["Statement"] if _KNOWLEDGE_BASE_ROLE_ACTIONS & set(stmt["Action"]) + ) + assert _granted_actions(policy) >= _KNOWLEDGE_BASE_ROLE_ACTIONS + assert all(fnmatchcase(role_arn, stmt["Resource"]) for stmt in granting) + assert not any(fnmatchcase(_PROXY_ROLE_ARN, stmt["Resource"]) for stmt in granting) + + def test_config_bucket_reader_signs_with_a_read_only_ceiling(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("AWS_WEB_IDENTITY_TOKEN", "oidc/google/litellm") + monkeypatch.setenv("AWS_ROLE_NAME", "arn:aws:iam::123456789012:role/litellm-role") + monkeypatch.setenv("AWS_SESSION_NAME", "config-bucket") + monkeypatch.setenv("AWS_REGION_NAME", "us-east-1") + sts: Final = _sts_mock() + with ( + patch("boto3.client", return_value=sts), + patch("litellm.llms.bedrock.base_aws_llm.get_secret", return_value="oidc-jwt-token"), + ): + s3_object_reader("litellm-config") + policy: Final = _SESSION_POLICY_ADAPTER.validate_python( + json.loads(sts.assume_role_with_web_identity.call_args.kwargs["Policy"]) + ) + assert _granted_actions(policy) == frozenset({"s3:GetObject", "kms:Decrypt"}) + + def test_file_download_ceiling_cannot_write_list_or_delete(self): + (policy,) = _policies_sent_to_sts((BedrockFilesHandler(),), session_name="file-download-only") + assert _granted_actions(policy) == frozenset({"s3:GetObject", "kms:Decrypt"}) + + @pytest.mark.parametrize(("route", "via_service"), sorted(_KMS_VIA_SERVICE_BY_ROUTE.items())) + def test_kms_is_only_granted_through_the_routes_storage_service(self, route: str, via_service: str): + make_caller, _actions = _CALLER_ROUTE_ACTIONS[route] + (policy,) = _policies_sent_to_sts((make_caller(),), session_name=f"kms-via-{route}") + granting: Final = tuple(stmt for stmt in policy["Statement"] if _grants_kms(stmt)) + assert granting + assert all(stmt["Condition"].get("StringLike") == {"kms:ViaService": via_service} for stmt in granting) + + +_OPENSEARCH_ENDPOINT_ACTIONS: Final = MappingProxyType( + { + "https://search-docs-abc123.us-east-1.es.amazonaws.com": "es:ESHttpPost", + "https://abc123.us-east-1.aoss.amazonaws.com": "aoss:APIAccessAll", + } +) + + +def _opensearch_policy(api_base: str) -> WebIdentitySessionPolicy: + sts: Final = _sts_mock() + with ( + patch("boto3.client", return_value=sts), + patch("litellm.llms.bedrock.base_aws_llm.get_secret", return_value="oidc-jwt-token"), + ): + OpenSearchVectorStoreConfig().validate_environment( + {}, + GenericLiteLLMParams( + api_base=api_base, + aws_web_identity_token="oidc/google/litellm", + aws_role_name="arn:aws:iam::123456789012:role/litellm-role", + aws_session_name=f"opensearch-{api_base}", + aws_region_name="us-east-1", + ), + ) + return _SESSION_POLICY_ADAPTER.validate_python( + json.loads(sts.assume_role_with_web_identity.call_args.kwargs["Policy"]) + ) + + +class TestOpenSearchSearchGetsItsOwnCeiling: + @pytest.mark.parametrize(("api_base", "action"), sorted(_OPENSEARCH_ENDPOINT_ACTIONS.items())) + def test_search_action_is_granted_by_the_opensearch_ceiling(self, api_base: str, action: str): + assert action in _granted_actions(_opensearch_policy(api_base)) + + @pytest.mark.parametrize("api_base", sorted(_OPENSEARCH_ENDPOINT_ACTIONS)) + def test_opensearch_ceiling_grants_only_opensearch_services(self, api_base: str): + allowed: Final = frozenset(_service(action) for action in _OPENSEARCH_ENDPOINT_ACTIONS.values()) + extra: Final = sorted(a for a in _granted_actions(_opensearch_policy(api_base)) if _service(a) not in allowed) + assert not extra, f"the OpenSearch session policy also grants {extra}" + + +class TestInvokeModelTransport: + def test_chat_ceiling_grants_invoke_model_only_over_tls(self): + (policy,) = _policies_sent_to_sts((BedrockConverseLLM(),), session_name="chat-transport") + assert "bedrock:InvokeModel" in _granted_actions(policy) + assert not _invoke_model_grants_over_any_transport(policy) + + def test_realtime_ceiling_grants_invoke_model_over_any_transport_and_nothing_else(self): + (policy,) = _policies_sent_to_sts((BedrockRealtime(),), session_name="realtime-transport") + assert _invoke_model_grants_over_any_transport(policy), ( + "Bedrock authorizes the Nova Sonic bidirectional stream as bedrock:InvokeModel with aws:SecureTransport " + "false (live probe on amazon.nova-2-sonic-v1:0, 2026-11-08), " + "so a TLS-only grant denies every realtime session" + ) + assert _granted_actions(policy) == frozenset({"bedrock:InvokeModel", "bedrock:ApplyGuardrail"}) + + +_MODEL_INVOKING_CALLERS: Final[MappingProxyType[str, Callable[[], BaseAWSLLM]]] = MappingProxyType( + { + "bedrock chat": BedrockConverseLLM, + **{ + route: make_caller + for route, (make_caller, actions) in _CALLER_ROUTE_ACTIONS.items() + if "bedrock:InvokeModel" in actions + }, + } +) + + +class TestEnforcedGuardrails: + @pytest.mark.parametrize(("route", "make_caller"), sorted(_MODEL_INVOKING_CALLERS.items())) + def test_every_grant_to_invoke_models_also_grants_apply_guardrail( + self, route: str, make_caller: Callable[[], BaseAWSLLM] + ): + (policy,) = _policies_sent_to_sts((make_caller(),), session_name=f"guardrail-{route}") + invoking: Final = tuple(stmt for stmt in policy["Statement"] if "bedrock:InvokeModel" in stmt["Action"]) + assert invoking + assert all("bedrock:ApplyGuardrail" in stmt["Action"] for stmt in invoking), ( + f"{route} can invoke models but not apply an account or organization enforced guardrail, " + "which Bedrock answers with AccessDenied" + ) + + +def _agent_runtime_operations_the_passthrough_signs_and_forwards() -> tuple[str, ...]: + service_model: Final = get_session().get_service_model("bedrock-agent-runtime") + http_bindings: Final = tuple( + (name, service_model.operation_model(name).http) for name in service_model.operation_names + ) + return tuple( + name + for name, http in http_bindings + if http["method"] == "POST" and _is_bedrock_agent_runtime_route(f"{http['requestUri'].lstrip('/')}/") + ) + + +class TestAgentRuntimePassThroughCeiling: + def test_ceiling_grants_every_post_operation_the_passthrough_forwards(self): + operations: Final = _agent_runtime_operations_the_passthrough_signs_and_forwards() + (policy,) = _policies_sent_to_sts((BedrockAgentRuntimeAWS(),), session_name="agent-runtime-operations") + assert operations + missing: Final = _ungranted(tuple(f"bedrock:{operation}" for operation in operations), policy) + assert not missing, ( + f"the agent runtime pass-through forwards {sorted(missing)}, which its ceiling does not grant" + ) + + +class TestWebIdentityCredentialCacheIsPerCeiling: + def test_callers_with_different_ceilings_get_their_own_sts_session(self): + bedrock_policy, sagemaker_policy = _policies_sent_to_sts( + (BaseAWSLLM(), SagemakerLLM()), session_name="cache-split" + ) + assert "sagemaker:InvokeEndpoint" not in _granted_actions(bedrock_policy) + assert "sagemaker:InvokeEndpoint" in _granted_actions(sagemaker_policy) + + def test_callers_with_the_same_ceiling_share_one_sts_session(self): + policies: Final = _policies_sent_to_sts((SagemakerLLM(), SagemakerLLM()), session_name="cache-share") + assert len(policies) == 1 + + +class TestBedrockBatchRequestSigning: + def test_create_batch_job_is_signed_with_the_batch_ceiling(self): + sts: Final = _sts_mock() + with ( + patch("boto3.client", return_value=sts), + patch("litellm.llms.bedrock.base_aws_llm.get_secret", return_value="oidc-jwt-token"), + ): + BedrockBatchesConfig().common_utils.sign_aws_request( + service_name="bedrock", + data={"jobName": "litellm-batch-job"}, + endpoint_url="https://bedrock.us-east-1.amazonaws.com/model-invocation-job", + optional_params={ + "aws_web_identity_token": "oidc/google/litellm", + "aws_role_name": "arn:aws:iam::123456789012:role/litellm-role", + "aws_session_name": "batch-sign", + "aws_region_name": "us-east-1", + }, + ) + policy: Final = _SESSION_POLICY_ADAPTER.validate_python( + json.loads(sts.assume_role_with_web_identity.call_args.kwargs["Policy"]) + ) + assert "bedrock:CreateModelInvocationJob" in _granted_actions(policy) diff --git a/tests/unit/llms/custom_httpx/test_http_handler.py b/tests/unit/llms/custom_httpx/test_http_handler.py index 3e583bf8dff..0de8efa2662 100644 --- a/tests/unit/llms/custom_httpx/test_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_http_handler.py @@ -26,6 +26,8 @@ from yarl import URL import litellm from litellm.caching.llm_caching_handler import LLMClientCache from litellm.llms.custom_httpx.aiohttp_transport import LiteLLMAiohttpTransport +from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine, default_token_exchange_engine +from litellm.llms.base_llm.auth.types import MintedToken, TokenExchangeSpec from litellm.llms.custom_httpx.http_handler import ( _CLIENT_REFCOUNT_WHEN_HANDLER_IS_SOLE_REFERRER, AsyncHTTPHandler, @@ -1513,9 +1515,14 @@ async def test_finalizer_on_live_loop_disposes_foreign_loop_session_without_sche class _RetryClientHandler(AsyncHTTPHandler): - def __init__(self, first: httpx.AsyncClient, retry: httpx.AsyncClient) -> None: + def __init__( + self, + first: httpx.AsyncClient, + retry: httpx.AsyncClient, + token_exchange_engine: JwtBearerTokenExchangeEngine = default_token_exchange_engine, + ) -> None: self._retry_client: Final = retry - super().__init__() + super().__init__(token_exchange_engine=token_exchange_engine) self.client = first def create_client( @@ -1552,6 +1559,135 @@ async def test_connection_error_retry_forwards_content(method: str): await handler.close() +class _WifClock: + def __init__(self) -> None: + self.now = 1_000.0 # mutable-ok: fake clock advanced by the test + + def __call__(self) -> float: + return self.now + + +class _CountingIssuer: + def __init__(self) -> None: + self.mints = 0 # mutable-ok: counts tokens minted so each one is distinct + + def post(self, url: str, *, content: bytes, headers: Mapping[str, str], timeout: float) -> httpx.Response: + self.mints += 1 + return httpx.Response( + 200, json={"access_token": f"wif-token-{self.mints}", "expires_in": 600, "token_type": "Bearer"} + ) + + +_WIF_SPEC: Final = TokenExchangeSpec( + token_url="https://api.anthropic.com/v1/oauth/token", + assertion_ref="oidc/env/ANTHROPIC_IDENTITY_TOKEN", + assertion_field="assertion", + static_body={"grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer", "federation_rule_id": "fdrl_1"}, + body_encoding="json", + request_headers={}, + cache_key_identity=("fdrl_1", "org-1", "", ""), + timeout_seconds=2.0, +) + + +def _minted_bearer(engine: JwtBearerTokenExchangeEngine) -> str: + minted: Final = engine.get_token(_WIF_SPEC) + assert isinstance(minted, MintedToken) + return minted.access_token.get_secret_value() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ["post", "put", "patch", "delete"]) +@pytest.mark.parametrize("refreshed_by_another_request", [False, True]) +async def test_connection_error_retry_sends_current_wif_token(method: str, refreshed_by_another_request: bool) -> None: + clock: Final = _WifClock() + engine: Final = JwtBearerTokenExchangeEngine( + poster=_CountingIssuer(), + assertion_reader=lambda ref: "identity-jwt", + clock=clock, + refresh_executor=ThreadPoolExecutor(max_workers=1), + ) + first_token: Final = _minted_bearer(engine) + sent: Final[list[str]] = [] # mutable-ok: async closure capture buffer + + async def expire_then_drop(request: httpx.Request) -> httpx.Response: + sent.append(request.headers["authorization"]) + clock.now += 601 + if refreshed_by_another_request: + _minted_bearer(engine) + raise httpx.ConnectError("connection dropped", request=request) + + async def capture_and_succeed(request: httpx.Request) -> httpx.Response: + sent.append(request.headers["authorization"]) + return httpx.Response(200, request=request) + + first: Final = httpx.AsyncClient(transport=httpx.MockTransport(expire_then_drop)) + retry: Final = httpx.AsyncClient(transport=httpx.MockTransport(capture_and_succeed)) + async with first, retry: + handler: Final = _RetryClientHandler(first=first, retry=retry, token_exchange_engine=engine) + await getattr(handler, method)( + "https://api.anthropic.com/v1/messages", + headers={"authorization": f"Bearer {first_token}", "x-trace": "kept"}, + content=b"{}", + ) + await handler.close() + + assert sent == [f"Bearer {first_token}", f"Bearer {_minted_bearer(engine)}"] + assert first_token != _minted_bearer(engine) + + +@pytest.mark.asyncio +async def test_connection_error_retry_keeps_bearer_the_engine_did_not_mint() -> None: + engine: Final = JwtBearerTokenExchangeEngine( + poster=_CountingIssuer(), assertion_reader=lambda ref: "identity-jwt", clock=_WifClock() + ) + _minted_bearer(engine) + sent: Final[list[str]] = [] # mutable-ok: async closure capture buffer + + async def raise_connection_error(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("connection dropped", request=request) + + async def capture_and_succeed(request: httpx.Request) -> httpx.Response: + sent.append(request.headers["authorization"]) + return httpx.Response(200, request=request) + + first: Final = httpx.AsyncClient(transport=httpx.MockTransport(raise_connection_error)) + retry: Final = httpx.AsyncClient(transport=httpx.MockTransport(capture_and_succeed)) + async with first, retry: + handler: Final = _RetryClientHandler(first=first, retry=retry, token_exchange_engine=engine) + await handler.post("https://api.anthropic.com/v1/messages", headers={"Authorization": "Bearer sk-static"}) + await handler.close() + + assert sent == ["Bearer sk-static"] + + +@pytest.mark.asyncio +async def test_connection_error_retry_keeps_wif_token_sent_to_another_host() -> None: + clock: Final = _WifClock() + engine: Final = JwtBearerTokenExchangeEngine( + poster=_CountingIssuer(), assertion_reader=lambda ref: "identity-jwt", clock=clock + ) + first_token: Final = _minted_bearer(engine) + sent: Final[list[str]] = [] # mutable-ok: async closure capture buffer + + async def expire_then_drop(request: httpx.Request) -> httpx.Response: + clock.now += 601 + raise httpx.ConnectError("connection dropped", request=request) + + async def capture_and_succeed(request: httpx.Request) -> httpx.Response: + sent.append(request.headers["authorization"]) + return httpx.Response(200, request=request) + + first: Final = httpx.AsyncClient(transport=httpx.MockTransport(expire_then_drop)) + retry: Final = httpx.AsyncClient(transport=httpx.MockTransport(capture_and_succeed)) + async with first, retry: + handler: Final = _RetryClientHandler(first=first, retry=retry, token_exchange_engine=engine) + await handler.post("https://elsewhere.example/v1/messages", headers={"authorization": f"Bearer {first_token}"}) + await handler.close() + + assert sent == [f"Bearer {first_token}"] + + @pytest.mark.asyncio @pytest.mark.parametrize("method", ["post", "patch"]) async def test_disconnect_after_send_is_not_replayed(method: str) -> None: diff --git a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py index e04dd35f851..053d60ef0e1 100644 --- a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py @@ -3220,3 +3220,18 @@ def test_databricks_claude_cache_control_cap_keeps_room_for_a_top_level_cache_co assert "cache_control" not in json.dumps(body["messages"][0]) assert json.dumps(body["messages"]).count('"cache_control"') == 3 + + +@pytest.mark.parametrize( + "model, expected", + [ + ("databricks-claude-sonnet-5", {"type": "disabled"}), + ("databricks-claude-haiku-4-5", None), + ("databricks-claude-sonnet-5-5", None), + ], +) +def test_none_disables_only_optional_default_thinking(model: str, expected: dict[str, str] | None) -> None: + params: Final = _map_reasoning_effort(model, "none") + assert params.get("thinking") == expected + assert "output_config" not in params + assert "reasoning_effort" not in params diff --git a/tests/unit/llms/databricks/test_databricks_cost_calculator.py b/tests/unit/llms/databricks/test_databricks_cost_calculator.py index c35aa6e4cf0..e9bc75119aa 100644 --- a/tests/unit/llms/databricks/test_databricks_cost_calculator.py +++ b/tests/unit/llms/databricks/test_databricks_cost_calculator.py @@ -6,6 +6,7 @@ from typing import Final import pytest import litellm +from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry from litellm.llms.databricks.cost_calculator import cost_per_token from litellm.types.utils import ModelInfo, Usage @@ -213,6 +214,7 @@ def test_every_model_without_published_cache_dbu_bills_cache_writes_at_its_own_i if model.startswith("databricks/") and info.get("input_cost_per_token") and model not in PUBLISHED_DBU_PER_MILLION + and not (isinstance(info, AliasedCostEntry) and info.names[0] != model) ] for model in without_published_rates: diff --git a/tests/unit/llms/moonshot/test_moonshot_chat_transformation.py b/tests/unit/llms/moonshot/test_moonshot_chat_transformation.py index b054dd84c22..2f58330e427 100644 --- a/tests/unit/llms/moonshot/test_moonshot_chat_transformation.py +++ b/tests/unit/llms/moonshot/test_moonshot_chat_transformation.py @@ -13,7 +13,7 @@ import pytest import litellm import litellm.utils -from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap, get_model_cost_map from litellm.litellm_core_utils.model_latest_aliases import get_model_latest_aliases from litellm.llms.moonshot.chat.transformation import MoonshotChatConfig @@ -715,7 +715,8 @@ class TestMoonshotReasoningEffort: @pytest.fixture(autouse=True) def force_local_model_cost(self, monkeypatch): - monkeypatch.setattr(litellm, "model_cost", GetModelCostMap.load_local_model_cost_map()) + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", get_model_cost_map(url=litellm.model_cost_map_url)) @pytest.mark.parametrize("model", ["kimi-k3", "kimi-k2.5", "kimi-k2.6"]) def test_reasoning_model_supports_reasoning_effort(self, model): diff --git a/tests/unit/llms/opensearch/vector_stores/test_opensearch_transformation.py b/tests/unit/llms/opensearch/vector_stores/test_opensearch_transformation.py index 07836d57af0..c55393f2e06 100644 --- a/tests/unit/llms/opensearch/vector_stores/test_opensearch_transformation.py +++ b/tests/unit/llms/opensearch/vector_stores/test_opensearch_transformation.py @@ -4,6 +4,7 @@ import re from collections.abc import Callable, Mapping from datetime import datetime, timezone from pathlib import Path +from types import MappingProxyType from typing import Final from unittest.mock import MagicMock @@ -12,7 +13,7 @@ import openai import pytest from botocore.auth import SigV4Auth from botocore.awsrequest import AWSRequest -from botocore.credentials import Credentials +from botocore.credentials import Credentials, RefreshableCredentials from botocore.exceptions import CredentialRetrievalError import litellm @@ -20,6 +21,12 @@ from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.opensearch.vector_stores.index_passthrough import ( + IndexPath, + IndexRefusal, + index_operation, + prepare_index_request, +) from litellm.llms.opensearch.vector_stores.transformation import OpenSearchVectorStoreConfig from litellm.types.llms.bedrock import AwsAuthParams from litellm.types.router import GenericLiteLLMParams @@ -286,6 +293,17 @@ _AMBIENT_AWS_ENV: Final = ( ) + +def _isolate_from_ambient_aws(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + for name in _AMBIENT_AWS_ENV: + monkeypatch.delenv(name, raising=False) + for name in ("AWS_CONFIG_FILE", "AWS_SHARED_CREDENTIALS_FILE", "BOTO_CONFIG"): + config_file: Path = tmp_path / name + config_file.write_text("") + monkeypatch.setenv(name, str(config_file)) + monkeypatch.setenv("AWS_EC2_METADATA_DISABLED", "true") + + @pytest.mark.parametrize("asynchronous", [False, True]) @pytest.mark.parametrize( "credentials,env", @@ -313,13 +331,7 @@ async def test_missing_aws_credentials_are_a_401_before_embedding_or_any_request monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: - for name in _AMBIENT_AWS_ENV: - monkeypatch.delenv(name, raising=False) - for name in ("AWS_CONFIG_FILE", "AWS_SHARED_CREDENTIALS_FILE", "BOTO_CONFIG"): - config_file: Path = tmp_path / name - config_file.write_text("") - monkeypatch.setenv(name, str(config_file)) - monkeypatch.setenv("AWS_EC2_METADATA_DISABLED", "true") + _isolate_from_ambient_aws(monkeypatch, tmp_path) for name, value in env.items(): monkeypatch.setenv(name, value) executor: Final = RecordingEmbeddingExecutor() @@ -353,6 +365,107 @@ def test_failing_credential_source_is_a_401_that_keeps_its_output_out_of_the_rep assert CREDENTIAL_SOURCE_OUTPUT not in str(error.value) +@pytest.mark.parametrize( + "sign", + [ + pytest.param( + lambda config: config.validate_environment({}, GenericLiteLLMParams(api_base=SERVERLESS)), + id="vector-store-search", + ), + pytest.param( + lambda config: config.sign_http_request( + "GET", f"{SERVERLESS}/docs/_mapping", b"", "application/json", {"api_base": SERVERLESS} + ), + id="index-pass-through", + ), + ], +) +def test_a_missing_container_credentials_token_file_is_a_500( + sign: Callable[[OpenSearchVectorStoreConfig], object], monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + _isolate_from_ambient_aws(monkeypatch, tmp_path) + monkeypatch.setenv("AWS_CONTAINER_CREDENTIALS_FULL_URI", "http://127.0.0.1:9/v1/credentials") + monkeypatch.setenv("AWS_CONTAINER_AUTHORIZATION_TOKEN_FILE", str(tmp_path / "missing-token")) + + with pytest.raises(litellm.InternalServerError, match="could not load or refresh the AWS credentials") as error: + sign(OpenSearchVectorStoreConfig()) + + assert error.value.status_code == 500 + + +EXPIRED_METADATA: Final = MappingProxyType( + {"access_key": "AKIDEXAMPLE", "secret_key": "secret", "token": "token", "expiry_time": "2020-01-01T00:00:00Z"} +) + + +def _refresh_returning_expired_credentials() -> dict[str, str]: + return dict(EXPIRED_METADATA) + + +def _refresh_after_the_token_file_is_gone() -> dict[str, str]: + raise CredentialRetrievalError(provider="container-role", error_msg=CREDENTIAL_SOURCE_OUTPUT) + + +class _CredentialSourceThatCannotRefresh(BaseAWSLLM): + def __init__(self, refresh: Callable[[], dict[str, str]]) -> None: + super().__init__() + self._refresh = refresh + + def resolve_credentials(self, auth_params: AwsAuthParams, aws_region_name: str | None) -> Credentials: + return RefreshableCredentials.create_from_metadata( + metadata=dict(EXPIRED_METADATA), refresh_using=self._refresh, method="container-role" + ) + + +@pytest.mark.parametrize( + "refresh", [_refresh_returning_expired_credentials, _refresh_after_the_token_file_is_gone] +) +@pytest.mark.parametrize( + "sign", + [ + pytest.param( + lambda config: config.validate_environment({}, GenericLiteLLMParams(api_base=SERVERLESS)), + id="vector-store-search", + ), + pytest.param( + lambda config: config.sign_http_request( + "GET", f"{SERVERLESS}/docs/_mapping", b"", "application/json", {"api_base": SERVERLESS} + ), + id="index-pass-through", + ), + ], +) +def test_credentials_that_fail_to_refresh_at_signing_are_a_500( + sign: Callable[[OpenSearchVectorStoreConfig], object], refresh: Callable[[], dict[str, str]] +) -> None: + config: Final = OpenSearchVectorStoreConfig(aws=_CredentialSourceThatCannotRefresh(refresh)) + + with pytest.raises(litellm.InternalServerError, match="could not load or refresh the AWS credentials") as error: + sign(config) + + assert error.value.status_code == 500 + assert CREDENTIAL_SOURCE_OUTPUT not in str(error.value) + + +@pytest.mark.asyncio +async def test_the_index_pass_through_answers_500_when_credentials_fail_to_refresh() -> None: + refused: Final = await prepare_index_request( + config=OpenSearchVectorStoreConfig(aws=_CredentialSourceThatCannotRefresh(_refresh_after_the_token_file_is_gone)), + method="GET", + path=IndexPath(index_name="docs", api=("_mapping",)), + operation=index_operation("GET", ("_mapping",)), + provider_index="prod-docs", + content_type=None, + query=(), + body=b"", + store_params={"api_base": SERVERLESS}, + can_read=True, + ) + + assert isinstance(refused, IndexRefusal) + assert (refused.status_code, CREDENTIAL_SOURCE_OUTPUT in refused.message) == (500, False) + + @pytest.mark.parametrize("asynchronous", [False, True]) @pytest.mark.asyncio async def test_opensearch_index_name_picks_the_index_so_the_store_id_can_differ(asynchronous: bool) -> None: diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index 8440d3d2ab2..70208bec0e5 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -5094,6 +5094,70 @@ async def test_centralized_common_checks_carries_team_and_user_budget_state_on_t } +@pytest.mark.asyncio +@pytest.mark.parametrize("prometheus_registered", [True, False], ids=["prometheus_logger", "no_prometheus_logger"]) +async def test_centralized_common_checks_carry_the_committed_spend_beside_the_stale_cached_spend( + monkeypatch: pytest.MonkeyPatch, prometheus_registered: bool +): + """The enforcement counters also hold budget holds, the committed twins only charged spend.""" + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as _proxy_server_mod + from litellm.integrations.prometheus import PrometheusLogger + from litellm.proxy.spend_tracking.committed_spend import committed_spend_key + from litellm.proxy.utils import ProxyLogging + from litellm.redis.stores.connection import connect + from litellm.redis.stores.spend_counters import SpendCounterStore + from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys + + for callback_list in ("success_callback", "failure_callback", "_async_success_callback", "_async_failure_callback"): + monkeypatch.setattr(litellm, callback_list, []) + monkeypatch.setattr(litellm, "callbacks", [MagicMock(spec=PrometheusLogger)] if prometheus_registered else []) + token: Final = UserAPIKeyAuth(api_key="sk-live-spend", max_budget=10.0, spend=0.0, team_id="t-live", team_spend=0.0) + user_api_key_cache: Final = UserApiKeyCache() + user_api_key_cache.set_cache( + key="team_id:t-live", value=LiteLLM_TeamTableCachedObj(team_id="t-live", max_budget=20.0, spend=0.0) + ) + spend_counter_store: Final = SpendCounterStore(connect(None).spend_counters()) + for counter_key, held, committed in ((f"spend:key:{token.token}", 0.67, 0.37), ("spend:team:t-live", 1.55, 1.25)): + spend_counter_store.set_local(counter_key, held) + spend_counter_store.set_local(committed_spend_key(counter_key), committed) + attrs: Final = { + **_proxy_attrs_for_centralized_checks(), + "llm_router": None, + "prisma_client": MagicMock(), + "user_api_key_cache": user_api_key_cache, + "spend_counter_store": spend_counter_store, + "proxy_logging_obj": ProxyLogging(user_api_key_cache=user_api_key_cache), + } + originals: Final = {attribute: getattr(_proxy_server_mod, attribute, None) for attribute in attrs} + request: Final = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + try: + for attribute, value in attrs.items(): + setattr(_proxy_server_mod, attribute, value) + await run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 1000}, + route="/chat/completions", + ) + finally: + for attribute, value in originals.items(): + setattr(_proxy_server_mod, attribute, value) + + assert token.committed_budget_spend == ( + CommittedBudgetSpend( + counter_keys=EntitySpendCounterKeys(key=f"spend:key:{token.token}", team=f"spend:team:{token.team_id}"), + key=0.37, + team=1.25, + ) + if prometheus_registered + else None + ) + assert (token.spend, token.team_spend) == (0.0, 0.0) + + def _end_user_budget_row(budget_id: str, max_budget: float) -> MagicMock: row = MagicMock() row.dict = lambda: {"budget_id": budget_id, "max_budget": max_budget} diff --git a/tests/unit/proxy/client/cli/test_auth_commands.py b/tests/unit/proxy/client/cli/test_auth_commands.py index 5d5ebcf5dfd..18c2c6ce8af 100644 --- a/tests/unit/proxy/client/cli/test_auth_commands.py +++ b/tests/unit/proxy/client/cli/test_auth_commands.py @@ -1,3 +1,4 @@ +import dataclasses import json import os import stat @@ -16,6 +17,7 @@ from litellm.litellm_core_utils.cli_keyring import ( DISABLE_KEYRING_ENV_VAR, KeyringDisabled, KeyringNotInstalled, + KeyringUnreachable, SecretErased, SecretStored, ) @@ -1344,6 +1346,22 @@ class TestKeychainBackedCommands: assert "could not be removed" in result.output assert not (isolated_home / ".litellm" / "token.json").exists() + @pytest.mark.parametrize("failure", [KeyringDisabled(), KeyringUnreachable(), KeyringNotInstalled()]) + def test_logout_without_a_usable_keychain_leaves_every_command_signed_out( + self, isolated_home, secret_vault_factory, failure + ): + vault = secret_vault_factory(available=False, failure=failure) + self._login(vault) + obj = {"base_url": "https://test.example.com", "secret_vault": vault} + self.runner.invoke(logout, obj=obj) + + whoami_result = self.runner.invoke(whoami, obj=obj) + print_result = self.runner.invoke(print_token, obj=obj) + + assert whoami_result.output == "Not authenticated. Run 'lite login' to authenticate.\n" + assert print_result.exit_code == 1 + assert print_result.output == "Not authenticated. Run 'lite login'.\n" + def test_print_token_explains_a_locked_keychain_instead_of_printing_nothing( self, isolated_home, secret_vault_factory ): @@ -1933,6 +1951,287 @@ class TestPkceLogoutCommand: assert "Logged out successfully" not in result.output clear.assert_not_called() + def test_a_logout_that_reaches_the_keychain_late_still_revokes_the_login( + self, isolated_home, secret_vault_factory + ): + unlocked = secret_vault_factory() + save_cli_token(CliTokenRecord(**_pkce_record()), vault=unlocked) + locked = secret_vault_factory(blob=unlocked.blob, available=False) + + with patch("litellm.proxy.client.cli.commands.auth.requests.Session", _FakeSession): + self.runner.invoke(logout, obj={"secret_vault": locked}) + result = self.runner.invoke(logout, obj={"secret_vault": unlocked}) + + assert result.output == "Logged out successfully. Authentication token cleared.\n" + assert unlocked.blob is None + assert _FakeSession.instances[-1].posts == [ + ( + f"{PKCE_BASE_URL}/revoke", + {"token": "llm_srefresh_old", "token_type_hint": "refresh_token", "client_id": "llm_dcrc_abc"}, + ) + ] + + def test_a_logout_that_cannot_note_the_keychain_login_leaves_it_for_the_next_logout_to_revoke( + self, isolated_home, secret_vault_factory + ): + unlocked = secret_vault_factory() + save_cli_token(CliTokenRecord(**_pkce_record()), vault=unlocked) + locked = secret_vault_factory(blob=unlocked.blob, available=False) + (isolated_home / ".litellm" / "keychain_not_cleared").mkdir() + + with patch("litellm.proxy.client.cli.commands.auth.requests.Session", _FakeSession): + locked_logout = self.runner.invoke(logout, obj={"secret_vault": locked}) + result = self.runner.invoke(logout, obj={"secret_vault": unlocked}) + + assert locked_logout.exit_code == 0, locked_logout.output + assert "Not logged out" in locked_logout.output + assert "Delete that file" not in locked_logout.output + assert result.output == "Logged out successfully. Authentication token cleared.\n" + assert unlocked.blob is None + assert [post for session in _FakeSession.instances for post in session.posts] == [ + ( + f"{PKCE_BASE_URL}/revoke", + {"token": "llm_srefresh_old", "token_type_hint": "refresh_token", "client_id": "llm_dcrc_abc"}, + ) + ] + + def test_a_login_after_a_logout_that_missed_the_keychain_revokes_the_login_it_replaces( + self, isolated_home, secret_vault_factory + ): + unlocked = secret_vault_factory() + save_cli_token(CliTokenRecord(**_pkce_record()), vault=unlocked) + locked = secret_vault_factory(blob=unlocked.blob, available=False) + + with ( + patch("litellm.proxy.client.cli.commands.auth.requests.Session", _FakeSession), + patch("litellm.proxy.client.cli.commands.auth.run_pkce_login", return_value=_pkce_credential()), + patch("litellm.proxy.client.cli.interface.show_commands"), + ): + self.runner.invoke(logout, obj={"secret_vault": locked}) + result = self.runner.invoke( + login, ["--pkce"], obj={"base_url": PKCE_BASE_URL, "secret_vault": unlocked} + ) + + assert result.exit_code == 0, result.output + assert "Could not revoke" not in result.output + assert [post for session in _FakeSession.instances for post in session.posts] == [ + ( + f"{PKCE_BASE_URL}/revoke", + {"token": "llm_srefresh_old", "token_type_hint": "refresh_token", "client_id": "llm_dcrc_abc"}, + ) + ] + + def test_a_second_logout_that_misses_the_keychain_still_lets_the_first_login_be_revoked( + self, isolated_home, secret_vault_factory + ): + unlocked = secret_vault_factory() + save_cli_token(CliTokenRecord(**_pkce_record()), vault=unlocked) + locked = secret_vault_factory(blob=unlocked.blob, available=False) + + with ( + patch("litellm.proxy.client.cli.commands.auth.requests.Session", _FakeSession), + patch("litellm.proxy.client.cli.commands.auth.run_pkce_login", return_value=_pkce_credential()), + patch("litellm.proxy.client.cli.interface.show_commands"), + ): + self.runner.invoke(logout, obj={"secret_vault": locked}) + self.runner.invoke(login, ["--pkce"], obj={"base_url": PKCE_BASE_URL, "secret_vault": locked}) + self.runner.invoke(logout, obj={"secret_vault": locked}) + result = self.runner.invoke(logout, obj={"secret_vault": unlocked}) + + assert result.output == "Logged out successfully. Authentication token cleared.\n" + assert unlocked.blob is None + assert [post["token"] for session in _FakeSession.instances for _, post in session.posts] == [ + "llm_srefresh_fresh", + "llm_srefresh_old", + ] + + @pytest.mark.parametrize("logout_while_locked", [True, False]) + def test_a_logout_after_a_locked_keychain_login_revokes_the_login_the_keychain_kept( + self, isolated_home, secret_vault_factory, logout_while_locked + ): + unlocked = secret_vault_factory() + save_cli_token(CliTokenRecord(**_pkce_record()), vault=unlocked) + locked = secret_vault_factory(blob=unlocked.blob, available=False) + + with ( + patch("litellm.proxy.client.cli.commands.auth.requests.Session", _FakeSession), + patch("litellm.proxy.client.cli.commands.auth.run_pkce_login", return_value=_pkce_credential()), + patch("litellm.proxy.client.cli.interface.show_commands"), + ): + if logout_while_locked: + self.runner.invoke(cli, ["logout"], obj={"secret_vault": locked}) + self.runner.invoke(cli, ["--base-url", PKCE_BASE_URL, "login", "--pkce"], obj={"secret_vault": locked}) + result = self.runner.invoke(cli, ["logout"], obj={"secret_vault": unlocked}) + + assert result.output == "Logged out successfully. Authentication token cleared.\n" + assert unlocked.blob is None + assert list((isolated_home / ".litellm").iterdir()) == [] + assert [post for session in _FakeSession.instances for post in session.posts] == [ + ( + f"{PKCE_BASE_URL}/revoke", + {"token": token, "token_type_hint": "refresh_token", "client_id": "llm_dcrc_abc"}, + ) + for token in ("llm_srefresh_old", "llm_srefresh_fresh") + ] + + def test_a_proxy_that_cannot_revoke_the_earlier_login_never_holds_back_the_current_logout( + self, isolated_home, secret_vault_factory + ): + class _EarlierProxyUnavailableSession(_FakeSession): + def post(self, url, *, data=None, json=None, timeout, allow_redirects): + self.posts.append((url, data)) + if data["token"] == "llm_srefresh_old": + return _FakeHttpResponse(503, {"error": "temporarily_unavailable"}) + return _FakeHttpResponse(200, {}) + + unlocked = secret_vault_factory() + save_cli_token(CliTokenRecord(**_pkce_record()), vault=unlocked) + locked = secret_vault_factory(blob=unlocked.blob, available=False) + + with ( + patch("litellm.proxy.client.cli.commands.auth.requests.Session", _EarlierProxyUnavailableSession), + patch("litellm.proxy.client.cli.commands.auth.run_pkce_login", return_value=_pkce_credential()), + patch("litellm.proxy.client.cli.interface.show_commands"), + ): + self.runner.invoke(cli, ["logout"], obj={"secret_vault": locked}) + self.runner.invoke(cli, ["--base-url", PKCE_BASE_URL, "login", "--pkce"], obj={"secret_vault": locked}) + result = self.runner.invoke(cli, ["logout"], obj={"secret_vault": unlocked}) + + assert result.exit_code == 0, result.output + assert ( + "Could not revoke the previous login's refresh token on the proxy (revocation failed with 503" + in result.output + ) + assert result.output.endswith("Logged out successfully. Authentication token cleared.\n") + assert unlocked.blob is None + assert list((isolated_home / ".litellm").iterdir()) == [] + assert [post["token"] for session in _FakeSession.instances for _, post in session.posts] == [ + "llm_srefresh_old", + "llm_srefresh_fresh", + ] + + def test_a_logout_whose_only_login_is_the_one_left_in_the_keychain_waits_for_its_revocation( + self, isolated_home, secret_vault_factory + ): + class _UnavailableSession(_FakeSession): + def __init__(self): + super().__init__() + self.response = _FakeHttpResponse(503, {"error": "temporarily_unavailable"}) + + unlocked = secret_vault_factory() + save_cli_token(CliTokenRecord(**_pkce_record()), vault=unlocked) + locked = secret_vault_factory(blob=unlocked.blob, available=False) + + with patch("litellm.proxy.client.cli.commands.auth.requests.Session", _UnavailableSession): + self.runner.invoke(cli, ["logout"], obj={"secret_vault": locked}) + result = self.runner.invoke(cli, ["logout"], obj={"secret_vault": unlocked}) + + assert result.exit_code == 1 + assert "Nothing was cleared; run `lite logout` again shortly." in result.output + assert unlocked.blob is not None + + def test_a_login_after_a_locked_keychain_login_revokes_the_login_the_keychain_kept( + self, isolated_home, secret_vault_factory + ): + unlocked = secret_vault_factory() + save_cli_token(CliTokenRecord(**_pkce_record()), vault=unlocked) + locked = secret_vault_factory(blob=unlocked.blob, available=False) + + with ( + patch("litellm.proxy.client.cli.commands.auth.requests.Session", _FakeSession), + patch("litellm.proxy.client.cli.commands.auth.run_pkce_login", return_value=_pkce_credential()), + patch("litellm.proxy.client.cli.interface.show_commands"), + ): + self.runner.invoke(cli, ["--base-url", PKCE_BASE_URL, "login", "--pkce"], obj={"secret_vault": locked}) + result = self.runner.invoke( + cli, ["--base-url", PKCE_BASE_URL, "login", "--pkce"], obj={"secret_vault": unlocked} + ) + + assert result.exit_code == 0, result.output + assert "Could not revoke" not in result.output + assert [post["token"] for session in _FakeSession.instances for _, post in session.posts] == [ + "llm_srefresh_old", + "llm_srefresh_fresh", + ] + + @pytest.mark.parametrize("seconds_the_clock_stepped_back", [0, 600]) + def test_a_command_between_the_unlock_and_the_logout_keeps_the_login_the_keychain_kept_revocable( + self, isolated_home, secret_vault_factory, seconds_the_clock_stepped_back + ): + unlocked = secret_vault_factory() + first_login = _pkce_record(timestamp=time.time() + seconds_the_clock_stepped_back) + save_cli_token(CliTokenRecord(**first_login), vault=unlocked) + locked = secret_vault_factory(blob=unlocked.blob, available=False) + + with ( + patch("litellm.proxy.client.cli.commands.auth.requests.Session", _FakeSession), + patch("litellm.proxy.client.cli.commands.auth.run_pkce_login", return_value=_pkce_credential()), + patch("litellm.proxy.client.cli.interface.show_commands"), + ): + self.runner.invoke(cli, ["logout"], obj={"secret_vault": locked}) + self.runner.invoke(cli, ["--base-url", PKCE_BASE_URL, "login", "--pkce"], obj={"secret_vault": locked}) + printed = self.runner.invoke(cli, ["auth", "print-token"], obj={"secret_vault": unlocked}) + result = self.runner.invoke(cli, ["logout"], obj={"secret_vault": unlocked}) + + assert printed.stdout == "sk-cli-fresh\n" + assert result.output == "Logged out successfully. Authentication token cleared.\n" + assert [post["token"] for session in _FakeSession.instances for _, post in session.posts] == [ + "llm_srefresh_old", + "llm_srefresh_fresh", + ] + + @pytest.mark.parametrize("revocation_status", [200, 503]) + def test_a_renewal_that_takes_the_keychain_entry_revokes_the_login_left_in_it( + self, isolated_home, secret_vault_factory, revocation_status + ): + class _RenewingSession(_FakeSession): + def post(self, url, *, data=None, json=None, timeout, allow_redirects): + self.posts.append((url, data)) + if url.endswith("/token"): + return _FakeHttpResponse(200, PKCE_TOKEN_RESPONSE) + return _FakeHttpResponse(revocation_status, {}) + + unlocked = secret_vault_factory() + save_cli_token(CliTokenRecord(**_pkce_record()), vault=unlocked) + locked = secret_vault_factory(blob=unlocked.blob, available=False) + expiring_login = dataclasses.replace(_pkce_credential(), expires_at=time.time() + 30) + + with ( + patch("litellm.proxy.client.cli.commands.auth.requests.Session", _RenewingSession), + patch("litellm.proxy.client.cli.commands.auth.run_pkce_login", return_value=expiring_login), + patch("litellm.proxy.client.cli.interface.show_commands"), + ): + self.runner.invoke(cli, ["logout"], obj={"secret_vault": locked}) + self.runner.invoke(cli, ["--base-url", PKCE_BASE_URL, "login", "--pkce"], obj={"secret_vault": locked}) + printed = self.runner.invoke(cli, ["auth", "print-token"], obj={"secret_vault": unlocked}) + self.runner.invoke(cli, ["logout"], obj={"secret_vault": unlocked}) + + assert printed.stdout == "sk-cli-rotated\n" + assert ("Could not revoke" in printed.stderr) is (revocation_status != 200) + assert [ + post["token"] for session in _FakeSession.instances for url, post in session.posts if url.endswith("/revoke") + ] == ["llm_srefresh_old", "llm_srefresh_rotated"] + + def test_a_logout_after_the_keychain_refused_the_erase_does_not_revoke_the_login_again( + self, isolated_home, secret_vault_factory + ): + vault = secret_vault_factory(erasable=False) + save_cli_token(CliTokenRecord(**_pkce_record()), vault=vault) + + with patch("litellm.proxy.client.cli.commands.auth.requests.Session", _FakeSession): + first = self.runner.invoke(logout, obj={"secret_vault": vault}) + second = self.runner.invoke(logout, obj={"secret_vault": vault}) + + assert first.exit_code == 0, first.output + assert second.exit_code == 0, second.output + assert "still in the OS keychain" in second.output + assert [post for session in _FakeSession.instances for post in session.posts] == [ + ( + f"{PKCE_BASE_URL}/revoke", + {"token": "llm_srefresh_old", "token_type_hint": "refresh_token", "client_id": "llm_dcrc_abc"}, + ) + ] + def test_logout_of_a_classic_token_makes_no_request(self): with ( patch("litellm.proxy.client.cli.commands.auth.load_token", return_value={"key": "sk-classic"}), diff --git a/tests/unit/proxy/client/cli/test_global_options.py b/tests/unit/proxy/client/cli/test_global_options.py index 8a1e9a3d141..2acffb704ca 100644 --- a/tests/unit/proxy/client/cli/test_global_options.py +++ b/tests/unit/proxy/client/cli/test_global_options.py @@ -253,6 +253,7 @@ def test_lite_goes_back_to_localhost_after_logout_without_a_keychain(cli_runner, result = cli_runner.invoke(cli, [], obj={"secret_vault": no_keychain}, input="quit\n") assert "Connected to LiteLLM server: http://localhost:4000" in result.output + assert "keychain" not in result.output def test_corrupt_config_file_falls_back_to_default(cli_runner, isolated_home): diff --git a/tests/unit/proxy/common_utils/test_check_responses_cost.py b/tests/unit/proxy/common_utils/test_check_responses_cost.py index 21685f71923..4404a0d926c 100644 --- a/tests/unit/proxy/common_utils/test_check_responses_cost.py +++ b/tests/unit/proxy/common_utils/test_check_responses_cost.py @@ -67,6 +67,7 @@ class _StoredResponse: id: str unified_object_id: str status: str + created_at: datetime = datetime(2026, 1, 2) model_object_id: str = "" file_purpose: str = "response" file_object: Mapping[str, object] = dataclasses.field(default_factory=lambda: {"model": "gpt-5"}) @@ -2081,7 +2082,14 @@ def _key_row(metadata: Mapping[str, object]) -> object: def _team_row(metadata: Mapping[str, object]) -> object: from types import SimpleNamespace - return SimpleNamespace(team_alias=None, organization_id=None, spend=0.0, max_budget=None, metadata=dict(metadata)) + return SimpleNamespace( + team_alias=None, + organization_id=None, + spend=0.0, + max_budget=None, + metadata=dict(metadata), + created_at=datetime(2026, 1, 1), + ) def _langfuse_callback_settings(callback_vars: Mapping[str, object]) -> Mapping[str, object]: diff --git a/tests/unit/proxy/common_utils/test_encrypt_decrypt_utils.py b/tests/unit/proxy/common_utils/test_encrypt_decrypt_utils.py index 1163090c510..5b0d5f014f8 100644 --- a/tests/unit/proxy/common_utils/test_encrypt_decrypt_utils.py +++ b/tests/unit/proxy/common_utils/test_encrypt_decrypt_utils.py @@ -6,8 +6,11 @@ gate, and the backward-compatibility guarantees that let legacy XSalsa20-Poly130 (nacl) ciphertext and new AES values coexist and decrypt correctly. """ +import _hashlib import base64 +import hmac import re +from typing import Final import pytest @@ -17,6 +20,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_bearer_token, decrypt_if_encrypted_with, decrypt_value_helper, + derive_key_from_salt_key, encrypt_bearer_token, encrypt_value, encrypt_value_helper, @@ -266,3 +270,29 @@ def test_bearer_token_uses_only_header_safe_characters(length: int): assert re.fullmatch(r"kind_a_[A-Za-z0-9_-]+", token), token assert decrypt_bearer_token(token, prefix="kind_a_") == "x" * length + + +_needs_shared_openssl: Final = pytest.mark.skipif( + not hasattr(_hashlib, "__file__"), reason="this Python links OpenSSL statically, out of ctypes reach" +) + + +@_needs_shared_openssl +@pytest.mark.usefixtures("fips_openssl") +@pytest.mark.parametrize("salt_key", ["sk-" + bytes(range(32)).hex(), bytes(range(32)).hex().upper()]) +def test_derive_key_from_salt_key_is_sp800_108_counter_mode_hmac_sha256( + monkeypatch: pytest.MonkeyPatch, salt_key: str +) -> None: + label: Final = b"litellm-test:label:" + context: Final = b"sk-test-master-key" + monkeypatch.setenv("LITELLM_SALT_KEY", salt_key) + fixed_input: Final = (1).to_bytes(4, "big") + label + b"\x00" + context + (256).to_bytes(4, "big") + assert derive_key_from_salt_key(label, context) == hmac.digest(bytes(range(32)), fixed_input, "sha256").hex() + + +@_needs_shared_openssl +@pytest.mark.skipif(_hashlib.get_fips_mode() == 1, reason="this OpenSSL runs a FIPS provider") +def test_derive_key_from_salt_key_refuses_an_openssl_without_a_fips_provider(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-" + bytes(range(32)).hex()) + with pytest.raises(ValueError, match="OpenSSL could not run KBKDF"): + derive_key_from_salt_key(b"litellm-test:label:", b"sk-test-master-key") diff --git a/tests/unit/proxy/common_utils/test_load_config_utils.py b/tests/unit/proxy/common_utils/test_load_config_utils.py index 927e6ee511d..d9abe5bbc0d 100644 --- a/tests/unit/proxy/common_utils/test_load_config_utils.py +++ b/tests/unit/proxy/common_utils/test_load_config_utils.py @@ -20,10 +20,9 @@ class TestGetFileContentsFromS3: """Test suite for S3 config loading functionality.""" @patch("boto3.client") - @patch("litellm.main.bedrock_converse_chat_completion") @patch("yaml.safe_load") def test_get_file_contents_from_s3_no_temp_file_creation( - self, mock_yaml_load, mock_bedrock, mock_boto3_client + self, mock_yaml_load, mock_boto3_client, monkeypatch ): """ Test that get_file_contents_from_s3 doesn't create temporary files @@ -33,12 +32,9 @@ class TestGetFileContentsFromS3: Relevant issue/PR: https://github.com/BerriAI/litellm/pull/12078 """ - # Mock credentials - mock_credentials = MagicMock() - mock_credentials.access_key = "test_access_key" - mock_credentials.secret_key = "test_secret_key" - mock_credentials.token = "test_token" - mock_bedrock.get_credentials.return_value = mock_credentials + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "test_access_key") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "test_secret_key") + monkeypatch.setenv("AWS_SESSION_TOKEN", "test_token") # Mock S3 client and response mock_s3_client = MagicMock() diff --git a/tests/unit/proxy/common_utils/test_registry_read_through.py b/tests/unit/proxy/common_utils/test_registry_read_through.py index da7b1f497c9..6c9104bb711 100644 --- a/tests/unit/proxy/common_utils/test_registry_read_through.py +++ b/tests/unit/proxy/common_utils/test_registry_read_through.py @@ -433,7 +433,7 @@ async def test_resync_model_deployments_mutates_router_under_model_reconcile_loc lock_states: list[bool] = [] def record_add_deployment(db_models) -> None: - lock_states.append(proxy_server.MODEL_RECONCILE_LOCK.locked()) + lock_states.append(proxy_server.DB_RECONCILE_LOCK.locked()) monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) monkeypatch.setattr(proxy_server, "store_model_in_db", True) @@ -443,7 +443,7 @@ async def test_resync_model_deployments_mutates_router_under_model_reconcile_loc assert await _resync_model_deployments("lock-scope-model") is True assert lock_states == [True] - assert not proxy_server.MODEL_RECONCILE_LOCK.locked() + assert not proxy_server.DB_RECONCILE_LOCK.locked() @pytest.mark.asyncio @@ -601,7 +601,7 @@ async def test_resync_model_deployments_by_ids_installs_only_the_published_rows( assert _served_ids(router.registry().deployments()) == {"primary-id", "m-1", "m-2"} assert _served_ids(proxy_server.llm_model_list) == {"primary-id", "m-1", "m-2"} - assert not proxy_server.MODEL_RECONCILE_LOCK.locked() + assert not proxy_server.DB_RECONCILE_LOCK.locked() @pytest.mark.asyncio @@ -619,7 +619,7 @@ async def test_resync_model_deployments_by_ids_leaves_the_router_alone_for_delet assert _served_ids(router.registry().deployments()) == {"primary-id"} assert proxy_server.llm_model_list is None - assert not proxy_server.MODEL_RECONCILE_LOCK.locked() + assert not proxy_server.DB_RECONCILE_LOCK.locked() @pytest.mark.asyncio @@ -1021,7 +1021,7 @@ async def test_unknown_model_name_read_through_reads_the_replica_without_waiting monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) monkeypatch.setattr(proxy_server, "store_model_in_db", True) monkeypatch.setattr(proxy_server, "llm_router", MagicMock()) - monkeypatch.setattr(proxy_server, "MODEL_RECONCILE_LOCK", reload_in_progress) + monkeypatch.setattr(proxy_server, "DB_RECONCILE_LOCK", reload_in_progress) async with reload_in_progress: found: Final = await asyncio.wait_for(_resync_model_deployments("no-such-model"), timeout=2) @@ -1042,7 +1042,7 @@ async def test_resync_model_deployments_by_ids_installs_the_row_a_concurrent_rel router: Final = _router_serving("primary-id") _serve_db_models(monkeypatch, _Db(table), router) reload_in_progress: Final = asyncio.Lock() - monkeypatch.setattr(proxy_server, "MODEL_RECONCILE_LOCK", reload_in_progress) + monkeypatch.setattr(proxy_server, "DB_RECONCILE_LOCK", reload_in_progress) async with reload_in_progress: targeted_load: Final = asyncio.create_task(resync_model_deployments_by_ids(("m-1",))) @@ -1132,4 +1132,4 @@ async def test_resync_model_deployments_by_ids_installs_the_ordinary_rows_queued assert _served_ids(proxy_server.llm_model_list) == {"primary-id", "backup-id"} assert table.reads == 1 assert db.other_tables_read == [] - assert not proxy_server.MODEL_RECONCILE_LOCK.locked() + assert not proxy_server.DB_RECONCILE_LOCK.locked() diff --git a/tests/unit/proxy/common_utils/test_reset_budget_job.py b/tests/unit/proxy/common_utils/test_reset_budget_job.py index 4a4490afc33..aaa686f2a98 100644 --- a/tests/unit/proxy/common_utils/test_reset_budget_job.py +++ b/tests/unit/proxy/common_utils/test_reset_budget_job.py @@ -1438,7 +1438,47 @@ def test_reset_budget_for_keys_invalidates_redis_counter(reset_budget_job, mock_ asyncio.run(reset_budget_job.reset_budget_for_litellm_keys()) - assert "spend:key:sk-abc" in counter_cache.deleted() + assert {"spend:key:sk-abc", "spend:committed:key:sk-abc"} <= counter_cache.deleted() + + +def _expired_key_row() -> object: + return type( + "Key", + (), + { + "spend": 100.0, + "budget_duration": "30d", + "budget_reset_at": datetime.now(timezone.utc), + "id": "key-1", + "token": "sk-abc", + }, + ) + + +def test_reset_budget_for_keys_without_redis_clears_the_counter_and_its_committed_twin( + reset_budget_job, mock_prisma_client, monkeypatch +): + counter_cache: Final = _make_counter_invalidation_job(monkeypatch) + counter_cache.redis = None + counter_cache.memory.update({"spend:key:sk-abc": 100.0, "spend:committed:key:sk-abc": 100.0}) + mock_prisma_client.data["key"] = [_expired_key_row()] + + asyncio.run(reset_budget_job.reset_budget_for_litellm_keys()) + + assert counter_cache.memory == {} + + +def test_a_redis_failure_on_the_key_reset_still_drops_its_committed_twin_locally( + reset_budget_job, mock_prisma_client, monkeypatch +): + counter_cache: Final = _make_counter_invalidation_job(monkeypatch) + counter_cache.failing = frozenset({"delete"}) + counter_cache.memory.update({"spend:key:sk-abc": 100.0, "spend:committed:key:sk-abc": 100.0}) + mock_prisma_client.data["key"] = [_expired_key_row()] + + asyncio.run(reset_budget_job.reset_budget_for_litellm_keys()) + + assert counter_cache.memory == {} def test_reset_budget_for_users_invalidates_redis_counter(reset_budget_job, mock_prisma_client, monkeypatch): @@ -1790,7 +1830,9 @@ def test_enduser_invalidation_is_paged_and_batched(reset_budget_job, mock_prisma assert counter_cache.user_api_key_cache.async_delete_cache_keys.await_count == 6 invalidated: Final = counter_cache.deleted() - assert invalidated == {f"spend:end_user:cust-{i:06d}" for i in range(population)} + assert invalidated == {f"spend:end_user:cust-{i:06d}" for i in range(population)} | { + f"spend:committed:end_user:cust-{i:06d}" for i in range(population) + } evicted: Final = { key for call in counter_cache.user_api_key_cache.async_delete_cache_keys.await_args_list for key in call.args[0] } diff --git a/tests/unit/proxy/conftest.py b/tests/unit/proxy/conftest.py index d08a98242ef..8d8781be0a8 100644 --- a/tests/unit/proxy/conftest.py +++ b/tests/unit/proxy/conftest.py @@ -523,6 +523,7 @@ def postgres_span_names() -> Iterator[Callable[[], Awaitable[tuple[str, ...]]]]: def fips_openssl(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(_hashlib, "get_fips_mode", lambda: 1) monkeypatch.setattr(_hashlib, "scrypt", MagicMock(side_effect=ValueError("[digital envelope routines] unsupported"))) + monkeypatch.setattr("litellm.proxy.common_utils.encrypt_decrypt_utils._KBKDF_PROPERTIES", None) @pytest.fixture diff --git a/tests/unit/proxy/credential_endpoints/test_endpoints.py b/tests/unit/proxy/credential_endpoints/test_endpoints.py index 10751def88d..f08fab25c57 100644 --- a/tests/unit/proxy/credential_endpoints/test_endpoints.py +++ b/tests/unit/proxy/credential_endpoints/test_endpoints.py @@ -1,6 +1,8 @@ """Tests for the credential management endpoints.""" +import asyncio import json +from collections.abc import Awaitable, Callable from contextlib import contextmanager from functools import reduce from types import SimpleNamespace @@ -19,11 +21,22 @@ import litellm from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.litellm_logging import get_masked_values from litellm.llms.github_copilot.per_user_auth import _SESSION_CACHE, _session_cache_key -from litellm.models.credentials import CredentialSource, UserConnectionFlowHandle +from litellm.models.credentials import ( + CreateCredentialItem, + CredentialSource, + UpdateCredentialItem, + UserConnectionFlowHandle, +) from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper, encrypt_value_helper -from litellm.proxy.credential_endpoints.endpoints import _decode_flow_handle, get_llm_router +from litellm.proxy.credential_endpoints.endpoints import ( + _decode_flow_handle, + create_credential, + delete_credential, + get_llm_router, + update_credential, +) from litellm.proxy.proxy_server import app from litellm.types.utils import CredentialItem @@ -3345,3 +3358,101 @@ def test_get_credential_by_model_masks_the_deployment_credentials(monkeypatch): with pytest.raises(ProxyException) as missing: asyncio.run(get_credential_by_model(MagicMock(), MagicMock(), "no-such-id", caller)) assert missing.value.code == "404" + + +class TestCredentialWritesHoldTheReconcileLock: + """A config reload snapshots the credential table, awaits the config rows, then prunes and + upserts ``litellm.credential_list`` from that snapshot. A credential write that lands inside + that gap is undone in memory until the next reload, so each write has to queue behind + DB_RECONCILE_LOCK, the lock the reload already holds.""" + + @staticmethod + async def _assert_write_waits_for_the_lock( + monkeypatch: pytest.MonkeyPatch, call_endpoint: Callable[[], Awaitable[dict[str, object]]], write: str + ) -> dict[str, object]: + lock: Final = asyncio.Lock() + monkeypatch.setattr("litellm.proxy.proxy_server.DB_RECONCILE_LOCK", lock) + async with lock: + task: Final = asyncio.create_task(call_endpoint()) + for _ in range(50): + await asyncio.sleep(0) + assert not task.done(), ( + f"{write} did not wait for DB_RECONCILE_LOCK, so an in-flight config reload can undo it in memory" + ) + return await asyncio.wait_for(task, timeout=5) + + async def test_create_credential_waits_for_an_in_flight_reload( + self, monkeypatch: pytest.MonkeyPatch, credential_store: Callable[..., None] + ) -> None: + credential_store(create=AsyncMock(return_value=None)) + + async def call_endpoint() -> dict[str, object]: + return await create_credential( + request=MagicMock(spec=Request), + fastapi_response=MagicMock(), + credential=CreateCredentialItem( + credential_name="fresh", + credential_values={"api_key": "sk-new"}, + credential_info={"custom_llm_provider": "openai"}, + ), + user_api_key_dict=_as_admin(), + llm_router=None, + ) + + response: Final = await self._assert_write_waits_for_the_lock(monkeypatch, call_endpoint, "create") + + assert response["success"] is True + assert [credential.credential_name for credential in litellm.credential_list] == ["fresh"] + + async def test_update_credential_waits_for_an_in_flight_reload( + self, monkeypatch: pytest.MonkeyPatch, credential_store: Callable[..., None] + ) -> None: + stored: Final = CredentialItem( + credential_name="existing", + display_name="before", + credential_values={"api_key": "sk-old"}, + credential_info={"custom_llm_provider": "openai"}, + ) + credential_store( + in_memory=(stored,), + find_by_name=AsyncMock(return_value=stored), + update_by_name=AsyncMock(return_value=None), + ) + + async def call_endpoint() -> dict[str, object]: + return await update_credential( + request=MagicMock(spec=Request), + fastapi_response=MagicMock(), + credential=UpdateCredentialItem(display_name="after", credential_info={}), + credential_name="existing", + user_api_key_dict=_as_admin(), + llm_router=None, + ) + + response: Final = await self._assert_write_waits_for_the_lock(monkeypatch, call_endpoint, "update") + + assert response["success"] is True + assert [credential.display_name for credential in litellm.credential_list] == ["after"] + + async def test_delete_credential_waits_for_an_in_flight_reload( + self, monkeypatch: pytest.MonkeyPatch, credential_store: Callable[..., None] + ) -> None: + stored: Final = CredentialItem( + credential_name="doomed", + credential_values={"api_key": "sk-old"}, + credential_info={"custom_llm_provider": "openai"}, + ) + credential_store(in_memory=(stored,), delete_by_name=AsyncMock(return_value=MagicMock())) + + async def call_endpoint() -> dict[str, object]: + return await delete_credential( + request=MagicMock(spec=Request), + fastapi_response=MagicMock(), + credential_name="doomed", + user_api_key_dict=_as_admin(), + ) + + response: Final = await self._assert_write_waits_for_the_lock(monkeypatch, call_endpoint, "delete") + + assert response["success"] is True + assert litellm.credential_list == [] diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py b/tests/unit/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py index 9a14222c684..552806d309e 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py @@ -247,6 +247,88 @@ class TestContentFilterGuardrail: assert "[EMAIL_REDACTED]" in result[0] assert "test@example.com" not in result[0] + @pytest.mark.asyncio + @pytest.mark.parametrize( + "pattern_names", + [("aws_secret_key", "openai_api_key"), ("openai_api_key", "aws_secret_key")], + ) + async def test_apply_guardrail_masks_whole_key_when_another_pattern_matches_inside_it( + self, pattern_names: tuple[str, str] + ): + guardrail: Final = ContentFilterGuardrail( + guardrail_name="test-overlapping-patterns", + patterns=[ + ContentFilterPattern(pattern_type="prebuilt", pattern_name=name, action=ContentFilterAction.MASK) + for name in pattern_names + ], + ) + openai_key: Final = "sk-proj-" + "Ab3" * 13 + "x-" + "Cd4" * 11 + "T3BlbkFJ" + "Ef5" * 24 + "gh" + aws_secret_key: Final = "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY" + + guardrailed_inputs: Final = await guardrail.apply_guardrail( + inputs={"texts": [f"key {openai_key} and secret {aws_secret_key}"]}, + request_data={}, + input_type="request", + ) + + assert guardrailed_inputs.get("texts") == [ + "key [OPENAI_API_KEY_REDACTED] and secret [AWS_SECRET_KEY_REDACTED]" + ] + + @pytest.mark.asyncio + async def test_apply_guardrail_does_not_block_a_match_an_earlier_mask_pattern_already_covers(self): + guardrail: Final = ContentFilterGuardrail( + guardrail_name="test-mask-then-block", + patterns=[ + ContentFilterPattern( + pattern_type="prebuilt", pattern_name="generic_api_key", action=ContentFilterAction.MASK + ), + ContentFilterPattern( + pattern_type="prebuilt", pattern_name="aws_secret_key", action=ContentFilterAction.BLOCK + ), + ], + ) + guardrailed_inputs: Final = await guardrail.apply_guardrail( + inputs={"texts": ["config api_key=" + "A1b2C3d4E5" * 4]}, + request_data={}, + input_type="request", + ) + assert guardrailed_inputs.get("texts") == ["config [GENERIC_API_KEY_REDACTED]"] + + @pytest.mark.asyncio + async def test_apply_guardrail_masks_a_contextual_match_an_earlier_mask_moved_next_to_its_keyword(self): + guardrail: Final = ContentFilterGuardrail( + guardrail_name="test-contextual-after-mask", + patterns=[ + ContentFilterPattern(pattern_type="prebuilt", pattern_name="email", action=ContentFilterAction.MASK), + ContentFilterPattern(pattern_type="prebuilt", pattern_name="au_tfn", action=ContentFilterAction.MASK), + ], + ) + guardrailed_inputs: Final = await guardrail.apply_guardrail( + inputs={"texts": ["TFN jane.doe@example.com 123456782"]}, + request_data={}, + input_type="request", + ) + assert guardrailed_inputs.get("texts") == ["TFN [EMAIL_REDACTED] [AU_TFN_REDACTED]"] + + @pytest.mark.asyncio + async def test_apply_guardrail_counts_only_the_unnamed_regex_that_matched(self): + guardrail: Final = ContentFilterGuardrail( + guardrail_name="test-unnamed-regex-count", + patterns=[ + ContentFilterPattern(pattern_type="regex", pattern=r"Key\d+", action=ContentFilterAction.MASK), + ContentFilterPattern(pattern_type="regex", pattern=r"Token\d+", action=ContentFilterAction.MASK), + ], + ) + request_data: Final = {"metadata": {}} + await guardrail.apply_guardrail( + inputs={"texts": ["use Key1 here"]}, + request_data=request_data, + input_type="request", + ) + entry: Final = request_data["metadata"]["standard_logging_guardrail_information"][0] + assert entry["masked_entity_count"] == {"custom_regex": 1} + @pytest.mark.asyncio async def test_apply_guardrail_blocked_word_mask(self): """ @@ -919,12 +1001,21 @@ class TestContentFilterGuardrail: # masked_entity_count for email is the real count, not N×. assert entry["masked_entity_count"].get("email") == 1 - @staticmethod + @classmethod async def _collect_streamed_text( + cls, guardrail: ContentFilterGuardrail, chunks: list[str], metadata: dict[str, list[StandardLoggingGuardrailInformation]], ) -> str: + return "".join(await cls._collect_streamed_deltas(guardrail, chunks, metadata)) + + @staticmethod + async def _collect_streamed_deltas( + guardrail: ContentFilterGuardrail, + chunks: list[str], + metadata: dict[str, list[StandardLoggingGuardrailInformation]], + ) -> list[str]: from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices async def mock_stream(): @@ -951,7 +1042,7 @@ class TestContentFilterGuardrail: request_data={"messages": [], "model": "gpt-4o", "metadata": metadata}, ): yielded.append(chunk.choices[0].delta.content or "") - return "".join(yielded) + return yielded @pytest.mark.asyncio async def test_streaming_hook_scans_bounded_window_per_chunk(self): @@ -1068,6 +1159,53 @@ class TestContentFilterGuardrail: assert entry["guardrail_status"] == "guardrail_intervened" assert [d["keyword"] for d in entry["guardrail_response"]] == [phrase] + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("pattern_name", "key"), + [ + ("openai_api_key", "sk-proj-" + "a" * 74 + "T3BlbkFJ" + "b" * 74), + ("openai_api_key", "sk-svcacct-" + "a" * 200 + "T3BlbkFJ" + "b" * 200), + ("anthropic_api_key", "sk-ant-api03-" + "a" * 93 + "AA"), + ("gemini_api_key", "AQ." + "a" * 55), + ], + ) + async def test_streaming_hook_masks_whole_provider_key_split_across_chunks(self, pattern_name: str, key: str): + guardrail: Final = ContentFilterGuardrail( + guardrail_name="test-streaming-provider-key", + patterns=[ + ContentFilterPattern( + pattern_type="prebuilt", + pattern_name=pattern_name, + action=ContentFilterAction.MASK, + ) + ], + event_hook=GuardrailEventHooks.post_call, + ) + text: Final = f"Use {key} to call the API." + chunks: Final = [text[i : i + 4] for i in range(0, len(text), 4)] + + streamed: Final = await self._collect_streamed_text(guardrail, chunks, {}) + + assert streamed == f"Use [{pattern_name.upper()}_REDACTED] to call the API." + + @pytest.mark.asyncio + async def test_streaming_hook_keeps_default_holdback_without_provider_key_patterns(self): + guardrail: Final = ContentFilterGuardrail( + guardrail_name="test-streaming-older-credential-patterns", + patterns=[ + ContentFilterPattern(pattern_type="prebuilt", pattern_name=name, action=ContentFilterAction.MASK) + for name in ("aws_access_key", "aws_secret_key", "github_token", "slack_token", "generic_api_key") + ], + event_hook=GuardrailEventHooks.post_call, + ) + text: Final = "plain words with nothing secret in them. " * 12 + chunks: Final = [text[i : i + 10] for i in range(0, len(text), 10)] + + deltas: Final = await self._collect_streamed_deltas(guardrail, chunks, {}) + + assert "".join(deltas) == text + assert len(next(delta for delta in reversed(deltas) if delta)) <= CONTENT_FILTER_STREAMING_HOLDBACK_CHARS + @pytest.mark.asyncio async def test_streaming_hook_blocks_keyword_longer_than_scan_context(self): """ diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/content_filter/test_patterns.py b/tests/unit/proxy/guardrails/guardrail_hooks/content_filter/test_patterns.py index d702b9e0116..3c71e814fe1 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/content_filter/test_patterns.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/content_filter/test_patterns.py @@ -4,6 +4,8 @@ Tests for content filter pattern loading from JSON import json import os +import time +from typing import Final import pytest @@ -183,3 +185,169 @@ def test_eu_patterns_have_category(): assert ( pattern_name in eu_category_patterns ), f"Pattern {pattern_name} not in EU PII Patterns category" + + +def test_provider_key_patterns_are_offered_in_the_ui_as_credential_patterns() -> None: + ui_categories: Final = {pattern["name"]: pattern["category"] for pattern in get_pattern_metadata()} + assert [ui_categories.get(name) for name in ("openai_api_key", "anthropic_api_key", "gemini_api_key")] == [ + "Credential Patterns" + ] * 3 + + +@pytest.mark.parametrize( + ("pattern_name", "key"), + [ + pytest.param("openai_api_key", "sk-" + "a" * 20 + "T3BlbkFJ" + "b" * 20, id="openai-legacy"), + pytest.param("openai_api_key", "sk-proj-" + "a" * 74 + "T3BlbkFJ" + "b" * 74, id="openai-project"), + pytest.param("openai_api_key", "sk-proj-" + "a" * 20 + "T3BlbkFJ" + "b" * 20, id="openai-project-shortest"), + pytest.param("openai_api_key", "sk-svcacct-" + "a_b-" * 15 + "T3BlbkFJ" + "c-d_" * 15, id="openai-service-account"), + pytest.param("openai_api_key", "sk-admin-" + "a" * 58 + "T3BlbkFJ" + "b" * 58, id="openai-admin"), + pytest.param("openai_api_key", "sk-None-" + "a" * 20 + "T3BlbkFJ" + "b" * 20, id="openai-user-key"), + pytest.param("openai_api_key", "sk-proj-" + "a" * 20 + "T3BlbkFJ" + "b" * 19 + "-", id="openai-trailing-hyphen"), + pytest.param( + "anthropic_api_key", "sk-ant-api03-abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOP", id="anthropic-versioned" + ), + pytest.param( + "anthropic_api_key", "sk-ant-abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTU", id="anthropic-unversioned" + ), + pytest.param("anthropic_api_key", "sk-ant-" + "a" * 20, id="anthropic-shortest"), + pytest.param("anthropic_api_key", "sk-ant-api03-" + "a_b-" * 23 + "AA", id="anthropic-urlsafe"), + pytest.param( + "anthropic_api_key", "sk-ant-api03-abcdefghijklmnopqrstuvwxyz0123456789-", id="anthropic-trailing-hyphen" + ), + pytest.param("gemini_api_key", "AIzaSyD-9tSrke72I6_AQp2h8HI4-uE6QNKQV0", id="gemini-legacy"), + pytest.param("gemini_api_key", "AIza" + "a" * 33, id="gemini-legacy-shortest"), + pytest.param("gemini_api_key", "AIza" + "a" * 36, id="gemini-legacy-longest"), + pytest.param("gemini_api_key", "AIza" + "a" * 34 + "-", id="gemini-legacy-trailing-hyphen"), + pytest.param("gemini_api_key", "AQ.Ab8RN6I-7xK9pL2mQ0vW5yZ1uX3tS4rE8qD7cB6aN5mP4oI3uY2", id="gemini-auth-key"), + pytest.param("gemini_api_key", "AQ." + "a" * 45, id="gemini-auth-key-shortest"), + pytest.param("gemini_api_key", "AQ." + "a" * 55, id="gemini-auth-key-longest"), + pytest.param("gemini_api_key", "AQ." + "a" * 49 + "-", id="gemini-auth-key-trailing-hyphen"), + ], +) +def test_credential_pattern_matches_whole_key(pattern_name: str, key: str) -> None: + match: Final = get_compiled_pattern(pattern_name).search(f"I keep getting 401 with {key}, is it expired?") + assert match is not None + assert match.group(0) == key + + +@pytest.mark.parametrize("separator", ["\\n", "\\t", "\\u003c", "Bearer%20", "?key%3D", "?key%253D", "GOOGLE_KEY_"]) +@pytest.mark.parametrize( + ("pattern_name", "key"), + [ + pytest.param("openai_api_key", "sk-proj-" + "a" * 20 + "T3BlbkFJ" + "b" * 20, id="openai"), + pytest.param("anthropic_api_key", "sk-ant-api03-" + "a" * 20, id="anthropic"), + pytest.param("gemini_api_key", "AIza" + "a" * 35, id="gemini-legacy"), + pytest.param("gemini_api_key", "AQ." + "a" * 50, id="gemini-auth-key"), + ], +) +def test_credential_pattern_matches_key_after_escaped_separator(pattern_name: str, key: str, separator: str) -> None: + match: Final = get_compiled_pattern(pattern_name).search(f'{{"log": "retrying with{separator}{key} failed"}}') + assert match is not None + assert match.group(0) == key + + +@pytest.mark.parametrize( + ("pattern_name", "text"), + [ + pytest.param("openai_api_key", "sk-short", id="openai-too-short"), + pytest.param("openai_api_key", "sk-proj-" + "a" * 19, id="openai-project-too-short"), + pytest.param("openai_api_key", "sk-" + "a" * 20 + "T3BlbkFJ" + "b" * 21, id="openai-legacy-too-long"), + pytest.param("openai_api_key", "sk-proj-" + "a" * 48, id="openai-project-without-marker"), + pytest.param("openai_api_key", "sk-proj-" + "a" * 19 + "T3BlbkFJ" + "b" * 20, id="openai-project-marker-too-early"), + pytest.param("openai_api_key", "sk-proj-" + "a" * 20 + "T3BlbkFJ" + "b" * 19, id="openai-project-marker-too-late"), + pytest.param("openai_api_key", "SK-ADMIN-dashboard-redesign-phase-two", id="openai-prefixed-slug"), + pytest.param("openai_api_key", "desk-proj-management-dashboard-v2", id="openai-inside-longer-word"), + pytest.param( + "openai_api_key", "sk-abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRST", id="openai-legacy-without-marker" + ), + pytest.param("openai_api_key", "sk-0123456789abcdef0123456789abcdef", id="openai-other-vendor-sk-key"), + pytest.param("openai_api_key", "sk-Xq3vR9tLm2Pa8KdW5yZgQw", id="openai-litellm-virtual-key"), + pytest.param( + "openai_api_key", + "sk-AbCdEfGhIjKlMnOpQrStUvWxYz0123456789AbCd-_x", + id="openai-litellm-virtual-key-raised-length", + ), + pytest.param("openai_api_key", "not-a-key", id="openai-no-prefix"), + pytest.param( + "openai_api_key", "sk-ant-api03-thisisanthropicnotopenaiabc123", id="openai-ignores-anthropic-key" + ), + pytest.param("openai_api_key", "ordinary sentence without any api key in it", id="openai-plain-text"), + pytest.param("anthropic_api_key", "sk-ant-" + "a" * 19, id="anthropic-too-short"), + pytest.param("anthropic_api_key", "flask-ant-design-system-components", id="anthropic-inside-longer-word"), + pytest.param( + "anthropic_api_key", + "sk-proj-abcdefghijklmnopqrstuvwxyzABCDEF1234567890", + id="anthropic-ignores-openai-project-key", + ), + pytest.param( + "anthropic_api_key", + "sk-admin-abcdefghijklmnopqrstuvwxyzABCDEF123456", + id="anthropic-ignores-openai-admin-key", + ), + pytest.param("anthropic_api_key", "sk-ant", id="anthropic-prefix-only"), + pytest.param("anthropic_api_key", "plain text with nothing sensitive", id="anthropic-plain-text"), + pytest.param("gemini_api_key", "AIza", id="gemini-prefix-only"), + pytest.param("gemini_api_key", "AIza" + "a" * 32, id="gemini-legacy-too-short"), + pytest.param("gemini_api_key", "AIza" + "a" * 37, id="gemini-legacy-too-long"), + pytest.param("gemini_api_key", "BBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB", id="gemini-wrong-prefix"), + pytest.param("gemini_api_key", "AQ.", id="gemini-auth-prefix-only"), + pytest.param("gemini_api_key", "AQ." + "a" * 44, id="gemini-auth-key-too-short"), + pytest.param("gemini_api_key", "AQ." + "a" * 56, id="gemini-auth-key-too-long"), + pytest.param("gemini_api_key", "AQ-" + "a" * 50, id="gemini-auth-key-wrong-separator"), + pytest.param("gemini_api_key", "x" + "AIza" + "a" * 35, id="gemini-inside-longer-token"), + pytest.param("gemini_api_key", "The AQ. was reported at 42 ppm today.", id="gemini-aq-acronym"), + pytest.param("gemini_api_key", "Is the question AQ. Maybe it is BQ.", id="gemini-aq-sentence-end"), + pytest.param("gemini_api_key", "plain text with no google key", id="gemini-plain-text"), + pytest.param( + "gemini_api_key", + "https://example.com/blog/aizawa-takeshi-interview-2024-tokyo-01", + id="gemini-lowercase-legacy-prefix", + ), + pytest.param( + "gemini_api_key", + "ssh aq.internal-observability-metrics-collector-eu-west-01", + id="gemini-lowercase-auth-prefix", + ), + ], +) +def test_credential_pattern_ignores_non_keys(pattern_name: str, text: str) -> None: + assert get_compiled_pattern(pattern_name).search(text) is None + + +@pytest.mark.parametrize("suffix", ["请检查", "é", "ü der Schlüssel"]) +@pytest.mark.parametrize( + ("pattern_name", "key"), + [ + pytest.param("openai_api_key", "sk-" + "a" * 20 + "T3BlbkFJ" + "b" * 20, id="openai-legacy"), + pytest.param("gemini_api_key", "AIza" + "a" * 35, id="gemini-legacy"), + pytest.param("gemini_api_key", "AQ." + "a" * 50, id="gemini-auth-key"), + ], +) +def test_credential_pattern_matches_key_followed_by_non_ascii_letters(pattern_name: str, key: str, suffix: str) -> None: + match: Final = get_compiled_pattern(pattern_name).search(f"我的密钥是{key}{suffix}") + assert match is not None + assert match.group(0) == key + + +@pytest.mark.parametrize( + "leading", + [ + pytest.param("sk-proj-" + "a" * 74 + "T3BlbkFJ" + "b" * 74, id="another-key"), + pytest.param("sk-proj-" + "x" * 20 + "T3BlbkFJ" + "y" * 192, id="key-shaped-filler"), + pytest.param("sk-proj-" + "x" * 100 + "T3BlbkFJ" + "y" * 100, id="long-key-shaped-filler"), + ], +) +def test_openai_key_pattern_masks_a_key_glued_to_preceding_key_shaped_text(leading: str) -> None: + key: Final = "sk-proj-" + "c" * 74 + "T3BlbkFJ" + "d" * 74 + assert get_compiled_pattern("openai_api_key").sub("[KEY]", leading + key) == "[KEY][KEY]" + + +@pytest.mark.parametrize("prefix", ["sk-proj-", "sk-svcacct-", "sk-admin-", "sk-None-"]) +def test_openai_key_pattern_scans_repeated_prefixes_without_stalling(prefix: str) -> None: + prompt: Final = prefix * (128_000 // len(prefix)) + started: Final = time.perf_counter() + match: Final = get_compiled_pattern("openai_api_key").search(prompt) + elapsed: Final = time.perf_counter() - started + assert match is None + assert elapsed < 1.0 diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index 64e7cde4034..741b6cb74b5 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -31,6 +31,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.model_latest_aliases import get_model_latest_aliases from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_alias from litellm.litellm_core_utils.ptu_shares import PTUTeamCeiling +from litellm.litellm_core_utils.total_deadline import TotalTimeoutExpired, total_deadline from litellm.llms.azure.ptu_capacity import AZURE_PTU_CAPACITY from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError @@ -66,6 +67,7 @@ from litellm.redis.stores.rate_limits import ( RateLimitStore, ) from litellm.responses.streaming_iterator import _record_usage_for_failure +from tests._support.clock_loops import AutojumpClockLoop from tests.unit.redis.stores.fake_rate_limits import FakeRateLimits, scope_key from litellm.responses.utils import ResponseAPILoggingUtils from litellm.types.extraction import ExtractionResponse @@ -7951,6 +7953,34 @@ async def _admit(handler, auth, data=None): ) +def test_a_total_timeout_never_interrupts_a_token_reservation_so_the_failure_path_can_refund_it() -> None: + async def reserve_outlasting_the_budget(keys: list[str], args: list[object]) -> list[object]: + await asyncio.sleep(120) + return [0, *[value for index in range(0, len(args), 4) for value in (0, args[index + 1])]] + + async def admit_under_a_total_timeout() -> tuple[bool, int]: + redis: Final = FakeRateLimits(check_and_increment_by_n_script=reserve_outlasting_the_budget) + handler: Final = _handler_with_redis(redis) + auth: Final = UserAPIKeyAuth(api_key=hash_token("sk-total-timeout-reservation"), tpm_limit=1000) + stash: Final = get_or_create_request_stash() + try: + async with total_deadline(60, from_body=True): + await _admit(handler, auth) + await asyncio.sleep(3600) + except TotalTimeoutExpired: + return True, stash.reserved_tokens + return False, stash.reserved_tokens + + loop: Final = AutojumpClockLoop() + try: + expired, reserved = loop.run_until_complete(admit_under_a_total_timeout()) + finally: + loop.close() + + assert expired + assert reserved > 0 + + async def _read_only_check(handler, auth): descriptors = handler.create_rate_limit_descriptors( user_api_key_dict=auth, diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index d06f95111aa..0283c90f3d6 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -44,6 +44,7 @@ from litellm.proxy.utils import ( BilledFailedAttempt, ProxyUpdateSpend, ) +from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys from litellm.types.router import Deployment, LiteLLM_Params from litellm.types.utils import CallTypes, LiteLLMBatch, ModelResponse, PromptTokensDetailsWrapper, Usage from litellm.types.videos.main import VideoObject @@ -1129,6 +1130,95 @@ def test_get_budget_reservation_from_metadata_handles_dict_auth_object(): ) +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("metadata", "expected_scope"), + [ + ( + { + "user_api_key_auth": UserAPIKeyAuth( + committed_budget_spend=CommittedBudgetSpend( + counter_keys=EntitySpendCounterKeys(key="spend:key:test_api_key") + ) + ) + }, + EntitySpendCounterKeys(key="spend:key:test_api_key"), + ), + ({}, EntitySpendCounterKeys()), + ], +) +async def test_charge_moves_only_the_committed_counters_auth_named_as_budgeted(metadata, expected_scope): + proxy_logging_obj = MagicMock() + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock() + increment_spend_counters = AsyncMock() + + await _update_database_and_spend_counters( + proxy_logging_obj=proxy_logging_obj, + increment_spend_counters=increment_spend_counters, + user_api_key="test_api_key", + user_id="test_user_id", + end_user_id=None, + team_id="test_team_id", + org_id="test_org_id", + kwargs={"litellm_params": {"metadata": metadata}}, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + response_cost=0.2, + budget_reservation=None, + ) + + assert increment_spend_counters.await_args.kwargs["committed_spend_scope"] == expected_scope + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("tracked", "expected_key", "expected_org"), + [(True, pytest.approx(1.2), pytest.approx(2.2)), (False, 1.0, 2.0)], +) +async def test_charge_with_no_auth_on_its_row_moves_the_committed_counters_that_exist( + tracked, expected_key, expected_org +): + """A background response is billed from its creator's stored attribution, after the request's auth is gone.""" + from litellm.proxy.proxy_server import increment_spend_counters + + proxy_logging_obj = MagicMock() + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock() + spend_counter_store: Final = SpendCounterStore(connect(None).spend_counters()) + spend_counter_store.set_local("spend:committed:key:test_api_key", 1.0) + spend_counter_store.set_local("spend:committed:org:test_org_id", 2.0) + + with ( + patch("litellm.proxy.hooks.proxy_track_cost_callback.committed_spend_tracked", return_value=tracked), + patch("litellm.proxy.proxy_server.spend_counter_store", spend_counter_store), + patch( + "litellm.proxy.spend_tracking.committed_spend.SpendCounterReseed.from_db", + AsyncMock(return_value=50.0), + ), + ): + await _update_database_and_spend_counters( + proxy_logging_obj=proxy_logging_obj, + increment_spend_counters=increment_spend_counters, + user_api_key="test_api_key", + user_id="test_user_id", + end_user_id=None, + team_id="test_team_id", + org_id="test_org_id", + kwargs={"litellm_params": {"metadata": {"user_api_key_hash": "test_api_key"}}}, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + response_cost=0.2, + budget_reservation=None, + ) + + cached: Final = spend_counter_store.get_local + assert cached("spend:committed:key:test_api_key") == expected_key + assert cached("spend:committed:org:test_org_id") == expected_org + assert cached("spend:committed:team:test_team_id") is None + assert cached("spend:committed:user:test_user_id") is None + + @pytest.mark.asyncio async def test_update_database_and_spend_counters_releases_reservation_when_db_update_fails(): proxy_logging_obj = MagicMock() @@ -1257,6 +1347,7 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda request_started_at=start_time, model_access_groups=("premium",), project_id=None, + committed_spend_scope=EntitySpendCounterKeys(), ) @@ -2776,7 +2867,9 @@ _LIVE_COUNTERS: Final = { } -async def _fail_request(logger: ProxyDBLogger, failure: Mapping[str, object]) -> None: +async def _fail_request( + logger: ProxyDBLogger, failure: Mapping[str, object], committed: CommittedBudgetSpend | None = None +) -> None: await logger.async_post_call_failure_hook( request_data={ "model": "gpt-4.1", @@ -2793,6 +2886,7 @@ async def _fail_request(logger: ProxyDBLogger, failure: Mapping[str, object]) -> end_user_id="end-user-1", org_id="org-1", project_id="project-1", + committed_budget_spend=committed, ), ) @@ -2831,6 +2925,43 @@ async def test_failure_spend_rows_raise_the_live_spend_counters_by_their_cost( ) +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("failure", "charged"), + [ + ({"combined_usage_object": Usage(prompt_tokens=30), "response_cost": 0.125}, 0.125), + (_failure_with_billed_attempts(), 1.25), + ({"standard_logging_object": {"cost_finalized": True, "response_cost": 0.125}}, 0.125), + ], + ids=["recovered_stream_cost", "billed_retries", "completed_cost"], +) +async def test_a_failed_charge_moves_only_the_committed_counters_auth_named_as_budgeted( + live_spend_counters: SpendCounterStore, failure: Mapping[str, object], charged: float +): + """The remaining budget gauges read the committed counters, so a charge missing from them shows as budget the + request never left behind, until the next reseed from the database.""" + budgeted: Final = EntitySpendCounterKeys.of( + token="hashed-key", team_id="team-1", user_id=None, org_id=None, end_user_id=None + ) + starting: Final = { + "spend:committed:key:hashed-key": 0.2, + "spend:committed:team:team-1": 0.3, + "spend:committed:user:user-1": 0.4, + } + for key, spend in starting.items(): + live_spend_counters.set_local(key, spend) + + await _fail_request(ProxyDBLogger(), failure, CommittedBudgetSpend(counter_keys=budgeted)) + + assert {key: live_spend_counters.get_local(key) for key in starting} == pytest.approx( + { + "spend:committed:key:hashed-key": 0.2 + charged, + "spend:committed:team:team-1": 0.3 + charged, + "spend:committed:user:user-1": 0.4, + } + ) + + @pytest.mark.asyncio async def test_a_completed_cost_failure_raises_the_live_spend_counters_once(live_spend_counters: SpendCounterStore): """A failure whose cost was already finalized settles through the success path's write, which raises the diff --git a/tests/unit/proxy/management_endpoints/test_customer_endpoints.py b/tests/unit/proxy/management_endpoints/test_customer_endpoints.py index e9f10e8c0db..c42690050da 100644 --- a/tests/unit/proxy/management_endpoints/test_customer_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_customer_endpoints.py @@ -249,6 +249,29 @@ def test_delete_customer_not_found(mock_prisma_client, mock_user_api_key_auth): assert response_json["error"]["code"] == "404" +def test_delete_customer_drops_its_committed_spend_so_a_recreated_customer_starts_from_its_own_row( + monkeypatch, mock_prisma_client, mock_user_api_key_auth +): + from types import SimpleNamespace + + from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters + + spend: Final = FakeSpendCounters( + redis=None, + memory={"spend:committed:end_user:cust-gone": 7.0, "spend:committed:end_user:cust-stays": 3.0}, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_store", spend.store()) + mock_prisma_client.db.litellm_endusertable.find_many = AsyncMock(return_value=[SimpleNamespace(user_id="cust-gone")]) + mock_prisma_client.db.litellm_endusertable.delete_many = AsyncMock(return_value=1) + + response: Final = client.post( + "/customer/delete", json={"user_ids": ["cust-gone"]}, headers={"Authorization": "Bearer test-key"} + ) + + assert response.status_code == 200 + assert spend.memory == {"spend:committed:end_user:cust-stays": 3.0} + + def test_error_schema_consistency(mock_prisma_client, mock_user_api_key_auth): """ Test that all customer endpoints return the same error schema format. diff --git a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py index 120cf8a2b84..b7be80da31a 100644 --- a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py @@ -4222,7 +4222,10 @@ async def test_admin_user_update_spend_invalidates_counter(mocker): admin_caller = UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) await _update_single_user_helper(user_request=user_request, user_api_key_dict=admin_caller) - mock_invalidate.assert_awaited_once_with(counter_key="spend:user:target-user") + assert [call.kwargs["counter_key"] for call in mock_invalidate.await_args_list] == [ + "spend:user:target-user", + "spend:committed:user:target-user", + ] @pytest.mark.asyncio @@ -5223,6 +5226,39 @@ async def test_delete_user_evicts_cached_user_rows(mocker: MockerFixture) -> Non broadcast.assert_awaited_once_with(cache_key=deleted.user_id) +@pytest.mark.asyncio +async def test_delete_user_drops_the_committed_spend_of_the_users_it_deleted(mocker: MockerFixture) -> None: + from litellm.proxy._types import DeleteUserRequest, LiteLLM_UserTable + from litellm.proxy.management_endpoints.internal_user_endpoints import delete_user + from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters + + spend: Final = FakeSpendCounters( + redis=None, memory={"spend:committed:user:user-gone": 7.0, "spend:committed:user:user-stays": 3.0} + ) + mocker.patch("litellm.proxy.proxy_server.spend_counter_store", spend.store()) # test-quality-ok: substitute the spend counter dependency + deleted: Final = LiteLLM_UserTable(user_id="user-gone", user_email="gone@example.test", teams=[]) + prisma_client: Final = mocker.MagicMock() + prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(return_value=deleted) + prisma_client.db.litellm_teamtable.find_many = mocker.AsyncMock(return_value=[]) + prisma_client.db.litellm_jwtkeymapping.find_many = mocker.AsyncMock(return_value=[]) + prisma_client.db.litellm_verificationtoken.find_many = mocker.AsyncMock(return_value=[]) + prisma_client.db.litellm_verificationtoken.delete_many = mocker.AsyncMock(return_value=0) + prisma_client.db.litellm_invitationlink.delete_many = mocker.AsyncMock(return_value=0) + prisma_client.db.litellm_organizationmembership.delete_many = mocker.AsyncMock(return_value=0) + prisma_client.db.litellm_teammembership.delete_many = mocker.AsyncMock(return_value=0) + prisma_client.db.litellm_usertable.delete_many = mocker.AsyncMock(return_value=1) + _serve_locked_user_delete(prisma_client, deleted) + mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency + mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj", None) # test-quality-ok: delete_user reads it off proxy_server at call time + + await delete_user( + data=DeleteUserRequest(user_ids=[deleted.user_id]), + user_api_key_dict=UserAPIKeyAuth(user_id="proxy-admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert spend.memory == {"spend:committed:user:user-stays": 3.0} + + _DB_OUTAGE_503_BODY: Final = { "error": { "message": "Service Unavailable, the authentication database is temporarily unreachable. Please retry shortly.", diff --git a/tests/unit/proxy/management_endpoints/test_key_generate_prisma.py b/tests/unit/proxy/management_endpoints/test_key_generate_prisma.py index 26f1603ccca..3a05e3e638b 100644 --- a/tests/unit/proxy/management_endpoints/test_key_generate_prisma.py +++ b/tests/unit/proxy/management_endpoints/test_key_generate_prisma.py @@ -24,7 +24,6 @@ import re import traceback from litellm._uuid import uuid from datetime import datetime, timezone -from unittest import mock from dotenv import load_dotenv from fastapi import Request @@ -53,12 +52,10 @@ from litellm.proxy.auth.auth_checks import get_key_object from litellm.proxy.management_endpoints.key_management_endpoints import ( delete_key_fn, generate_key_fn, - generate_key_helper_fn, info_key_fn, list_keys, regenerate_key_fn, update_key_fn, - key_aliases, ) from litellm.proxy.management_endpoints.team_endpoints import ( new_team, @@ -76,13 +73,7 @@ from litellm.proxy.proxy_server import ( user_api_key_auth, ) from litellm.proxy.image_endpoints import image_generation -from litellm.proxy.management_endpoints.customer_endpoints import ( - new_end_user, -) from litellm.proxy.spend_tracking.spend_management_endpoints import ( - global_spend, - spend_key_fn, - spend_user_fn, view_spend_logs, ) from litellm.proxy.utils import PrismaClient, ProxyLogging, hash_token, update_spend @@ -98,7 +89,6 @@ from litellm.proxy._types import ( DynamoDBArgs, GenerateKeyRequest, KeyRequest, - NewCustomerRequest, NewTeamRequest, NewUserRequest, ProxyErrorTypes, @@ -297,121 +287,6 @@ def test_generate_and_call_with_valid_key(prisma_client, api_route): pytest.fail(f"An exception occurred - {str(e)}") -@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") -def test_call_with_invalid_key(prisma_client): - # 2. Make a call with invalid key, expect it to fail - setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) - setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY) - - async def test(): - await litellm.proxy.proxy_server.prisma_client.connect() - generated_key = "sk-126666" - bearer_token = "Bearer " + generated_key - - request = Request(scope={"type": "http"}, receive=None) - request._url = URL(url="/chat/completions") - - # use generated key to auth in - result = await user_api_key_auth(request=request, api_key=bearer_token) - print("got result", result) - pytest.fail(f"This should have failed!. IT's an invalid key") - - with pytest.raises(Exception, match="Authentication Error, Invalid API key provided") as exc_info: - asyncio.run(test()) - e = exc_info.value - print("Got Exception", e) - print(e.message) - assert "Authentication Error, Invalid API key provided" in e.message - - -@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") -def test_call_with_invalid_model(prisma_client): - litellm.set_verbose = True - # 3. Make a call to a key with an invalid model - expect to fail - setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) - setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY) - - async def test(): - await litellm.proxy.proxy_server.prisma_client.connect() - request = NewUserRequest(models=["mistral"]) - key = await new_user( - data=request, - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key=MASTER_KEY, - user_id="1234", - ), - ) - print(key) - - generated_key = key.key - bearer_token = "Bearer " + generated_key - - request = Request(scope={"type": "http"}) - request._url = URL(url="/chat/completions") - - async def return_body(): - return b'{"model": "gemini-pro-vision"}' - - request.body = return_body - - # use generated key to auth in - print( - "Bearer token being sent to user_api_key_auth() - {}".format( - bearer_token - ) - ) - result = await user_api_key_auth(request=request, api_key=bearer_token) - pytest.fail(f"This should have failed!. IT's an invalid model") - - with pytest.raises(ProxyException) as exc_info: - asyncio.run(test()) - e = exc_info.value - assert isinstance(e, ProxyException) - assert e.type == ProxyErrorTypes.key_model_access_denied - assert e.param == "model" - - -@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") -def test_call_with_valid_model(prisma_client): - # 4. Make a call to a key with a valid model - expect to pass - setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) - setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY) - try: - - async def test(): - await litellm.proxy.proxy_server.prisma_client.connect() - request = NewUserRequest(models=["mistral"]) - key = await new_user( - request, - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key=MASTER_KEY, - user_id="1234", - ), - ) - print(key) - - generated_key = key.key - bearer_token = "Bearer " + generated_key - - request = Request(scope={"type": "http"}) - request._url = URL(url="/chat/completions") - - async def return_body(): - return b'{"model": "mistral"}' - - request.body = return_body - - # use generated key to auth in - result = await user_api_key_auth(request=request, api_key=bearer_token) - print("result from user auth with new key", result) - - asyncio.run(test()) - except Exception as e: - pytest.fail(f"An exception occurred - {str(e)}") - - @pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") @pytest.mark.asyncio async def test_call_with_valid_model_using_all_models(prisma_client): @@ -482,89 +357,6 @@ async def test_call_with_valid_model_using_all_models(prisma_client): pytest.fail(f"An exception occurred - {str(e)}") -@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") -def test_call_with_user_over_budget(prisma_client): - # 5. Make a call with a key over budget, expect to fail - setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) - setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY) - - async def test(): - await litellm.proxy.proxy_server.prisma_client.connect() - request = NewUserRequest(max_budget=0.00001) - key = await new_user( - data=request, - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key=MASTER_KEY, - user_id="1234", - ), - ) - print(key) - - generated_key = key.key - user_id = key.user_id - bearer_token = "Bearer " + generated_key - - request = Request(scope={"type": "http"}) - request._url = URL(url="/chat/completions") - - # use generated key to auth in - result = await user_api_key_auth(request=request, api_key=bearer_token) - print("result from user auth with new key", result) - - # update spend using track_cost callback, make 2nd request, it should fail - from litellm import Choices, Message, ModelResponse, Usage - from litellm.proxy.proxy_server import ProxyDBLogger - - proxy_db_logger = ProxyDBLogger() - - resp = ModelResponse( - id="chatcmpl-e41836bb-bb8b-4df2-8e70-8f3e160155ac", - choices=[ - Choices( - finish_reason=None, - index=0, - message=Message( - content=" Sure! Here is a short poem about the sky:\n\nA canvas of blue, a", - role="assistant", - ), - ) - ], - model="gpt-35-turbo", # azure always has model written like this - usage=Usage(prompt_tokens=210, completion_tokens=200, total_tokens=410), - ) - await proxy_db_logger._PROXY_track_cost_callback( - kwargs={ - "stream": False, - "litellm_params": { - "metadata": { - "user_api_key": generated_key, - "user_api_key_user_id": user_id, - } - }, - "response_cost": 0.00002, - }, - completion_response=resp, - start_time=datetime.now(), - end_time=datetime.now(), - ) - await asyncio.sleep(5) - # use generated key to auth in - result = await user_api_key_auth(request=request, api_key=bearer_token) - print("result from user auth with new key", result) - pytest.fail("This should have failed!. They key crossed it's budget") - - with pytest.raises(ProxyException) as exc_info: - asyncio.run(test()) - e = exc_info.value - print("got an errror=", e) - error_detail = e.message - assert "ExceededBudget:" in error_detail - assert isinstance(e, ProxyException) - assert e.type == ProxyErrorTypes.budget_exceeded - print(vars(e)) - - def test_end_user_cache_write_unit_test(): """ assert end user object is being written to cache as expected @@ -572,210 +364,6 @@ def test_end_user_cache_write_unit_test(): pass -@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") -def test_call_with_end_user_over_budget(prisma_client): - # Test if a user passed to /chat/completions is tracked & fails when they cross their budget - # we only check this when litellm.max_end_user_budget is set - import random - - setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) - setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY) - setattr(litellm, "max_end_user_budget", 0.00001) - - async def test(): - await litellm.proxy.proxy_server.prisma_client.connect() - user = f"ishaan {uuid.uuid4().hex}" - request = NewCustomerRequest( - user_id=user, max_budget=0.000001 - ) # create a key with no budget - await new_end_user( - request, - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key=MASTER_KEY, - user_id="1234", - ), - ) - - request = Request(scope={"type": "http"}) - request._url = URL(url="/chat/completions") - bearer_token = "Bearer sk-9876" - - async def return_body(): - return_string = f'{{"model": "gemini-pro-vision", "user": "{user}"}}' - # return string as bytes - return return_string.encode() - - request.body = return_body - - result = await user_api_key_auth(request=request, api_key=bearer_token) - - # update spend using track_cost callback, make 2nd request, it should fail - from litellm import Choices, Message, ModelResponse, Usage - from litellm.proxy.proxy_server import ProxyDBLogger - - proxy_db_logger = ProxyDBLogger() - - resp = ModelResponse( - id="chatcmpl-e41836bb-bb8b-4df2-8e70-8f3e160155ac", - choices=[ - Choices( - finish_reason=None, - index=0, - message=Message( - content=" Sure! Here is a short poem about the sky:\n\nA canvas of blue, a", - role="assistant", - ), - ) - ], - model="gpt-35-turbo", # azure always has model written like this - usage=Usage(prompt_tokens=210, completion_tokens=200, total_tokens=410), - ) - await proxy_db_logger._PROXY_track_cost_callback( - kwargs={ - "stream": False, - "litellm_params": { - "metadata": { - "user_api_key": MASTER_KEY, - "user_api_key_end_user_id": user, - }, - "proxy_server_request": { - "body": { - "user": user, - } - }, - }, - "response_cost": 10, - }, - completion_response=resp, - start_time=datetime.now(), - end_time=datetime.now(), - ) - - await asyncio.sleep(10) - await update_spend( - prisma_client=prisma_client, - db_writer_client=None, - proxy_logging_obj=proxy_logging_obj, - ) - - # use generated key to auth in - result = await user_api_key_auth(request=request, api_key=bearer_token) - print("result from user auth with new key", result) - pytest.fail("This should have failed!. They key crossed it's budget") - - with pytest.raises(ProxyException) as exc_info: - asyncio.run(test()) - e = exc_info.value - print(f"raised error: {e}, traceback: {traceback.format_exc()}") - # Handle DataError and other exceptions that don't have .message attribute - error_detail = getattr(e, "message", str(e)) - assert "ExceededBudget: End User=" in error_detail - assert "over budget" in error_detail - assert isinstance(e, ProxyException) - assert e.type == ProxyErrorTypes.budget_exceeded - print(vars(e)) - - -@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") -def test_call_with_proxy_over_budget(prisma_client): - # 5.1 Make a call with a proxy over budget, expect to fail - setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) - setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY) - litellm_proxy_budget_name = f"litellm-proxy-budget-{time.time()}" - setattr( - litellm.proxy.proxy_server, - "litellm_proxy_admin_name", - litellm_proxy_budget_name, - ) - setattr(litellm, "max_budget", 0.00001) - from litellm.proxy.proxy_server import user_api_key_cache - - user_api_key_cache.set_cache( - key="{}:spend".format(litellm_proxy_budget_name), value=0 - ) - setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache) - - async def test(): - await litellm.proxy.proxy_server.prisma_client.connect() - request = NewUserRequest() - key = await new_user( - data=request, - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key=MASTER_KEY, - user_id="1234", - ), - ) - print(key) - - generated_key = key.key - user_id = key.user_id - bearer_token = "Bearer " + generated_key - - request = Request(scope={"type": "http"}) - request._url = URL(url="/chat/completions") - - # use generated key to auth in - result = await user_api_key_auth(request=request, api_key=bearer_token) - print("result from user auth with new key", result) - - # update spend using track_cost callback, make 2nd request, it should fail - from litellm import Choices, Message, ModelResponse, Usage - from litellm.proxy.proxy_server import ProxyDBLogger - - proxy_db_logger = ProxyDBLogger() - - resp = ModelResponse( - id="chatcmpl-e41836bb-bb8b-4df2-8e70-8f3e160155ac", - choices=[ - Choices( - finish_reason=None, - index=0, - message=Message( - content=" Sure! Here is a short poem about the sky:\n\nA canvas of blue, a", - role="assistant", - ), - ) - ], - model="gpt-35-turbo", # azure always has model written like this - usage=Usage(prompt_tokens=210, completion_tokens=200, total_tokens=410), - ) - await proxy_db_logger._PROXY_track_cost_callback( - kwargs={ - "stream": False, - "litellm_params": { - "metadata": { - "user_api_key": generated_key, - "user_api_key_user_id": user_id, - } - }, - "response_cost": 0.00002, - }, - completion_response=resp, - start_time=datetime.now(), - end_time=datetime.now(), - ) - - await asyncio.sleep(5) - # use generated key to auth in - result = await user_api_key_auth(request=request, api_key=bearer_token) - print("result from user auth with new key", result) - pytest.fail(f"This should have failed!. They key crossed it's budget") - - with pytest.raises(ProxyException) as exc_info: - asyncio.run(test()) - e = exc_info.value - if hasattr(e, "message"): - error_detail = e.message - else: - error_detail = traceback.format_exc() - assert "Budget has been exceeded" in error_detail - assert isinstance(e, ProxyException) - assert e.type == ProxyErrorTypes.budget_exceeded - print(vars(e)) - - @pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") def test_call_with_user_over_budget_stream(prisma_client): # 6. Make a call with a key over budget, expect to fail @@ -970,287 +558,6 @@ def test_call_with_proxy_over_budget_stream(prisma_client): print(vars(e)) -@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") -def test_generate_and_call_with_valid_key_never_expires(prisma_client): - # 7. Make a call with an key that never expires, expect to pass - - print("prisma client=", prisma_client) - - setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) - setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY) - try: - - async def test(): - await litellm.proxy.proxy_server.prisma_client.connect() - request = NewUserRequest(duration=None) - key = await new_user( - data=request, - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key=MASTER_KEY, - user_id="1234", - ), - ) - print(key) - - generated_key = key.key - bearer_token = "Bearer " + generated_key - - request = Request(scope={"type": "http"}) - request._url = URL(url="/chat/completions") - - # use generated key to auth in - result = await user_api_key_auth(request=request, api_key=bearer_token) - print("result from user auth with new key", result) - - asyncio.run(test()) - except Exception as e: - pytest.fail(f"An exception occurred - {str(e)}") - - -@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") -def test_generate_and_call_with_expired_key(prisma_client): - # 8. Make a call with an expired key, expect to fail - - print("prisma client=", prisma_client) - - setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) - setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY) - - async def test(): - await litellm.proxy.proxy_server.prisma_client.connect() - request = NewUserRequest(duration="0s") - key = await new_user( - data=request, - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key=MASTER_KEY, - user_id="1234", - ), - ) - print(key) - - generated_key = key.key - bearer_token = "Bearer " + generated_key - - request = Request(scope={"type": "http"}) - request._url = URL(url="/chat/completions") - - # use generated key to auth in - result = await user_api_key_auth(request=request, api_key=bearer_token) - print("result from user auth with new key", result) - pytest.fail("This should have failed!. It's an expired key") - - with pytest.raises(Exception, match="Authentication Error") as exc_info: - asyncio.run(test()) - e = exc_info.value - print("Got Exception", e) - print(e.message) - assert "Authentication Error" in e.message - assert e.type == ProxyErrorTypes.expired_key - - -@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") -def test_delete_key(prisma_client): - # 9. Generate a Key, delete it. Check if deletion works fine - - print("prisma client=", prisma_client) - - setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) - setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY) - setattr(litellm.proxy.proxy_server, "user_custom_auth", None) - try: - - async def test(): - await litellm.proxy.proxy_server.prisma_client.connect() - from litellm.proxy.proxy_server import user_api_key_cache - - request = NewUserRequest() - key = await new_user( - data=request, - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key=MASTER_KEY, - user_id="1234", - ), - ) - print(key) - - generated_key = key.key - bearer_token = "Bearer " + generated_key - - delete_key_request = KeyRequest(keys=[generated_key]) - - bearer_token = "Bearer sk-9876" - - request = Request(scope={"type": "http"}) - request._url = URL(url="/key/delete") - - # use generated key to auth in - result = await user_api_key_auth(request=request, api_key=bearer_token) - print(f"result: {result}") - result.user_role = LitellmUserRoles.PROXY_ADMIN - # delete the key - result_delete_key = await delete_key_fn( - data=delete_key_request, user_api_key_dict=result - ) - print("result from delete key", result_delete_key) - assert result_delete_key == {"deleted_keys": [generated_key]} - - assert generated_key not in user_api_key_cache.in_memory_cache.cache_dict - assert ( - hash_token(generated_key) - not in user_api_key_cache.in_memory_cache.cache_dict - ) - - asyncio.run(test()) - except Exception as e: - pytest.fail(f"An exception occurred - {str(e)}") - - -@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") -def test_delete_key_auth(prisma_client): - # 10. Generate a Key, delete it, use it to make a call -> expect fail - - print("prisma client=", prisma_client) - - setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) - setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY) - try: - - async def test(): - await litellm.proxy.proxy_server.prisma_client.connect() - from litellm.proxy.proxy_server import user_api_key_cache - - request = NewUserRequest() - key = await new_user( - data=request, - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key=MASTER_KEY, - user_id="1234", - ), - ) - print(key) - - generated_key = key.key - bearer_token = "Bearer " + generated_key - - delete_key_request = KeyRequest(keys=[generated_key]) - - # delete the key - bearer_token = "Bearer sk-9876" - - request = Request(scope={"type": "http"}) - request._url = URL(url="/key/delete") - - # use generated key to auth in - result = await user_api_key_auth(request=request, api_key=bearer_token) - print(f"result: {result}") - result.user_role = LitellmUserRoles.PROXY_ADMIN - - result_delete_key = await delete_key_fn( - data=delete_key_request, user_api_key_dict=result - ) - - print("result from delete key", result_delete_key) - assert result_delete_key == {"deleted_keys": [generated_key]} - - request = Request(scope={"type": "http"}, receive=None) - request._url = URL(url="/chat/completions") - - assert generated_key not in user_api_key_cache.in_memory_cache.cache_dict - assert ( - hash_token(generated_key) - not in user_api_key_cache.in_memory_cache.cache_dict - ) - - # use generated key to auth in - bearer_token = "Bearer " + generated_key - result = await user_api_key_auth(request=request, api_key=bearer_token) - print("got result", result) - pytest.fail(f"This should have failed!. IT's an invalid key") - - asyncio.run(test()) - except Exception as e: - print("Got Exception", e) - # Handle different exception types - ProxyException has .message, others might have .detail or str(e) - error_message = ( - getattr(e, "message", None) or getattr(e, "detail", None) or str(e) - ) - print(f"Error message: {error_message}") - assert ( - "Authentication Error" in error_message - or "Invalid proxy server token" in error_message - or "not found in db" in error_message - ) - pass - - -@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") -def test_generate_and_call_key_info(prisma_client): - # 10. Generate a Key, cal key/info - - print("prisma client=", prisma_client) - - setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) - setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY) - try: - - async def test(): - await litellm.proxy.proxy_server.prisma_client.connect() - request = NewUserRequest( - metadata={"team": "litellm-team3", "project": "litellm-project3"} - ) - key = await new_user( - data=request, - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key=MASTER_KEY, - user_id="1234", - ), - ) - print(key) - - generated_key = key.key - - # use generated key to auth in - result = await info_key_fn( - key=generated_key, - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - ), - ) - print("result from info_key_fn", result) - assert result["key"] == generated_key - print("\n info for key=", result["info"]) - assert result["info"]["max_parallel_requests"] == None - assert result["info"]["metadata"] == { - "team": "litellm-team3", - "project": "litellm-project3", - } - - # cleanup - delete key - delete_key_request = KeyRequest(keys=[generated_key]) - bearer_token = "Bearer sk-9876" - - request = Request(scope={"type": "http"}) - request._url = URL(url="/key/delete") - - # use generated key to auth in - result = await user_api_key_auth(request=request, api_key=bearer_token) - print(f"result: {result}") - result.user_role = LitellmUserRoles.PROXY_ADMIN - - result_delete_key = await delete_key_fn( - data=delete_key_request, user_api_key_dict=result - ) - - asyncio.run(test()) - except Exception as e: - pytest.fail(f"An exception occurred - {str(e)}") - - @pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") def test_generate_and_update_key(prisma_client): # 11. Generate a Key, cal key/info, call key/update, call key/info @@ -1540,123 +847,6 @@ def test_key_generate_with_custom_auth(prisma_client): pytest.fail(f"An exception occurred - {str(e)}") -@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") -def test_call_with_key_over_budget(prisma_client): - # 12. Make a call with a key over budget, expect to fail - setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) - setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY) - - async def test(): - await litellm.proxy.proxy_server.prisma_client.connect() - request = GenerateKeyRequest(max_budget=0.00001) - key = await generate_key_fn( - request, - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key=MASTER_KEY, - user_id="1234", - ), - ) - print(key) - - generated_key = key.key - user_id = key.user_id - bearer_token = "Bearer " + generated_key - - request = Request(scope={"type": "http"}) - request._url = URL(url="/chat/completions") - - # use generated key to auth in - result = await user_api_key_auth(request=request, api_key=bearer_token) - print("result from user auth with new key", result) - - # update spend using track_cost callback, make 2nd request, it should fail - from litellm import Choices, Message, ModelResponse, Usage - from litellm.caching.caching import Cache - from litellm.proxy.proxy_server import ProxyDBLogger - - proxy_db_logger = ProxyDBLogger() - - litellm.cache = Cache() - import time - from litellm._uuid import uuid - - request_id = f"chatcmpl-e41836bb-bb8b-4df2-8e70-8f3e160155ac{uuid.uuid4()}" - - resp = ModelResponse( - id=request_id, - choices=[ - Choices( - finish_reason=None, - index=0, - message=Message( - content=" Sure! Here is a short poem about the sky:\n\nA canvas of blue, a", - role="assistant", - ), - ) - ], - model="gpt-35-turbo", # azure always has model written like this - usage=Usage(prompt_tokens=210, completion_tokens=200, total_tokens=410), - ) - await proxy_db_logger._PROXY_track_cost_callback( - kwargs={ - "model": "chatgpt-v-3", - "stream": False, - "litellm_params": { - "metadata": { - "user_api_key": hash_token(generated_key), - "user_api_key_user_id": user_id, - } - }, - "response_cost": 0.00002, - }, - completion_response=resp, - start_time=datetime.now(), - end_time=datetime.now(), - ) - await update_spend( - prisma_client=prisma_client, - db_writer_client=None, - proxy_logging_obj=proxy_logging_obj, - ) - # test spend_log was written and we can read it - spend_logs = await view_spend_logs( - request_id=request_id, - user_api_key_dict=UserAPIKeyAuth(api_key=generated_key), - ) - - print("read spend logs", spend_logs) - assert len(spend_logs) == 1 - - spend_log = spend_logs[0] - - assert spend_log.request_id == request_id - assert spend_log.spend == float("2e-05") - assert spend_log.model == "chatgpt-v-3" - assert ( - spend_log.cache_key - == "509ba0554a7129ae4f4fd13d11c141acce5549bb6aaf1f629ed543101615658e" - ) - - # use generated key to auth in - result = await user_api_key_auth(request=request, api_key=bearer_token) - print("result from user auth with new key", result) - pytest.fail("This should have failed!. They key crossed it's budget") - - with pytest.raises(ProxyException) as exc_info: - asyncio.run(test()) - e = exc_info.value - traceback.print_exc() - if hasattr(e, "message"): - error_detail = e.message - else: - error_detail = str(e) - assert "Budget has been exceeded" in error_detail - assert isinstance(e, ProxyException) - assert e.type == ProxyErrorTypes.budget_exceeded - print(vars(e)) - - @pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") def test_call_with_key_over_budget_no_cache(prisma_client): # 12. Make a call with a key over budget, expect to fail @@ -1908,246 +1098,6 @@ async def test_aasync_call_with_key_over_model_budget( assert exc_info.value.type == ProxyErrorTypes.budget_exceeded -@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") -@pytest.mark.asyncio() -async def test_call_with_key_never_over_budget(prisma_client): - # Make a call with a key with budget=None, it should never fail - setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) - setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY) - try: - await litellm.proxy.proxy_server.prisma_client.connect() - request = GenerateKeyRequest(max_budget=None) - key = await generate_key_fn( - request, - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key=MASTER_KEY, - user_id="1234", - ), - ) - print(key) - - generated_key = key.key - user_id = key.user_id - bearer_token = "Bearer " + generated_key - - request = Request(scope={"type": "http"}) - request._url = URL(url="/chat/completions") - - # use generated key to auth in - result = await user_api_key_auth(request=request, api_key=bearer_token) - print("result from user auth with new key: {result}") - - # update spend using track_cost callback, make 2nd request, it should fail - import time - from litellm._uuid import uuid - - from litellm import Choices, Message, ModelResponse, Usage - from litellm.proxy.proxy_server import ProxyDBLogger - - proxy_db_logger = ProxyDBLogger() - - request_id = f"chatcmpl-{uuid.uuid4()}" - - resp = ModelResponse( - id=request_id, - choices=[ - Choices( - finish_reason=None, - index=0, - message=Message( - content=" Sure! Here is a short poem about the sky:\n\nA canvas of blue, a", - role="assistant", - ), - ) - ], - model="gpt-35-turbo", # azure always has model written like this - usage=Usage( - prompt_tokens=210000, completion_tokens=200000, total_tokens=41000 - ), - ) - await proxy_db_logger._PROXY_track_cost_callback( - kwargs={ - "model": "chatgpt-v-3", - "stream": False, - "litellm_params": { - "metadata": { - "user_api_key": hash_token(generated_key), - "user_api_key_user_id": user_id, - } - }, - "response_cost": 200000, - }, - completion_response=resp, - start_time=datetime.now(), - end_time=datetime.now(), - ) - await update_spend( - prisma_client=prisma_client, - db_writer_client=None, - proxy_logging_obj=proxy_logging_obj, - ) - # use generated key to auth in - result = await user_api_key_auth(request=request, api_key=bearer_token) - print("result from user auth with new key", result) - except Exception as e: - pytest.fail(f"This should have not failed!. They key uses max_budget=None. {e}") - - -@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") -@pytest.mark.asyncio -async def test_call_with_key_over_budget_stream(prisma_client): - # 14. Make a call with a key over budget, expect to fail - setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) - setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY) - import logging - - from litellm._logging import verbose_proxy_logger - - litellm.set_verbose = True - verbose_proxy_logger.setLevel(logging.DEBUG) - await litellm.proxy.proxy_server.prisma_client.connect() - request = GenerateKeyRequest(max_budget=0.00001) - key = await generate_key_fn( - request, - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key=MASTER_KEY, - user_id="1234", - ), - ) - print(key) - - generated_key = key.key - user_id = key.user_id - bearer_token = "Bearer " + generated_key - print(f"generated_key: {generated_key}") - request = Request(scope={"type": "http"}) - request._url = URL(url="/chat/completions") - - # use generated key to auth in - result = await user_api_key_auth(request=request, api_key=bearer_token) - print("result from user auth with new key", result) - - # update spend using track_cost callback, make 2nd request, it should fail - import time - from litellm._uuid import uuid - - from litellm import Choices, Message, ModelResponse, Usage - from litellm.proxy.proxy_server import ProxyDBLogger - - proxy_db_logger = ProxyDBLogger() - - request_id = f"chatcmpl-e41836bb-bb8b-4df2-8e70-8f3e160155ac{uuid.uuid4()}" - resp = ModelResponse( - id=request_id, - choices=[ - Choices( - finish_reason=None, - index=0, - message=Message( - content=" Sure! Here is a short poem about the sky:\n\nA canvas of blue, a", - role="assistant", - ), - ) - ], - model="gpt-35-turbo", # azure always has model written like this - usage=Usage(prompt_tokens=210, completion_tokens=200, total_tokens=410), - ) - await proxy_db_logger._PROXY_track_cost_callback( - kwargs={ - "call_type": "acompletion", - "model": "sagemaker-chatgpt-v-3", - "stream": True, - "complete_streaming_response": resp, - "litellm_params": { - "metadata": { - "user_api_key": hash_token(generated_key), - "user_api_key_user_id": user_id, - } - }, - "response_cost": 0.00005, - }, - completion_response=resp, - start_time=datetime.now(), - end_time=datetime.now(), - ) - await update_spend( - prisma_client=prisma_client, - db_writer_client=None, - proxy_logging_obj=proxy_logging_obj, - ) - # use generated key to auth in - with pytest.raises(Exception, match="Budget has been exceeded") as exc_info: - await user_api_key_auth(request=request, api_key=bearer_token) - # Handle DataError and other exceptions that don't have .message attribute - assert "Budget has been exceeded" in getattr(exc_info.value, "message", str(exc_info.value)) - - -@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") -@pytest.mark.asyncio() -async def test_aview_spend_per_user(prisma_client): - setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) - setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY) - await litellm.proxy.proxy_server.prisma_client.connect() - try: - user_by_spend = await spend_user_fn(user_id=None) - assert type(user_by_spend) == list - assert len(user_by_spend) > 0 - first_user = user_by_spend[0] - - print("\nfirst_user=", first_user) - assert first_user["spend"] >= 0 - except Exception as e: - print("Got Exception", e) - pytest.fail(f"Got exception {e}") - - -@pytest.mark.asyncio() -@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") -async def test_view_spend_per_key(prisma_client): - """ - Test viewing spend per key. - """ - setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) - setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY) - await litellm.proxy.proxy_server.prisma_client.connect() - try: - # First create a key to ensure there's data to query - request = GenerateKeyRequest(models=["gpt-3.5-turbo"], max_budget=100) - key = await generate_key_fn( - request, - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key=MASTER_KEY, - user_id="test_user_spend", - ), - ) - print(f"Created test key: {key.key}") - - # Now query spend - key_by_spend = await spend_key_fn() - assert type(key_by_spend) == list - - # The list might be empty if no spend has been recorded yet - that's okay - if len(key_by_spend) > 0: - first_key = key_by_spend[0] - print("\nfirst_key=", first_key) - assert first_key.spend >= 0 - else: - print("No keys with spend found (expected for new database)") - except Exception as e: - print(f"Got Exception: {e}") - # If it's a 400 error with empty message, it might be an empty database - that's okay - error_str = str(e) - if "400" in error_str and ( - "error" in error_str.lower() or not error_str.strip() - ): - print("Empty database or no spend data - test passes") - else: - pytest.fail(f"Got unexpected exception {e}") - - @pytest.mark.asyncio() @pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") async def test_key_name_null(prisma_client): @@ -2185,40 +1135,6 @@ async def test_key_name_null(prisma_client): os.environ["DISABLE_KEY_NAME"] = "False" -@pytest.mark.asyncio() -@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") -async def test_key_name_set(prisma_client): - """ - - create key - - get key info - - assert key_name is not null - """ - setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) - setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY) - setattr(litellm.proxy.proxy_server, "general_settings", {"allow_user_auth": True}) - await litellm.proxy.proxy_server.prisma_client.connect() - try: - request = GenerateKeyRequest() - key = await generate_key_fn( - request, - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key=MASTER_KEY, - user_id="1234", - ), - ) - generated_key = key.key - result = await info_key_fn( - key=generated_key, - user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), - ) - print("result from info_key_fn", result) - assert isinstance(result["info"]["key_name"], str) - except Exception as e: - print("Got Exception", e) - pytest.fail(f"Got exception {e}") - - @pytest.mark.asyncio() @pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") async def test_default_key_params(prisma_client): @@ -2374,34 +1290,6 @@ def test_get_bearer_token(): assert result == "sk-9876", f"Expected 'sk-9876', got '{result}'" -@pytest.mark.asyncio -@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") -async def test_update_logs_with_spend_logs_url(prisma_client): - """ - Unit test for making sure spend logs list is still updated when url passed in - """ - from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter - - db_spend_update_writer = DBSpendUpdateWriter() - - payload = {"startTime": datetime.now(), "endTime": datetime.now()} - await db_spend_update_writer._insert_spend_log_to_db( - payload=payload, prisma_client=prisma_client - ) - - assert len(prisma_client.spend_log_transactions) > 0 - - prisma_client.spend_log_transactions = [] - - spend_logs_url = "" - payload = {"startTime": datetime.now(), "endTime": datetime.now()} - await db_spend_update_writer._insert_spend_log_to_db( - payload=payload, spend_logs_url=spend_logs_url, prisma_client=prisma_client - ) - - assert len(prisma_client.spend_log_transactions) > 0 - - @pytest.mark.asyncio @pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") async def test_user_api_key_auth(prisma_client): @@ -2442,241 +1330,6 @@ async def test_user_api_key_auth(prisma_client): ) -@pytest.mark.asyncio -@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") -async def test_user_api_key_auth_without_master_key(prisma_client): - # if master key is not set, expect all calls to go through - try: - from litellm.proxy.proxy_server import ProxyException - - setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) - setattr(litellm.proxy.proxy_server, "master_key", None) - setattr( - litellm.proxy.proxy_server, "general_settings", {"allow_user_auth": True} - ) - await litellm.proxy.proxy_server.prisma_client.connect() - - request = Request(scope={"type": "http"}) - request._url = URL(url="/chat/completions") - # Test case: No API Key passed in - - await user_api_key_auth(request, api_key=None) - await user_api_key_auth(request, api_key="my_token") - await user_api_key_auth(request, api_key="") - await user_api_key_auth(request, api_key="Bearer " + "1234") - except Exception as e: - print("Got Exception", e) - pytest.fail(f"Got exception {e}") - - -@pytest.mark.asyncio -@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") -async def test_key_with_no_permissions(prisma_client): - """ - - create key - - get key info - - assert key_name is null - """ - setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) - setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY) - setattr(litellm.proxy.proxy_server, "general_settings", {"allow_user_auth": False}) - await litellm.proxy.proxy_server.prisma_client.connect() - try: - response = await generate_key_helper_fn( - request_type="key", - **{"duration": "1hr", "key_max_budget": 0, "models": [], "aliases": {}, "config": {}, "spend": 0, "user_id": "ishaan", "team_id": "litellm-dashboard"}, # type: ignore - ) - - print(response) - key = response["token"] - - # make a /chat/completions call -> it should fail - request = Request(scope={"type": "http"}) - request._url = URL(url="/chat/completions") - - # use generated key to auth in - result = await user_api_key_auth(request=request, api_key="Bearer " + key) - print("result from user auth with new key", result) - pytest.fail(f"This should have failed!. IT's an invalid key") - except Exception as e: - print("Got Exception", e) - print(e.message) - - -async def track_cost_callback_helper_fn(generated_key: str, user_id: str): - from litellm._uuid import uuid - - from litellm import Choices, Message, ModelResponse, Usage - from litellm.proxy.proxy_server import ProxyDBLogger - - request_id = f"chatcmpl-e41836bb-bb8b-4df2-8e70-8f3e160155ac{uuid.uuid4()}" - resp = ModelResponse( - id=request_id, - choices=[ - Choices( - finish_reason=None, - index=0, - message=Message( - content=" Sure! Here is a short poem about the sky:\n\nA canvas of blue, a", - role="assistant", - ), - ) - ], - model="gpt-35-turbo", # azure always has model written like this - usage=Usage(prompt_tokens=210, completion_tokens=200, total_tokens=410), - ) - proxy_db_logger = ProxyDBLogger() - await proxy_db_logger._PROXY_track_cost_callback( - kwargs={ - "call_type": "acompletion", - "model": "sagemaker-chatgpt-v-3", - "stream": True, - "complete_streaming_response": resp, - "litellm_params": { - "metadata": { - "user_api_key": hash_token(generated_key), - "user_api_key_user_id": user_id, - } - }, - "response_cost": 0.00005, - }, - completion_response=resp, - start_time=datetime.now(), - end_time=datetime.now(), - ) - - -@pytest.mark.skip(reason="High traffic load test for spend tracking") -@pytest.mark.asyncio -async def test_proxy_load_test_db(prisma_client): - """ - Run 1500 req./s against track_cost_callback function - """ - setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) - setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY) - import logging - import time - - from litellm._logging import verbose_proxy_logger - - litellm.set_verbose = True - verbose_proxy_logger.setLevel(logging.DEBUG) - try: - start_time = time.time() - await litellm.proxy.proxy_server.prisma_client.connect() - request = GenerateKeyRequest(max_budget=0.00001) - key = await generate_key_fn( - request, - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key=MASTER_KEY, - user_id="1234", - ), - ) - print(key) - - generated_key = key.key - user_id = key.user_id - bearer_token = "Bearer " + generated_key - - request = Request(scope={"type": "http"}) - request._url = URL(url="/chat/completions") - - # use generated key to auth in - result = await user_api_key_auth(request=request, api_key=bearer_token) - print("result from user auth with new key", result) - # update spend using track_cost callback, make 2nd request, it should fail - n = 5000 - tasks = [ - track_cost_callback_helper_fn(generated_key=generated_key, user_id=user_id) - for _ in range(n) - ] - completions = await asyncio.gather(*tasks) - await asyncio.sleep(120) - try: - # call spend logs - spend_logs = await view_spend_logs( - api_key=generated_key, - user_api_key_dict=UserAPIKeyAuth(api_key=generated_key), - ) - - print(f"len responses: {len(spend_logs)}") - assert len(spend_logs) == n - print(n, time.time() - start_time, len(spend_logs)) - except Exception: - print(n, time.time() - start_time, 0) - raise Exception(f"it worked! key={key.key}") - except Exception as e: - pytest.fail(f"An exception occurred - {str(e)}") - - -@pytest.mark.asyncio() -@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") -async def test_master_key_hashing(prisma_client): - try: - from litellm._uuid import uuid - - print("prisma client=", prisma_client) - - master_key = MASTER_KEY - - setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) - setattr(litellm.proxy.proxy_server, "master_key", master_key) - - await litellm.proxy.proxy_server.prisma_client.connect() - from litellm.proxy.proxy_server import user_api_key_cache - - _team_id = "ishaans-special-team_{}".format(uuid.uuid4()) - user_api_key_dict = UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key=MASTER_KEY, - user_id="1234", - ) - await new_team( - NewTeamRequest(team_id=_team_id), - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key=MASTER_KEY, - user_id="1234", - ), - http_request=Request(scope={"type": "http"}), - ) - - _response = await new_user( - data=NewUserRequest( - models=["azure-gpt-3.5"], - team_id=_team_id, - tpm_limit=20, - ), - user_api_key_dict=user_api_key_dict, - ) - print(_response) - assert _response.models == ["azure-gpt-3.5"] - assert _response.team_id == _team_id - assert _response.tpm_limit == 20 - - bearer_token = "Bearer " + master_key - - request = Request(scope={"type": "http"}) - request._url = URL(url="/chat/completions") - - # use generated key to auth in - result: UserAPIKeyAuth = await user_api_key_auth( - request=request, api_key=bearer_token - ) - - # Master-key auth substitutes a stable alias so the master key (or - # its hash) never propagates into spend logs / metrics / audit trails. - from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS - - assert result.api_key == LITELLM_PROXY_MASTER_KEY_ALIAS - assert result.api_key != hash_token(master_key) - - except Exception as e: - print("Got Exception", e) - pytest.fail(f"Got exception {e}") - - @pytest.mark.asyncio @pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") async def test_reset_spend_authentication(prisma_client): @@ -2884,67 +1537,6 @@ async def test_create_update_team(prisma_client): ) -@pytest.mark.asyncio() -@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") -async def test_update_user_role(prisma_client): - """ - Tests if we update user role, incorrect values are not stored in cache - -> create a user with role == INTERNAL_USER - -> access an Admin only route -> expect to fail - -> update user role to == PROXY_ADMIN - -> access an Admin only route -> expect to succeed - """ - setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) - setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY) - await litellm.proxy.proxy_server.prisma_client.connect() - key = await new_user( - data=NewUserRequest( - user_role=LitellmUserRoles.INTERNAL_USER, - ) - ) - - print(key) - api_key = "Bearer " + key.key - - api_route = APIRoute(path="/global/spend", endpoint=global_spend) - request = Request( - { - "type": "http", - "route": api_route, - "path": "/global/spend", - "headers": [("Authorization", api_key)], - } - ) - - request._url = URL(url="/global/spend") - - # use generated key to auth in - try: - result = await user_api_key_auth(request=request, api_key=api_key) - print("result from user auth with new key", result) - except Exception as e: - print(e) - pass - - await user_update( - data=UpdateUserRequest( - user_id=key.user_id, user_role=LitellmUserRoles.PROXY_ADMIN - ), - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key=MASTER_KEY, - user_id="1234", - ), - ) - - # await asyncio.sleep(3) - - # use generated key to auth in - print("\n\nMAKING NEW REQUEST WITH UPDATED USER ROLE\n\n") - result = await user_api_key_auth(request=request, api_key=api_key) - print("result from user auth with new key", result) - - @pytest.mark.asyncio() @pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") async def test_update_user_unit_test(prisma_client): @@ -3625,71 +2217,6 @@ async def test_list_keys(prisma_client): assert _key in response["keys"] -@pytest.mark.asyncio -@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") -async def test_key_aliases(prisma_client): - """ - Test the key_aliases function: - - Returns a paginated response - - Includes alias from a newly created key - - Aliases are sorted - - Pagination and search params work correctly - """ - import asyncio - import uuid - import litellm - from litellm.proxy._types import LitellmUserRoles - - # Wire up test prisma client - setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) - setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY) - await litellm.proxy.proxy_server.prisma_client.connect() - - # Basic call - check pagination response shape - response = await key_aliases(page=1, size=50) - assert "aliases" in response - assert isinstance(response["aliases"], list) - assert "total_count" in response - assert "current_page" in response - assert "total_pages" in response - assert "size" in response - assert response["current_page"] == 1 - assert response["size"] == 50 - - # Create a new user (and key) with a unique alias - unique_id = str(uuid.uuid4()) - test_alias = f"key-aliases-test-{unique_id}" - test_user_id = f"key-aliases-user-{unique_id}" - - await new_user( - data=NewUserRequest( - user_id=test_user_id, - user_role=LitellmUserRoles.INTERNAL_USER, - key_alias=test_alias, - ), - user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), - ) - - # Allow async DB writes to settle - await asyncio.sleep(2) - - # Call again and validate alias is present - response_after = await key_aliases(page=1, size=50) - aliases = response_after["aliases"] - assert test_alias in aliases - assert aliases == sorted(aliases) - - # Search by partial alias - partial = test_alias[:10] - search_response = await key_aliases(page=1, size=50, search=partial) - assert test_alias in search_response["aliases"] - - # Search with no match - no_match_response = await key_aliases(page=1, size=50, search="__no_match_xyz__") - assert len(no_match_response["aliases"]) == 0 - assert no_match_response["total_count"] == 0 - - @pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") @pytest.mark.asyncio async def test_auth_vertex_ai_route(prisma_client): @@ -3827,123 +2354,6 @@ async def test_user_api_key_auth_db_unavailable_not_allowed(monkeypatch): ## E2E Virtual Key + Secret Manager Tests ######################################### -@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") -@pytest.mark.asyncio -@mock.patch( - "litellm.secret_managers.aws_secret_manager_v2.AWSSecretsManagerV2.async_write_secret" -) -@mock.patch( - "litellm.secret_managers.aws_secret_manager_v2.AWSSecretsManagerV2.async_read_secret" -) -@mock.patch( - "litellm.secret_managers.aws_secret_manager_v2.AWSSecretsManagerV2.async_delete_secret" -) -async def test_key_generate_with_secret_manager_call( - mock_delete_secret, mock_read_secret, mock_write_secret, prisma_client -): - """ - Generate a key - assert it exists in the secret manager - - delete the key - assert it is deleted from the secret manager - """ - from litellm.secret_managers.aws_secret_manager_v2 import AWSSecretsManagerV2 - from litellm.types.secret_managers.main import ( - KeyManagementSystem, - KeyManagementSettings, - ) - - from litellm.proxy.hooks.key_management_event_hooks import ( - LITELLM_PREFIX_STORED_VIRTUAL_KEYS, - ) - - litellm.set_verbose = True - - #### Test Setup ############################################################ - aws_secret_manager_client = AWSSecretsManagerV2() - litellm.secret_manager_client = aws_secret_manager_client - litellm._key_management_system = KeyManagementSystem.AWS_SECRET_MANAGER - litellm._key_management_settings = KeyManagementSettings( - store_virtual_keys=True, - ) - general_settings = { - "key_management_system": "aws_secret_manager", - "key_management_settings": { - "store_virtual_keys": True, - }, - } - - setattr(litellm.proxy.proxy_server, "general_settings", general_settings) - setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY) - setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) - await litellm.proxy.proxy_server.prisma_client.connect() - ############################################################################ - - # generate new key - key_alias = f"test_alias_secret_manager_key-{uuid.uuid4()}" - spend = 100 - max_budget = 400 - models = ["fake-openai-endpoint"] - - # Mock write_secret to return success - mock_write_secret.return_value = None - - new_key = await generate_key_fn( - data=GenerateKeyRequest( - key_alias=key_alias, spend=spend, max_budget=max_budget, models=models - ), - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key=MASTER_KEY, - user_id="1234", - ), - ) - - generated_key = new_key.key - print(generated_key) - - await asyncio.sleep(2) - - # read from the secret manager - # Mock read_secret to return the generated key - mock_read_secret.return_value = generated_key - - result = await aws_secret_manager_client.async_read_secret( - secret_name=f"{litellm._key_management_settings.prefix_for_stored_virtual_keys}{key_alias}" - ) - - # Assert the correct key is stored in the secret manager - print("response from AWS Secret Manager") - print(result) - assert result == generated_key - - # Mock delete_secret to return success - mock_delete_secret.return_value = None - - # delete the key - await delete_key_fn( - data=KeyRequest(keys=[generated_key]), - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, api_key=MASTER_KEY, user_id="1234" - ), - ) - - await asyncio.sleep(2) - - # Assert the key is deleted from the secret manager - # Mock read_secret to return None after deletion - mock_read_secret.return_value = None - - result = await aws_secret_manager_client.async_read_secret( - secret_name=f"{litellm._key_management_settings.prefix_for_stored_virtual_keys}{key_alias}" - ) - assert result is None - - # cleanup - setattr(litellm.proxy.proxy_server, "general_settings", {}) - - ################################################################################ @@ -4120,248 +2530,3 @@ def test_should_track_cost_callback(): ) -@pytest.mark.asyncio -@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") -async def test_get_paginated_teams(prisma_client): - """ - Test the get_paginated_teams function: - 1. Test pagination returns valid results - 2. Test total count matches across pages - 3. Test page size is respected - """ - from litellm.proxy.management_endpoints.team_endpoints import get_paginated_teams - - setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) - setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY) - await litellm.proxy.proxy_server.prisma_client.connect() - - try: - # Get first page with page_size=2 - teams_page_1, total_count_1 = await get_paginated_teams( - prisma_client=prisma_client, page_size=2, page=1 - ) - - print("teams_page_1=", teams_page_1) - print("total_count_1=", total_count_1) - - # Get second page - teams_page_2, total_count_2 = await get_paginated_teams( - prisma_client=prisma_client, page_size=2, page=2 - ) - - print("teams_page_2=", teams_page_2) - print("total_count_2=", total_count_2) - - # Verify results - assert isinstance(teams_page_1, list) # Should return a list - assert isinstance(total_count_1, int) # Should return an integer count - assert ( - total_count_1 == total_count_2 - ) # Total count should be consistent across pages - assert len(teams_page_1) <= 2 # Should respect page_size limit - - except Exception as e: - print(f"Error occurred: {e}") - pytest.fail(f"Test failed with exception: {e}") - - -@pytest.mark.asyncio -@pytest.mark.flaky(retries=3, delay=1) -@pytest.mark.parametrize("entity_type", ["key", "user", "team"]) -@pytest.mark.skip( - reason="Skipping reset budget job test. Fails on ci/cd due to db timeout errors. Need to replace with mock db." -) -async def test_reset_budget_job(prisma_client, entity_type): - """ - Test that the ResetBudgetJob correctly resets budgets for keys, users, and teams. - - For each entity type: - 1. Create a new entity with max_budget=100, spend=99, budget_duration=5s - 2. Call the reset_budget function - 3. Verify the entity's spend is reset to 0 and budget_reset_at is updated - """ - from datetime import datetime, timedelta - import time - - from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob - from litellm.proxy.utils import ProxyLogging - - # Setup - setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) - setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY) - await litellm.proxy.proxy_server.prisma_client.connect() - - proxy_logging_obj = ProxyLogging(user_api_key_cache=None) - reset_budget_job = ResetBudgetJob( - proxy_logging_obj=proxy_logging_obj, prisma_client=prisma_client - ) - - # Create entity based on type - entity_id = None - if entity_type == "key": - # Create a key with specific budget settings - key = await generate_key_fn( - data=GenerateKeyRequest( - max_budget=100, - budget_duration="5s", - ), - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key=MASTER_KEY, - user_id="1234", - ), - ) - entity_id = key.token_id - print("generated key=", key) - - # Update the key to set spend and reset_at to now - updated = await prisma_client.db.litellm_verificationtoken.update_many( - where={"token": key.token_id}, - data={ - "spend": 99.0, - }, - ) - print("Updated key=", updated) - - elif entity_type == "user": - # Create a user with specific budget settings - user = await new_user( - data=NewUserRequest( - max_budget=100, - budget_duration="5s", - ), - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key=MASTER_KEY, - user_id="1234", - ), - ) - entity_id = user.user_id - - # Update the user to set spend and reset_at to now - await prisma_client.db.litellm_usertable.update_many( - where={"user_id": user.user_id}, - data={ - "spend": 99.0, - }, - ) - - elif entity_type == "team": - # Create a team with specific budget settings - team_id = f"test-team-{uuid.uuid4()}" - team = await new_team( - NewTeamRequest( - team_id=team_id, - max_budget=100, - budget_duration="5s", - ), - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key=MASTER_KEY, - user_id="1234", - ), - http_request=Request(scope={"type": "http"}), - ) - entity_id = team_id - - # Update the team to set spend and reset_at to now - current_time = datetime.utcnow() - await prisma_client.db.litellm_teamtable.update( - where={"team_id": team_id}, - data={ - "spend": 99.0, - }, - ) - - # Verify entity was created and updated with spend - if entity_type == "key": - entity_before = await prisma_client.db.litellm_verificationtoken.find_unique( - where={"token": entity_id} - ) - elif entity_type == "user": - entity_before = await prisma_client.db.litellm_usertable.find_unique( - where={"user_id": entity_id} - ) - elif entity_type == "team": - entity_before = await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": entity_id} - ) - - assert entity_before is not None - assert entity_before.spend == 99.0 - - # Wait for 5 seconds to pass - print("sleeping for 5 seconds") - time.sleep(5) - - # Call the reset_budget function - await reset_budget_job.reset_budget() - - # Verify the entity's spend is reset and budget_reset_at is updated - if entity_type == "key": - entity_after = await prisma_client.db.litellm_verificationtoken.find_unique( - where={"token": entity_id} - ) - elif entity_type == "user": - entity_after = await prisma_client.db.litellm_usertable.find_unique( - where={"user_id": entity_id} - ) - elif entity_type == "team": - entity_after = await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": entity_id} - ) - - assert entity_after is not None - assert entity_after.spend == 0.0 - - -@pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") -def test_delete_nonexistent_key_returns_404(prisma_client): - # Try to delete a key that does not exist, expect a 404 error - import random, string - from litellm.proxy._types import ( - KeyRequest, - UserAPIKeyAuth, - LitellmUserRoles, - ProxyException, - ) - from litellm.proxy.management_endpoints.key_management_endpoints import ( - delete_key_fn, - ) - from starlette.datastructures import URL - from fastapi import Request - - print("prisma client=", prisma_client) - setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) - setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY) - try: - - async def test(): - await litellm.proxy.proxy_server.prisma_client.connect() - # Generate a random key that does not exist - random_key = "sk-" + "".join( - random.choices(string.ascii_letters + string.digits, k=24) - ) - delete_key_request = KeyRequest(keys=[random_key]) - bearer_token = "Bearer sk-9876" - request = Request(scope={"type": "http"}) - request._url = URL(url="/key/delete") - # use admin to auth in - result = await litellm.proxy.proxy_server.user_api_key_auth( - request=request, api_key=bearer_token - ) - result.user_role = LitellmUserRoles.PROXY_ADMIN - with pytest.raises(ProxyException) as exc_info: - await delete_key_fn(data=delete_key_request, user_api_key_dict=result) - e = exc_info.value - print("Caught ProxyException:", e) - assert str(e.code) == "404" - assert "No keys found" in str( - e.message - ) or "No matching keys or aliases found to delete" in str(e.message) - - import asyncio - - asyncio.run(test()) - except Exception as e: - pytest.fail(f"An exception occurred - {str(e)}") diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index 6ae74e1b740..2e18a0c1099 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -6983,6 +6983,22 @@ async def test_info_key_fn_unknown_key_still_404s(monkeypatch): assert exc_info.value.code == "404" +@pytest.mark.asyncio +async def test_info_key_fn_no_key_returns_400_not_500(monkeypatch): + """#43571: when both ?key= and auth api_key are None, return a clean 400 + instead of letting Prisma raise MissingRequiredValueError as a 500.""" + from litellm.proxy.management_endpoints.key_management_endpoints import info_key_fn + + mock_prisma_client = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key=None) + with pytest.raises(ProxyException) as exc_info: + await info_key_fn(key=None, user_api_key_dict=auth) + assert exc_info.value.code == "400" + mock_prisma_client.db.litellm_verificationtoken.find_unique.assert_not_called() + + @pytest.mark.asyncio @pytest.mark.parametrize( ("blocked", "expires", "expected_status"), @@ -8680,6 +8696,7 @@ async def test_reset_key_spend_success(monkeypatch): mock_spend_counter_cache.in_memory_cache.set_cache.assert_any_call( key=f"spend_db_floor:spend:key:{hashed_key}", value=50.0, ttl=5 ) + assert spend.deleted() == {f"spend:committed:key:{hashed_key}"} @pytest.mark.asyncio @@ -9000,7 +9017,10 @@ async def test_update_key_spend_updates_counter(monkeypatch): mock_delete_cache.assert_awaited_once() assert spend.memory == {f"spend:key:{hashed_key}": 0.0} - assert spend.commands == [f"SET spend:key:{hashed_key} 0.0 60"] + assert spend.commands == [ + f"SET spend:key:{hashed_key} 0.0 60", + f"DEL spend:committed:key:{hashed_key}", + ] @pytest.mark.asyncio @@ -19780,6 +19800,150 @@ async def test_bulk_key_update_batch_enqueued_token_limit_rejected_for_non_admin assert "Only proxy admins can set" in str(exc.value.detail) + +_INDEX_GRANT = [{"index_name": "docs", "index_permissions": ["read", "write"]}] + + +@pytest.mark.parametrize( + "label, top_level, sent, stored, allowed", + [ + ("granted top-level on a key with none stored", _INDEX_GRANT, _UNTOUCHED, None, False), + ("granted in metadata on a key with none stored", None, _INDEX_GRANT, None, False), + ("widened from read to read and write", None, _INDEX_GRANT, [{"index_name": "docs", "index_permissions": ["read"]}], False), + ("cleared by replacing the blob", None, None, _INDEX_GRANT, False), + ("resent unchanged in metadata", None, _INDEX_GRANT, _INDEX_GRANT, True), + ("resent unchanged top-level", _INDEX_GRANT, _UNTOUCHED, _INDEX_GRANT, True), + ("left untouched", None, _UNTOUCHED, _INDEX_GRANT, True), + ("cleared top-level with an empty list", [], _UNTOUCHED, _INDEX_GRANT, False), + ("empty top-level list on a key with none stored", [], _UNTOUCHED, None, True), + ], +) +def test_vector_store_index_grants_admin_gate_matrix( + label: str, + top_level: list[dict[str, object]] | None, + sent: object, + stored: list[dict[str, object]] | None, + allowed: bool, +) -> None: + from litellm.proxy.auth.auth_utils import ( + enforce_vector_store_index_grants_are_admin_only, + ) + + request_body = {} if sent is _UNTOUCHED else {"metadata": {} if sent is None else {"allowed_vector_store_indexes": sent}} + if top_level is not None: + request_body["allowed_vector_store_indexes"] = top_level + data = UpdateKeyRequest(key="sk-1", **request_body) + existing_metadata = None if stored is None else {"allowed_vector_store_indexes": stored} + + def _call(caller: UserAPIKeyAuth) -> None: + enforce_vector_store_index_grants_are_admin_only( + data=data, + top_level_grants=data.allowed_vector_store_indexes, + existing_metadata=existing_metadata, + user_api_key_dict=caller, + entity="key", + ) + + non_admin = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-non-admin", user_id="alice") + if allowed: + _call(non_admin) + else: + with pytest.raises(HTTPException) as exc: + _call(non_admin) + assert exc.value.status_code == 403 + assert "Only proxy admins can set allowed_vector_store_indexes" in str(exc.value.detail) + + _call(UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin")) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "request_fields", + [{"allowed_vector_store_indexes": _INDEX_GRANT}, {"metadata": {"allowed_vector_store_indexes": _INDEX_GRANT}}], + ids=["top-level", "metadata"], +) +async def test_generate_key_vector_store_index_grants_rejected_for_non_admin(request_fields: dict[str, object]) -> None: + """An index grant on a key opens document reads and writes on that index, so a + non-admin minting their own key cannot pick its grants.""" + with patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()): + with pytest.raises(HTTPException) as exc: + await _common_key_generation_helper( + data=GenerateKeyRequest(**request_fields), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-alice", + user_id="alice", + ), + litellm_changed_by=None, + team_table=None, + ) + assert int(getattr(exc.value, "status_code", 0)) == 403 + assert "Only proxy admins can set allowed_vector_store_indexes" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_update_key_vector_store_index_grants_unchanged_allows_non_admin_edit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A key owner renaming a key that carries admin-set index grants resends them + with the rest of the form, which must not count as setting them.""" + from litellm.proxy.management_endpoints.key_management_endpoints import ( + update_key_fn, + ) + + token = "b3b2c3d4e5f6789012345678901234567890123456789012345678901234abcd" + _wire_update_key_fn(monkeypatch, _estimate_key_row(token, {"allowed_vector_store_indexes": _INDEX_GRANT})) + + mock_request = MagicMock() + mock_request.query_params = {} + + result = await update_key_fn( + request=mock_request, + data=UpdateKeyRequest(key=token, key_alias="my-alias", metadata={"allowed_vector_store_indexes": _INDEX_GRANT}), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-internal", + user_id="internal_user", + ), + litellm_changed_by=None, + ) + + assert result is not None + + +@pytest.mark.asyncio +async def test_update_key_vector_store_index_grants_widened_rejected_for_non_admin( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy.management_endpoints.key_management_endpoints import ( + update_key_fn, + ) + + token = "c3b2c3d4e5f6789012345678901234567890123456789012345678901234abcd" + _wire_update_key_fn( + monkeypatch, + _estimate_key_row(token, {"allowed_vector_store_indexes": [{"index_name": "docs", "index_permissions": ["read"]}]}), + ) + + mock_request = MagicMock() + mock_request.query_params = {} + + with pytest.raises(ProxyException) as exc: + await update_key_fn( + request=mock_request, + data=UpdateKeyRequest(key=token, allowed_vector_store_indexes=_INDEX_GRANT), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-internal", + user_id="internal_user", + ), + litellm_changed_by=None, + ) + + assert str(exc.value.code) == "403" + assert "Only proxy admins can set allowed_vector_store_indexes" in str(exc.value.message) + + @pytest.mark.asyncio async def test_execute_virtual_key_regeneration_stamps_settings_updated_at(): """Regenerate rewrites the key's config, so it must move settings_updated_at.""" 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 2fdb77af8f9..2b269469f89 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -1571,7 +1571,7 @@ def deleted_auto_router_catalog(monkeypatch): config = proxy_server.ProxyConfig() config.auto_router_db_catalog = build_auto_router_catalog(rows) monkeypatch.setattr(proxy_server, "proxy_config", config) - monkeypatch.setattr(proxy_server, "MODEL_RECONCILE_LOCK", asyncio.Lock()) + monkeypatch.setattr(proxy_server, "DB_RECONCILE_LOCK", asyncio.Lock()) monkeypatch.setattr(proxy_server, "llm_router", Router(model_list=[])) monkeypatch.setattr(proxy_server, "_license_check", SimpleNamespace(auto_router_capability_limit=lambda: 1)) monkeypatch.setattr(proxy_server, "heuristic_v1_tuning_baselines", {}) @@ -4869,6 +4869,146 @@ class TestModelInfoCostMapEchoFilter: "cost-map metadata must not be persisted from an unchanged /model/info echo" ) + def test_echo_of_an_expanded_wildcard_row_is_not_persisted_on_the_wildcard(self): + import litellm + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + router: Final = litellm.Router( + model_list=[ + { + "model_name": "my-openai/*", + "litellm_params": {"model": "openai/*", "api_key": "x"}, + "model_info": {"id": "dep-echo-wildcard"}, + } + ] + ) + entry: Final = litellm.get_model_info("openai/gpt-5.6") + echo: Final = {**entry, "id": "dep-echo-wildcard", "db_model": True, "access_groups": ["prod"]} + db_model: Final = Deployment( + model_name="my-openai/*", + litellm_params=LiteLLM_Params(model="openai/*"), + model_info=ModelInfo(id="dep-echo-wildcard"), + ) + + result: Final = update_db_model( + db_model=db_model, + updated_patch=updateDeployment(model_name="my-openai/gpt-5.6", model_info=ModelInfo(**echo)), + llm_router=router, + ) + + info: Final = json.loads(result["model_info"]) + assert info["access_groups"] == ["prod"] + assert set(info).isdisjoint(entry), "one listed model's metadata must not be saved onto the whole wildcard" + + def test_echo_saved_on_a_worker_that_has_not_loaded_the_wildcard_is_not_persisted(self): + import litellm + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + router_without_the_wildcard: Final = litellm.Router( + model_list=[ + { + "model_name": "openai-mini", + "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "x"}, + "model_info": {"id": "dep-other"}, + } + ] + ) + entry: Final = litellm.get_model_info("openai/gpt-5.6") + db_model: Final = Deployment( + model_name="my-openai/*", + litellm_params=LiteLLM_Params(model="openai/*"), + model_info=ModelInfo(id="dep-echo-wildcard"), + ) + + result: Final = update_db_model( + db_model=db_model, + updated_patch=updateDeployment( + model_name="my-openai/gpt-5.6", + model_info=ModelInfo(**{**entry, "id": "dep-echo-wildcard", "access_groups": ["prod"]}), + ), + llm_router=router_without_the_wildcard, + ) + + info: Final = json.loads(result["model_info"]) + assert info["access_groups"] == ["prod"] + assert set(info).isdisjoint(entry), "one listed model's metadata must not be saved onto the whole wildcard" + + def test_legacy_update_echo_of_an_expanded_wildcard_row_is_not_persisted_on_the_wildcard(self): + import litellm + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _ListedRow, + legacy_update_model_info_changes, + ) + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + router: Final = litellm.Router( + model_list=[ + { + "model_name": "my-openai/*", + "litellm_params": {"model": "openai/*", "api_key": "x"}, + "model_info": {"id": "dep-echo-wildcard"}, + } + ] + ) + entry: Final = litellm.get_model_info("openai/gpt-5.6") + db_model: Final = Deployment( + model_name="my-openai/*", + litellm_params=LiteLLM_Params(model="openai/*"), + model_info=ModelInfo(id="dep-echo-wildcard"), + ) + + changes: Final = legacy_update_model_info_changes( + db_model, + ModelInfo(**{**entry, "id": "dep-echo-wildcard", "access_groups": ["prod"]}), + loaded_catalog=lambda: {}, + listed_row=_ListedRow("my-openai/gpt-5.6", router), + ) + + assert dict(changes) == {"access_groups": ["prod"]} + + @pytest.mark.asyncio + async def test_echo_of_a_team_wildcard_row_is_not_persisted_on_the_wildcard(self): + import litellm + from litellm.proxy.management_endpoints.model_management_endpoints import _update_team_model_in_db + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + router: Final = litellm.Router( + model_list=[ + { + "model_name": "team-openai/*", + "litellm_params": {"model": "openai/*", "api_key": "x"}, + "model_info": {"id": "dep-team-wildcard", "team_id": "team-a"}, + } + ] + ) + entry: Final = litellm.get_model_info("openai/gpt-5.6") + echo: Final = {**entry, "id": "dep-team-wildcard", "team_id": "team-a", "access_groups": ["prod"]} + db_model: Final = Deployment( + model_name="team-openai/*", + litellm_params=LiteLLM_Params(model="openai/*"), + model_info=ModelInfo(id="dep-team-wildcard", team_id="team-a"), + ) + + with ( + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.management_endpoints.model_management_endpoints.team_model_add"), + patch("litellm.proxy.management_endpoints.model_management_endpoints.team_model_delete"), + ): + result: Final = await _update_team_model_in_db( + db_model=db_model, + patch_data=updateDeployment(model_name="team-openai/gpt-5.6", model_info=ModelInfo(**echo)), + user_api_key_dict=UserAPIKeyAuth(user_id="test_user", user_role=LitellmUserRoles.PROXY_ADMIN), + prisma_client=MockPrismaClient(team_exists=True), + write_row=_passthrough_row, + llm_router=router, + ) + + info: Final = json.loads(result["model_info"]) + assert info["access_groups"] == ["prod"] + assert set(info).isdisjoint(entry), "one listed model's metadata must not be saved onto a team's wildcard" + def test_an_edited_value_survives_the_echo_filter(self): import litellm from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model @@ -5959,7 +6099,7 @@ class TestConcurrentModelWritesDoNotEvictEachOther: @pytest.mark.asyncio async def test_reconciles_serialize_so_no_stale_snapshot_can_evict(self, monkeypatch): - """MODEL_RECONCILE_LOCK admits one reconcile at a time. + """DB_RECONCILE_LOCK admits one reconcile at a time. The fake body awaits, which is the whole point: without the lock the gather below parks all five inside the critical section at that await and observed depth goes @@ -6088,7 +6228,7 @@ class TestConcurrentModelWritesDoNotEvictEachOther: class TestDeleteEvictionsHoldTheReconcileLock: """A delete evicts from ``llm_router`` directly instead of reconciling, so it must - take MODEL_RECONCILE_LOCK to do it. + take DB_RECONCILE_LOCK to do it. The db row is gone by then, but a reconcile that snapshotted the db BEFORE the row was deleted still lists that id as desired, and its ``_add_deployment`` upserts the @@ -6102,7 +6242,7 @@ class TestDeleteEvictionsHoldTheReconcileLock: async def _assert_evicts_under_lock(monkeypatch, call_endpoint, model_id: str, config) -> None: """Run ``call_endpoint`` with the lock already held and assert it blocks. - Holding MODEL_RECONCILE_LOCK stands in for a reconcile that is mid-flight. If + Holding DB_RECONCILE_LOCK stands in for a reconcile that is mid-flight. If the eviction takes the lock it cannot run until we release; if it does not, it runs straight through and the deployment is evicted while the "reconcile" is still in its critical section -- exactly the interleaving that resurrects it. @@ -6115,7 +6255,7 @@ class TestDeleteEvictionsHoldTheReconcileLock: own, exactly as here. """ lock = asyncio.Lock() - monkeypatch.setattr("litellm.proxy.proxy_server.MODEL_RECONCILE_LOCK", lock) + monkeypatch.setattr("litellm.proxy.proxy_server.DB_RECONCILE_LOCK", lock) stale_catalog = config.auto_router_db_catalog async with lock: @@ -6124,7 +6264,7 @@ class TestDeleteEvictionsHoldTheReconcileLock: for _ in range(50): await asyncio.sleep(0) assert not task.done(), ( - f"deleting {model_id} did not wait for MODEL_RECONCILE_LOCK -- an " + f"deleting {model_id} did not wait for DB_RECONCILE_LOCK -- an " f"in-flight reconcile can resurrect the deployment it just evicted" ) config.auto_router_db_catalog = stale_catalog diff --git a/tests/unit/proxy/management_endpoints/test_organization_endpoints.py b/tests/unit/proxy/management_endpoints/test_organization_endpoints.py index 8d2c6c0054c..70cfaf4140d 100644 --- a/tests/unit/proxy/management_endpoints/test_organization_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_organization_endpoints.py @@ -1188,6 +1188,27 @@ async def _run_legacy_update_organization( return mock_prisma_client +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("body", "twin_after"), + [ + ({"organization_id": "org-1", "spend": 0.0}, None), + ({"organization_id": "org-1", "organization_alias": "renamed"}, 7.0), + ], +) +async def test_legacy_update_drops_the_committed_org_spend_only_when_it_writes_the_stored_spend( + monkeypatch, body, twin_after +): + from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters + + spend: Final = FakeSpendCounters(redis=None, memory={"spend:committed:org:org-1": 7.0}) + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_store", spend.store()) + + await _run_legacy_update_organization(monkeypatch, body=body, existing_budget_id="budget-1") + + assert spend.memory.get("spend:committed:org:org-1") == twin_after + + @pytest.mark.asyncio @pytest.mark.parametrize( "body", @@ -1582,6 +1603,49 @@ async def test_delete_organization_evicts_the_cache_of_the_keys_it_deletes(monke assert jwt_table.rows == (kept_row,) +@pytest.mark.asyncio +async def test_delete_organization_drops_the_committed_spend_of_the_org_and_its_teams(monkeypatch): + from litellm.proxy._types import DeleteOrganizationRequest, LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.management_endpoints.organization_endpoints import delete_organization + from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters + + spend: Final = FakeSpendCounters( + redis=None, + memory={ + "spend:committed:org:org-doomed": 7.0, + "spend:committed:team:team-in-org": 5.0, + "spend:committed:org:org-kept": 3.0, + "spend:committed:team:team-elsewhere": 2.0, + }, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_store", spend.store()) + prisma_client: Final = AsyncMock() + prisma_client.db.litellm_teamtable.find_many = AsyncMock( + return_value=[SimpleNamespace(team_id="team-in-org", organization_id="org-doomed")] + ) + prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + prisma_client.db.litellm_jwtkeymapping.find_many = AsyncMock(return_value=[]) + prisma_client.db.litellm_organizationtable.delete = AsyncMock(return_value=MagicMock()) + + @asynccontextmanager + async def transaction(): + yield prisma_client.db + + prisma_client.tx = transaction + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True, raising=False) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", None) + + await delete_organization( + data=DeleteOrganizationRequest(organization_ids=["org-doomed"]), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert spend.memory == {"spend:committed:org:org-kept": 3.0, "spend:committed:team:team-elsewhere": 2.0} + + @pytest.mark.asyncio async def test_delete_organization_is_refused_while_one_of_its_teams_holds_a_ptu_share(monkeypatch): """Deleting an organization deletes its teams, so a team still holding a PTU share stops it the diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index 94a8022fe96..0881c902a9d 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -9966,6 +9966,37 @@ async def test_delete_validated_team_rows_deletes_only_the_keys_its_rostered_mem assert deleted_keys.kwargs["where"] == {"token": {"in": ["hashed-member"]}} +@pytest.mark.asyncio +async def test_delete_validated_team_rows_drops_the_committed_spend_of_the_teams_it_deleted( + monkeypatch, + disable_audit_logging_for_mocked_team, +): + from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters + + spend: Final = FakeSpendCounters( + redis=None, memory={"spend:committed:team:team-doomed": 7.0, "spend:committed:team:team-kept": 3.0} + ) + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_store", spend.store()) + team = LiteLLM_TeamTable(team_id="team-doomed", members_with_roles=[], metadata={}, model_max_budget={}, model_spend={}) + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_teammembership.delete_many = AsyncMock() + mock_access_group_tx_cm = MagicMock() + mock_access_group_tx_cm.__aenter__ = AsyncMock(return_value=AsyncMock()) + mock_access_group_tx_cm.__aexit__ = AsyncMock(return_value=False) + mock_prisma_client.db.tx = MagicMock(return_value=mock_access_group_tx_cm) + _wire_team_delete_tx(mock_prisma_client, locked_teams=(team,), team_keys=()) + monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") + + await delete_validated_team_rows( + teams=(team,), + prisma_client=mock_prisma_client, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="hashed-admin"), + litellm_changed_by=None, + ) + + assert spend.memory == {"spend:committed:team:team-kept": 3.0} + + @pytest.mark.asyncio async def test_delete_validated_teams_writes_the_deleted_audit_row_only_for_a_team_still_there_under_the_lock( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unit/proxy/management_helpers/test_bulk_user_deletion.py b/tests/unit/proxy/management_helpers/test_bulk_user_deletion.py index e212a433393..880e5746150 100644 --- a/tests/unit/proxy/management_helpers/test_bulk_user_deletion.py +++ b/tests/unit/proxy/management_helpers/test_bulk_user_deletion.py @@ -2,11 +2,12 @@ import asyncio import copy import itertools import json -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence from contextlib import asynccontextmanager from typing import Final, cast import pytest +from pytest_mock import MockerFixture from pydantic import BaseModel, ConfigDict, ValidationError import litellm @@ -22,6 +23,7 @@ from litellm.proxy.management_helpers.bulk_user_deletion import ( ) from litellm.types.proxy.management_endpoints.internal_user_endpoints import BulkDeleteUserRequest from litellm.types.proxy.management_endpoints.team_endpoints import BulkTeamMemberDeleteRequest, TeamMemberRef +from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters ADMIN: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin") INTERNAL: Final = UserAPIKeyAuth(user_id="someone", user_role=LitellmUserRoles.INTERNAL_USER) @@ -697,6 +699,34 @@ async def test_bulk_delete_evicts_jwt_key_mappings_of_the_deleted_users_keys(): assert cache.get_cache(key=kept) is not None +async def _delete_through_scim_or_user_delete(prisma: _FakePrisma, user_ids: Sequence[str]) -> object: + return await delete_users_from_current_teams( + prisma_client=prisma, # pyright: ignore[reportArgumentType] # fake stands in for PrismaClient + user_ids=user_ids, + user_api_key_dict=ADMIN, + user_api_key_cache=UserApiKeyCache(), + proxy_logging_obj=None, + litellm_changed_by=None, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("delete", [_delete, _delete_through_scim_or_user_delete]) +async def test_user_deletes_drop_the_committed_spend_of_the_users_they_deleted( + delete: Callable[[_FakePrisma, Sequence[str]], Awaitable[object]], mocker: MockerFixture +) -> None: + spend: Final = FakeSpendCounters( + redis=None, memory={"spend:committed:user:u1": 7.0, "spend:committed:user:keep": 3.0} + ) + mocker.patch("litellm.proxy.proxy_server.spend_counter_store", spend.store()) # test-quality-ok: substitute the spend counter dependency + prisma: Final = _FakePrisma(users=[_user("u1"), _user("keep")]) + + _ = await delete(prisma, ["u1"]) + + assert sorted(prisma.db.litellm_usertable.rows) == ["keep"] + assert spend.memory == {"spend:committed:user:keep": 3.0} + + @pytest.mark.asyncio async def test_bulk_delete_rejects_non_admin_callers_before_touching_the_db(): prisma = _FakePrisma(users=[_user("u1")]) diff --git a/tests/unit/proxy/management_helpers/test_team_member_permission_checks.py b/tests/unit/proxy/management_helpers/test_team_member_permission_checks.py index 55d29724e38..544aa4ed21b 100644 --- a/tests/unit/proxy/management_helpers/test_team_member_permission_checks.py +++ b/tests/unit/proxy/management_helpers/test_team_member_permission_checks.py @@ -271,6 +271,7 @@ class TestEnforceMemberCanAssignAccessGroups: team = MagicMock() team.team_id = team_id team.team_member_permissions = team_member_permissions + team.access_group_ids = ["ag-1"] return team def test_no_access_group_ids_is_noop(self, monkeypatch): diff --git a/tests/unit/proxy/mcp_server/outbound_credentials/conftest.py b/tests/unit/proxy/mcp_server/outbound_credentials/conftest.py new file mode 100644 index 00000000000..bd80cba1c62 --- /dev/null +++ b/tests/unit/proxy/mcp_server/outbound_credentials/conftest.py @@ -0,0 +1,20 @@ +from collections.abc import Iterator + +import pytest + +from litellm.proxy.mcp_server.outbound_credentials.bridge_credentials import envelope_keys_from_master_key +from litellm.proxy.mcp_server.outbound_credentials.session_credentials import session_keys_from_master_key + + +@pytest.fixture +def openssl_enforcing_fips(fips_openssl: None) -> Iterator[None]: + envelope_keys_from_master_key.cache_clear() + session_keys_from_master_key.cache_clear() + yield + envelope_keys_from_master_key.cache_clear() + session_keys_from_master_key.cache_clear() + + +@pytest.fixture +def fips_mode_on(openssl_enforcing_fips: None, fips_mode: None) -> None: + return None diff --git a/tests/unit/proxy/mcp_server/outbound_credentials/test_bridge_credentials.py b/tests/unit/proxy/mcp_server/outbound_credentials/test_bridge_credentials.py index 47b6c3c2fa9..0ba19d13955 100644 --- a/tests/unit/proxy/mcp_server/outbound_credentials/test_bridge_credentials.py +++ b/tests/unit/proxy/mcp_server/outbound_credentials/test_bridge_credentials.py @@ -9,6 +9,7 @@ through the consumer; and no path leaks the upstream token in a repr. """ from datetime import datetime, timedelta, timezone +from typing import Final import pytest from pydantic import SecretStr @@ -318,3 +319,33 @@ def test_is_bridge_envelope_shaped_rejects_non_envelope_bearer(): assert is_bridge_envelope_shaped("Bearer sk-some-litellm-key") is False assert is_bridge_envelope_shaped("plain-upstream-token") is False assert is_bridge_envelope_shaped("") is False + + +def test_envelope_keys_without_fips_stay_the_scrypt_keys_of_the_master_key() -> None: + keys: Final = envelope_keys_from_master_key(_MASTER_KEY) + assert (keys.signing_key.get_secret_value(), keys.encryption_key.get_secret_value()) == ( + "34d3bb6ae2e03cddca3417209304e9603417b89258d3979d0569880b6fbd1bed", + "eebfbee8cbe767ee136ce3a9745b44de2b2ae35d1ea46dcf72c58b811fac3250", + ) + + +@pytest.mark.usefixtures("openssl_enforcing_fips") +def test_envelope_keys_on_a_fips_openssl_without_fips_mode_stay_the_scrypt_keys( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("LITELLM_SALT_KEY", raising=False) + keys: Final = envelope_keys_from_master_key(_MASTER_KEY) + assert (keys.signing_key.get_secret_value(), keys.encryption_key.get_secret_value()) == ( + "34d3bb6ae2e03cddca3417209304e9603417b89258d3979d0569880b6fbd1bed", + "eebfbee8cbe767ee136ce3a9745b44de2b2ae35d1ea46dcf72c58b811fac3250", + ) + + +@pytest.mark.usefixtures("fips_mode_on") +@pytest.mark.parametrize("salt_key", ["", "sk-salt-change-me", _MASTER_KEY, "sk-" + "5a" * 31]) +def test_envelope_keys_under_fips_refuse_a_salt_key_that_is_not_32_bytes_of_hex( + monkeypatch: pytest.MonkeyPatch, salt_key: str +) -> None: + monkeypatch.setenv("LITELLM_SALT_KEY", salt_key) + with pytest.raises(ValueError, match="LITELLM_SALT_KEY"): + envelope_keys_from_master_key(_MASTER_KEY) diff --git a/tests/unit/proxy/mcp_server/outbound_credentials/test_session_credentials.py b/tests/unit/proxy/mcp_server/outbound_credentials/test_session_credentials.py index f084d5c7d1d..795ded8528c 100644 --- a/tests/unit/proxy/mcp_server/outbound_credentials/test_session_credentials.py +++ b/tests/unit/proxy/mcp_server/outbound_credentials/test_session_credentials.py @@ -1,6 +1,8 @@ """Tests for the session-token KDF and the edge/token-endpoint resolvers.""" +import _hashlib from datetime import datetime, timedelta, timezone +from typing import Final import pytest from cryptography.hazmat.primitives import serialization @@ -230,3 +232,64 @@ def test_signing_config_error_detail_never_leaks_key_material(): ) assert isinstance(resolved, SessionSigningConfigError) assert pem.splitlines()[1] not in resolved.detail + + +def test_session_key_without_fips_stays_the_scrypt_key_of_the_master_key() -> None: + keys: Final = session_keys_from_master_key(MASTER_KEY) + assert keys.signing_key.get_secret_value() == "990dce331d97862e3af47815b707c42923fc760e6740d4a77a81583a206f339d" + + +@pytest.mark.usefixtures("openssl_enforcing_fips") +def test_session_key_on_a_fips_openssl_without_fips_mode_stays_the_scrypt_key(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("LITELLM_SALT_KEY", raising=False) + keys: Final = session_keys_from_master_key(MASTER_KEY) + assert keys.signing_key.get_secret_value() == "990dce331d97862e3af47815b707c42923fc760e6740d4a77a81583a206f339d" + + +def _fips_keys(master_key: str) -> tuple[str, str, str]: + envelope_keys_from_master_key.cache_clear() + session_keys_from_master_key.cache_clear() + envelope: Final = envelope_keys_from_master_key(master_key) + session: Final = session_keys_from_master_key(master_key) + return ( + envelope.signing_key.get_secret_value(), + envelope.encryption_key.get_secret_value(), + session.signing_key.get_secret_value(), + ) + + +_needs_shared_openssl: Final = pytest.mark.skipif( + not hasattr(_hashlib, "__file__"), reason="this Python links OpenSSL statically, out of ctypes reach" +) + + +@_needs_shared_openssl +@pytest.mark.usefixtures("fips_mode_on") +def test_keys_under_fips_come_from_the_salt_key_are_distinct_and_the_same_on_every_call( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-" + "5a" * 32) + first: Final = _fips_keys(MASTER_KEY) + assert first == _fips_keys(MASTER_KEY) + assert len(set(first)) == 3, first + assert first == ( + "ef694d415772aef09731c47709d4d9da47f7bbfcd6484836dea09c2bb387063b", + "d8a7cd70690fbf98c5581301b2e0460f2ad714f7c74b2af1b3dac3e0aa99db85", + "5e11e470bf9c838165b70d0b2b82a7ee5e2f00d9449eeb54f7c5bf9491b64354", + ) + + +@_needs_shared_openssl +@pytest.mark.usefixtures("fips_mode_on") +def test_keys_under_fips_all_change_when_the_master_key_rotates(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-" + "5a" * 32) + before: Final = _fips_keys(MASTER_KEY) + after: Final = _fips_keys("sk-a-different-master-key") + assert set(before).isdisjoint(after), (before, after) + + +@pytest.mark.usefixtures("fips_mode_on") +def test_session_key_under_fips_refuses_an_unset_salt_key(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("LITELLM_SALT_KEY", raising=False) + with pytest.raises(ValueError, match="LITELLM_SALT_KEY"): + session_keys_from_master_key(MASTER_KEY) diff --git a/tests/unit/proxy/middleware/test_billable_request_metrics_middleware.py b/tests/unit/proxy/middleware/test_billable_request_metrics_middleware.py index d2e961ad158..2488af7952c 100644 --- a/tests/unit/proxy/middleware/test_billable_request_metrics_middleware.py +++ b/tests/unit/proxy/middleware/test_billable_request_metrics_middleware.py @@ -105,6 +105,7 @@ def test_is_pure_asgi_not_base_http_middleware(): # SpendLogs-producing routes surfaced by the route-inventory audit ("/v1/search", (BillableCategory.LLM, "/search")), ("/v1/vector_stores/vs_1/search", (BillableCategory.LLM, "/search")), + ("/opensearch/docs/_search", (BillableCategory.LLM, "/opensearch")), ("/v1/rag/query", (BillableCategory.LLM, "/rag/query")), ("/rag/ingest", (BillableCategory.LLM, "/rag/ingest")), # Provider passthrough carries real inference and writes SpendLogs @@ -164,6 +165,18 @@ def test_classify_billable(path: str, expected: Tuple[BillableCategory, str]): "/v1beta/interactions/int_123/cancel", # observability passthrough writes no SpendLogs row "/langfuse/api/public/ingestion", + "/opensearch/docs/_doc", + "/opensearch/docs/_doc/_search", + "/opensearch/docs/_doc/1", + "/opensearch/docs/_bulk", + "/opensearch/docs/_update/1", + "/opensearch/docs/_count", + "/opensearch/docs/_mget", + "/opensearch/docs/_doc/search", + "/opensearch/docs/_update/search", + "/opensearch/docs/_create/embeddings", + "/opensearch/docs/_doc/completions", + "/opensearch/docs/_explain/rerank", # a bare provider prefix is not an inference call "/bedrock", "/eu_jev/v1/systemone/extra", diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 2fa3473d1c0..96d0a2ea8e0 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -1,11 +1,16 @@ import asyncio import base64 import contextlib +import hashlib +import itertools import json import logging import os +import re import traceback from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping +from datetime import datetime, timezone +from pathlib import Path from types import MappingProxyType, SimpleNamespace from typing import Final, Literal from unittest import mock @@ -15,6 +20,9 @@ from urllib.parse import parse_qs import httpx import pytest import respx +from botocore.auth import SigV4Auth +from botocore.awsrequest import AWSRequest +from botocore.credentials import Credentials from fastapi import FastAPI, HTTPException, Request, Response from fastapi.routing import APIRoute from fastapi.responses import StreamingResponse @@ -33,7 +41,7 @@ from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessin from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException, ModifyResponseException from litellm.proxy._types import AwsErrorTypeProxyException, ProxyException from litellm.proxy.spend_tracking.spend_tracking_utils import get_spend_logs_id -from litellm.proxy.utils import ProxyLogging +from litellm.proxy.utils import ProxyLogging, hash_token from litellm.passthrough.main import raise_for_provider_status from litellm.llms.bedrock.common_utils import BedrockError from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe @@ -2203,7 +2211,7 @@ class TestBedrockAgentRuntimePassthroughToggle: "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_secret_str", return_value="us-east-1", ), - patch("litellm.llms.bedrock.chat.BedrockConverseLLM", return_value=bedrock_llm), + patch("litellm.llms.bedrock.passthrough.transformation.BedrockAgentRuntimeAWS", return_value=bedrock_llm), patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_request_copy", Mock(), @@ -2492,6 +2500,67 @@ async def test_bedrock_agent_runtime_region_lookup_leaves_the_event_loop_free(mo assert retrieve.calls.last.request.headers["authorization"].startswith("AWS4-HMAC-SHA256") +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("endpoint", "action"), + [ + ("retrieveAndGenerate/", "bedrock:RetrieveAndGenerate"), + ("flows/FLOW123456/aliases/ALIAS12345", "bedrock:InvokeFlow"), + ("generateQuery/", "bedrock:GenerateQuery"), + ("optimize-prompt/", "bedrock:OptimizePrompt"), + ("agents/session-1", "bedrock:InvokeInlineAgent"), + ], +) +async def test_bedrock_agent_runtime_web_identity_session_grants_the_routes_action( + monkeypatch: pytest.MonkeyPatch, endpoint: str, action: str +) -> None: + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + for ambient in ( + *TestBedrockAgentRuntimePassthroughVirtualKeyLeak.AMBIENT_AWS_ENV, + "AWS_ACCESS_KEY_ID", + "AWS_SECRET_ACCESS_KEY", + ): + monkeypatch.delenv(ambient, raising=False) + monkeypatch.setenv("AWS_REGION_NAME", "us-east-1") + monkeypatch.setenv("AWS_WEB_IDENTITY_TOKEN", "oidc/google/litellm") + monkeypatch.setenv("AWS_ROLE_NAME", "arn:aws:iam::123456789012:role/litellm-role") + monkeypatch.setenv("AWS_SESSION_NAME", f"agent-runtime-{endpoint}") + sts: Final = MagicMock() + sts.assume_role_with_web_identity.return_value = { + "Credentials": { + "AccessKeyId": "k", + "SecretAccessKey": "s", + "SessionToken": "t", + "Expiration": datetime(2099, 1, 1, tzinfo=timezone.utc), + }, + "PackedPolicySize": 0, + } + request: Final = Mock() + request.method = "POST" + request.state = SimpleNamespace() + request.json = AsyncMock(return_value={"input": {"text": "hi"}}) + module: Final = "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints" + + with ( + patch("boto3.client", return_value=sts), + patch("litellm.llms.bedrock.base_aws_llm.get_secret", return_value="oidc-jwt-token"), + patch(f"{module}.create_request_copy", Mock()), + patch(f"{module}.create_pass_through_route", return_value=AsyncMock(return_value="forwarded")), + ): + _ = await bedrock_proxy_route( + endpoint=endpoint, + request=request, + fastapi_response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-litellm-caller"), + ) + + session_policy: Final = json.loads(sts.assume_role_with_web_identity.call_args.kwargs["Policy"]) + granted: Final = frozenset( + itertools.chain.from_iterable(statement["Action"] for statement in session_policy["Statement"]) + ) + assert action in granted, f"/bedrock/{endpoint} authorizes against {action}, which its session does not grant" + + class TestLLMPassthroughFactoryProxyRoute: @pytest.mark.asyncio async def test_llm_passthrough_factory_proxy_route_success(self): @@ -10106,3 +10175,914 @@ class TestTinyFishRouteTimeout: monkeypatch.setattr(proxy_server, "general_settings", {"pass_through_request_timeout": 30}, raising=False) assert _tinyfish_route_timeout() is None + + +OPENSEARCH_HOST: Final = "abc123.us-east-1.aoss.amazonaws.com" +OPENSEARCH_ACCESS_KEY: Final = "AKIDOPENSEARCHPASSTHROUGH" +OPENSEARCH_SECRET_KEY: Final = "opensearch-passthrough-secret" +DOCS_UPSTREAM: Final = f"https://{OPENSEARCH_HOST}/prod-docs-v2" +JSON_TYPE: Final = "application/json" +NDJSON_TYPE: Final = "application/x-ndjson" +_SIGV4_AUTHORIZATION: Final = re.compile( + r"\AAWS4-HMAC-SHA256 Credential=(?P[^/]+)/\d{8}/(?P[^/]+)/(?P[^/]+)/aws4_request, " + r"SignedHeaders=(?P[^,]+), Signature=(?P[0-9a-f]{64})\Z" +) + + +def _verified_sigv4_scope(sent: httpx.Request) -> tuple[str, str]: + authorization: Final = _SIGV4_AUTHORIZATION.match(sent.headers["authorization"]) + assert authorization is not None, sent.headers["authorization"] + signed_names: Final = authorization["signed"].split(";") + assert {"content-type", "host", "x-amz-content-sha256", "x-amz-date"} <= set(signed_names), signed_names + assert sent.headers["x-amz-content-sha256"] == hashlib.sha256(sent.content).hexdigest() + received: Final = AWSRequest( + method=sent.method, + url=str(sent.url), + data=sent.content, + headers={name: sent.headers[name] for name in signed_names}, + ) + received.context["timestamp"] = sent.headers["x-amz-date"] + signer: Final = SigV4Auth( + Credentials(OPENSEARCH_ACCESS_KEY, OPENSEARCH_SECRET_KEY), authorization["service"], authorization["region"] + ) + expected: Final = signer.signature(signer.string_to_sign(received, signer.canonical_request(received)), received) + assert (authorization["key"], authorization["signature"]) == (OPENSEARCH_ACCESS_KEY, expected) + return authorization["service"], authorization["region"] + + +def _opensearch_caller(*permissions: str, index_name: str = "docs") -> UserAPIKeyAuth: + grant: Final = { + "allowed_vector_store_indexes": [{"index_name": index_name, "index_permissions": list(permissions)}] + } + return UserAPIKeyAuth(api_key="sk-virtual", metadata=grant) + + +def _act_as(monkeypatch: pytest.MonkeyPatch, caller: UserAPIKeyAuth) -> None: + from litellm.proxy.proxy_server import app + + monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, lambda: caller) + + +@pytest.fixture +def opensearch_client(monkeypatch: pytest.MonkeyPatch) -> Iterator[TestClient]: + from litellm.proxy.proxy_server import app + from litellm.types.vector_stores import IndexCreateLiteLLMParams, LiteLLM_ManagedVectorStoreIndex + from litellm.vector_stores.vector_store_registry import VectorStoreIndexRegistry, VectorStoreRegistry + + monkeypatch.delenv("SERVER_ROOT_PATH", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + indexes: Final = ( + ("docs", "prod-docs-v2", "opensearch-store"), + ("hr", "hr-private", "opensearch-store"), + ("no-credentials", "prod-docs-v3", "opensearch-store-without-credentials"), + ("azure-docs", "azure-docs", "azure-store"), + ) + monkeypatch.setattr( + litellm, + "vector_store_index_registry", + VectorStoreIndexRegistry( + [ + LiteLLM_ManagedVectorStoreIndex( + id=f"index-{name}", + index_name=name, + litellm_params=IndexCreateLiteLLMParams(vector_store_index=provider_index, vector_store_name=store), + ) + for name, provider_index, store in indexes + ] + ), + ) + monkeypatch.setattr( + litellm, + "vector_store_registry", + VectorStoreRegistry( + [ + { + "vector_store_id": "opensearch-store", + "vector_store_name": "opensearch-store", + "custom_llm_provider": "opensearch", + "litellm_params": { + "api_base": f"https://{OPENSEARCH_HOST}", + "aws_access_key_id": OPENSEARCH_ACCESS_KEY, + "aws_secret_access_key": OPENSEARCH_SECRET_KEY, + }, + }, + { + "vector_store_id": "opensearch-store-without-credentials", + "vector_store_name": "opensearch-store-without-credentials", + "custom_llm_provider": "opensearch", + "litellm_params": { + "api_base": f"https://{OPENSEARCH_HOST}", + "aws_profile_name": "opensearch-missing-profile", + }, + }, + { + "vector_store_id": "azure-store", + "vector_store_name": "azure-store", + "custom_llm_provider": "azure_ai", + "litellm_params": {"api_base": "https://search.example.search.windows.net"}, + }, + ] + ), + ) + yield TestClient(app) + + +class TestOpenSearchProxyRoute: + SEARCH_BODY: Final = b'{"query":{"match":{"text":"remote work"}}}' + HITS: Final = MappingProxyType( + { + "took": 3, + "timed_out": False, + "hits": { + "total": {"value": 1, "relation": "eq"}, + "hits": [{"_index": "prod-docs-v2", "_id": "1", "_score": 1.2, "_source": {"text": "Remote work"}}], + }, + } + ) + + @pytest.mark.parametrize("content_type", [JSON_TYPE, "Application/JSON; charset=UTF-8"]) + def test_a_read_key_search_is_signed_for_the_mapped_index_and_returns_the_hits( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, content_type: str + ) -> None: + _act_as(monkeypatch, _opensearch_caller("read")) + with respx.mock(assert_all_called=True) as upstream: + route = upstream.post(f"{DOCS_UPSTREAM}/_search").mock( + return_value=httpx.Response(200, json=dict(self.HITS)) + ) + response = opensearch_client.post( + "/opensearch/docs/_search", + content=self.SEARCH_BODY, + headers={"Content-Type": content_type, "Authorization": "Bearer sk-virtual"}, + ) + + assert (response.status_code, response.json()) == (200, dict(self.HITS)) + sent: Final = route.calls.last.request + assert (sent.content, sent.headers["content-type"]) == (self.SEARCH_BODY, JSON_TYPE) + assert _verified_sigv4_scope(sent) == ("aoss", "us-east-1") + + def test_a_team_grant_lets_the_key_search( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch + ) -> None: + team_grant: Final = {"allowed_vector_store_indexes": [{"index_name": "docs", "index_permissions": ["read"]}]} + _act_as(monkeypatch, UserAPIKeyAuth(api_key="sk-virtual", team_id="team-a", team_metadata=team_grant)) + with respx.mock(assert_all_called=True) as upstream: + upstream.post(f"{DOCS_UPSTREAM}/_search").mock(return_value=httpx.Response(200, json=dict(self.HITS))) + response = opensearch_client.post( + "/opensearch/docs/_search", content=self.SEARCH_BODY, headers={"Content-Type": JSON_TYPE} + ) + + assert (response.status_code, response.json()) == (200, dict(self.HITS)) + + @pytest.mark.parametrize( + "method, api", + [ + ("POST", "_doc"), + ("PUT", "_doc/1"), + ("POST", "_doc/1"), + ("DELETE", "_doc/1"), + ("PUT", "_create/1"), + ("POST", "_update/1"), + ("POST", "_bulk"), + ("POST", "_delete_by_query"), + ("POST", "_update_by_query"), + ], + ) + def test_a_read_only_key_cannot_write( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, method: str, api: str + ) -> None: + _act_as(monkeypatch, _opensearch_caller("read")) + with respx.mock(assert_all_called=False) as upstream: + route = upstream.route(host=OPENSEARCH_HOST).mock(return_value=httpx.Response(201, json={})) + response = opensearch_client.request( + method, + f"/opensearch/docs/{api}", + content=b'{"title":"Remote work"}', + headers={"Content-Type": JSON_TYPE}, + ) + + assert (response.status_code, route.called) == (403, False) + + @pytest.mark.parametrize( + "document", + [b'{"title":"Remote work","index":"hr-private"}', b'{"title":"Remote work","stream":true}'], + ids=["field-named-index", "field-named-stream"], + ) + def test_a_read_write_key_writes_the_document_as_sent( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, document: bytes + ) -> None: + _act_as(monkeypatch, _opensearch_caller("read", "write")) + with respx.mock(assert_all_called=True) as upstream: + route = upstream.post(f"{DOCS_UPSTREAM}/_doc").mock( + return_value=httpx.Response(201, json={"_id": "abc", "result": "created"}) + ) + response = opensearch_client.post( + "/opensearch/docs/_doc", content=document, headers={"Content-Type": JSON_TYPE} + ) + + assert (response.status_code, response.json()["result"]) == (201, "created") + assert route.calls.last.request.content == document + assert _verified_sigv4_scope(route.calls.last.request) == ("aoss", "us-east-1") + + @pytest.mark.parametrize( + "path, body, status_code", + [ + pytest.param("_update/1", b'{"doc":{},"_source":true}', 403, id="update-body-source"), + pytest.param("_update/1?_source=text", b'{"doc":{"title":"b"}}', 403, id="update-source-param"), + pytest.param( + "_update/1", + b'{"script":{"source":"throw new IllegalArgumentException(ctx._source.text)"}}', + 403, + id="update-script", + ), + pytest.param("_update/1", b'{"doc":{} /* c */ ,"_source":true}', 400, id="update-body-json-with-comment"), + pytest.param("_bulk", b'{"update":{"_id":"1","_source":true}}\n{"doc":{}}\n', 403, id="bulk-action-source"), + pytest.param( + "_bulk", + b'{"index":{"_id":"2"}}\n{"title":"a"}\n{"update":{"_id":"1"}}\n{"doc":{},"_source":true}\n', + 403, + id="bulk-update-payload-source", + ), + pytest.param("_bulk?_source_includes=text", b'{"update":{"_id":"1"}}\n{"doc":{}}\n', 403, id="bulk-param"), + pytest.param("_delete_by_query", b'{"query":{"prefix":{"text":"a"}}}', 403, id="delete-by-query"), + pytest.param("_update_by_query", b'{"query":{"match_all":{}}}', 403, id="update-by-query"), + ], + ) + def test_a_write_only_key_cannot_read_documents_back_through_a_write( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, path: str, body: bytes, status_code: int + ) -> None: + _act_as(monkeypatch, _opensearch_caller("write")) + with respx.mock(assert_all_called=False) as upstream: + route = upstream.route(host=OPENSEARCH_HOST).mock( + return_value=httpx.Response(200, json={"get": {"_source": {"text": "Remote work"}}}) + ) + response = opensearch_client.post( + f"/opensearch/docs/{path}", content=body, headers={"Content-Type": NDJSON_TYPE} + ) + + assert (response.status_code, route.called) == (status_code, False) + assert "Remote work" not in response.text + + @pytest.mark.parametrize( + "caller, path, body", + [ + pytest.param( + _opensearch_caller("write"), "_update/1", b'{"doc":{"script":"a field","_source":"b"}}', id="write-only" + ), + pytest.param( + _opensearch_caller("read", "write"), "_update/1?_source=text", b'{"doc":{},"_source":true}', id="update" + ), + pytest.param( + _opensearch_caller("read", "write"), + "_bulk", + b'{"update":{"_id":"1","_source":true}}\n{"doc":{},"script":{"source":"ctx._source.n += 1"}}\n', + id="bulk", + ), + pytest.param( + _opensearch_caller("read", "write"), "_delete_by_query", b'{"query":{"term":{"a":1}}}', id="by-query" + ), + ], + ) + def test_a_write_that_reads_no_document_or_has_the_read_grant_is_sent_as_is( + self, + opensearch_client: TestClient, + monkeypatch: pytest.MonkeyPatch, + caller: UserAPIKeyAuth, + path: str, + body: bytes, + ) -> None: + _act_as(monkeypatch, caller) + with respx.mock(assert_all_called=True) as upstream: + route = upstream.post(url__startswith=f"{DOCS_UPSTREAM}/{path.partition('?')[0]}").mock( + return_value=httpx.Response(200, json={"result": "updated"}) + ) + response = opensearch_client.post( + f"/opensearch/docs/{path}", content=body, headers={"Content-Type": NDJSON_TYPE} + ) + + assert (response.status_code, route.calls.last.request.content) == (200, body) + + @pytest.mark.parametrize( + "caller, index_name", + [ + (_opensearch_caller("read", "write", index_name="hr"), "docs"), + (UserAPIKeyAuth(api_key="sk-virtual"), "docs"), + (UserAPIKeyAuth(api_key="sk-virtual"), "unregistered"), + ], + ids=["grant-on-another-index", "no-grant", "no-grant-unregistered-index"], + ) + def test_a_key_without_a_grant_on_the_index_is_refused( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, caller: UserAPIKeyAuth, index_name: str + ) -> None: + _act_as(monkeypatch, caller) + with respx.mock(assert_all_called=False) as upstream: + route = upstream.route(host=OPENSEARCH_HOST).mock(return_value=httpx.Response(200, json={})) + response = opensearch_client.post( + f"/opensearch/{index_name}/_search", content=self.SEARCH_BODY, headers={"Content-Type": JSON_TYPE} + ) + + assert (response.status_code, route.called) == (403, False) + assert index_name in response.json()["detail"] + + @pytest.mark.parametrize( + "method, path", + [ + ("POST", "/opensearch/docs/_search/template"), + ("POST", "/opensearch/docs/_msearch"), + ("GET", "/opensearch/docs/_msearch"), + ("PUT", "/opensearch/docs/_mapping"), + ("GET", "/opensearch/docs/_settings"), + ("POST", "/opensearch/docs/_close"), + ("POST", "/opensearch/docs/_refresh"), + ("GET", "/opensearch/docs/_doc/1/x"), + ("PATCH", "/opensearch/docs/_doc/1"), + ("PUT", "/opensearch/docs"), + ("DELETE", "/opensearch/docs"), + ], + ) + def test_an_unlisted_api_is_for_proxy_admins_only( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, method: str, path: str + ) -> None: + _act_as(monkeypatch, _opensearch_caller("read", "write")) + with respx.mock(assert_all_called=False) as upstream: + route = upstream.route(host=OPENSEARCH_HOST).mock(return_value=httpx.Response(200, json={})) + response = opensearch_client.request( + method, path, content=b'{"source":"{}"}', headers={"Content-Type": JSON_TYPE} + ) + + assert (response.status_code, route.called) == (403, False) + + @pytest.mark.parametrize( + "method, path, upstream_path, status", + [ + pytest.param("GET", "/opensearch/docs", "/prod-docs-v2", 200, id="get-index"), + pytest.param("HEAD", "/opensearch/docs", "/prod-docs-v2", 200, id="index-exists"), + pytest.param("HEAD", "/opensearch/docs/_doc/1", "/prod-docs-v2/_doc/1", 404, id="doc-exists"), + pytest.param("HEAD", "/opensearch/docs/_source/1", "/prod-docs-v2/_source/1", 200, id="source-exists"), + pytest.param( + "POST", "/opensearch/docs/_explain/_bulk1", "/prod-docs-v2/_explain/_bulk1", 200, id="id-named-bulk" + ), + ], + ) + def test_a_read_key_reaches_the_index_and_its_existence_checks( + self, + opensearch_client: TestClient, + monkeypatch: pytest.MonkeyPatch, + method: str, + path: str, + upstream_path: str, + status: int, + ) -> None: + _act_as(monkeypatch, _opensearch_caller("read")) + with respx.mock(assert_all_called=True) as upstream: + route = upstream.route(method=method, host=OPENSEARCH_HOST, path=upstream_path).mock( + return_value=httpx.Response(status, content=b"" if method == "HEAD" else b"{}") + ) + response = opensearch_client.request( + method, path, content=b"" if method == "HEAD" else self.SEARCH_BODY, headers={"Content-Type": JSON_TYPE} + ) + + assert response.status_code == status + assert _verified_sigv4_scope(route.calls.last.request) == ("aoss", "us-east-1") + + def test_a_proxy_admin_can_call_an_unlisted_api( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch + ) -> None: + settings: Final = {"prod-docs-v2": {"settings": {"index": {"number_of_shards": "2"}}}} + _act_as(monkeypatch, UserAPIKeyAuth(api_key="sk-admin", user_role=LitellmUserRoles.PROXY_ADMIN)) + with respx.mock(assert_all_called=True) as upstream: + route = upstream.get(f"{DOCS_UPSTREAM}/_settings").mock(return_value=httpx.Response(200, json=settings)) + response = opensearch_client.get("/opensearch/docs/_settings") + + assert (response.status_code, response.json()) == (200, settings) + assert _verified_sigv4_scope(route.calls.last.request) == ("aoss", "us-east-1") + + @pytest.mark.parametrize( + "index_name, detail", + [ + ("unknown", "Index unknown is not a litellm managed vector store index"), + ("azure-docs", "not an OpenSearch vector store"), + ], + ) + def test_an_index_that_is_not_a_managed_opensearch_index_is_a_400( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, index_name: str, detail: str + ) -> None: + _act_as(monkeypatch, _opensearch_caller("read", index_name=index_name)) + with respx.mock(assert_all_called=False) as upstream: + route = upstream.route().mock(return_value=httpx.Response(200, json={})) + response = opensearch_client.post( + f"/opensearch/{index_name}/_search", content=self.SEARCH_BODY, headers={"Content-Type": JSON_TYPE} + ) + + assert (response.status_code, route.called) == (400, False) + assert detail in response.json()["detail"] + + def test_an_index_whose_store_name_two_stores_share_is_a_400( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch + ) -> None: + from litellm.vector_stores.vector_store_registry import VectorStoreRegistry + + same_name: Final = ( + ("opensearch-store", f"https://{OPENSEARCH_HOST}"), + ("opensearch-store-copy", "https://other123.us-east-1.aoss.amazonaws.com"), + ) + monkeypatch.setattr( + litellm, + "vector_store_registry", + VectorStoreRegistry( + [ + { + "vector_store_id": store_id, + "vector_store_name": "opensearch-store", + "custom_llm_provider": "opensearch", + "litellm_params": {"api_base": api_base}, + } + for store_id, api_base in same_name + ] + ), + ) + _act_as(monkeypatch, _opensearch_caller("read")) + with respx.mock(assert_all_called=False) as upstream: + route = upstream.route().mock(return_value=httpx.Response(200, json={})) + response = opensearch_client.post( + "/opensearch/docs/_search", content=self.SEARCH_BODY, headers={"Content-Type": JSON_TYPE} + ) + + assert (response.status_code, route.called) == (400, False) + assert "2 vector stores are named opensearch-store" in response.json()["detail"] + + @pytest.mark.parametrize( + "path", + [ + "/opensearch/docs/", + "/opensearch/docs//_search", + "/opensearch/docs/%2E/_search", + "/opensearch/docs/%2E%2E/hr/_search", + "/opensearch/docs/_search/%2E%2E/%2E%2E/hr/_search", + ], + ) + def test_a_path_with_an_empty_or_dot_segment_is_a_400( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, path: str + ) -> None: + _act_as(monkeypatch, _opensearch_caller("read", "write")) + with respx.mock(assert_all_called=False) as upstream: + route = upstream.route(host=OPENSEARCH_HOST).mock(return_value=httpx.Response(200, json={})) + response = opensearch_client.post(path, content=self.SEARCH_BODY, headers={"Content-Type": JSON_TYPE}) + + assert (response.status_code, route.called) == (400, False) + + @pytest.mark.parametrize( + "api, content_type, body, detail", + [ + pytest.param( + "_search", + JSON_TYPE, + b'{"query":{"terms":{"user":{"index":"hr-private","id":"1","path":"followers"}}}}', + "'index'", + id="search-terms-lookup", + ), + pytest.param( + "_search", + JSON_TYPE, + b'{"query":{"more_like_this":{"like":[{"_index":"hr-private","_id":"1"}]}}}', + "'_index'", + id="search-more-like-this", + ), + pytest.param( + "_count", + JSON_TYPE, + b'{"query":{"terms":{"user":{"\\u0069ndex":"hr-private","id":"1","path":"followers"}}}}', + "'index'", + id="count-escaped-key", + ), + pytest.param( + "_search", + JSON_TYPE, + b'{"query":{"terms":{"user":{"index":"hr-private","id":"1","path":"f"}}},"query":{"match_all":{}}}', + "'index'", + id="search-duplicate-key", + ), + pytest.param( + "_delete_by_query", + JSON_TYPE, + b'{"query":{"terms":{"user":{"index":"hr-private","id":"1","path":"followers"}}}}', + "'index'", + id="delete-by-query-terms-lookup", + ), + pytest.param( + "_mget", + JSON_TYPE, + b'{"docs":[{"_id":"1"},{"_index":"hr-private","_id":"2"}]}', + "'_index'", + id="mget-docs", + ), + pytest.param( + "_bulk", + NDJSON_TYPE, + b'{"index":{"_id":"1"}}\n{"title":"a"}\n{"delete":{"_index":"hr-private","_id":"2"}}\n', + "'_index'", + id="bulk-action-index", + ), + pytest.param( + "_bulk", + NDJSON_TYPE, + b'{"index":{"_index":"prod-docs-v2","_id":"1"}}\n{"title":"a"}\n', + "'_index'", + id="bulk-action-names-the-provider-index", + ), + pytest.param( + "_bulk", + NDJSON_TYPE, + b'{"index":{"_index":"docs","_index":"hr-private","_id":"1"}}\n{"title":"a"}\n', + "'_index'", + id="bulk-action-duplicate-index", + ), + pytest.param( + "_bulk", + NDJSON_TYPE, + b'{"index":{"_id":"1"}}\n\n{"delete":{"_index":"hr-private","_id":"2"}}\n', + "'_index'", + id="bulk-blank-source-line", + ), + pytest.param( + "_bulk", + NDJSON_TYPE, + b'{"index":{"_id":"1","pipeline":"reroute"}}\n{"title":"a"}\n', + "'pipeline'", + id="bulk-action-pipeline", + ), + pytest.param("_bulk", NDJSON_TYPE, b"not json\n", "action line", id="bulk-unparseable-action"), + pytest.param( + "_search", + JSON_TYPE, + b'{"profile":true,"query":{"wrapper":{"query":"eyJ0ZXJtcyI6e319"}}}', + "'wrapper'", + id="search-wrapper-query", + ), + pytest.param( + "_search", + JSON_TYPE, + b'{"query":{"geo_shape":{"location":{"indexed_shape":{"id":"deu","path":"location"}}}}}', + "'indexed_shape'", + id="search-indexed-shape", + ), + pytest.param( + "_search", + JSON_TYPE, + b'{"query":{"match_all":{}},"search_pipeline":{"request_processors":[]}}', + "'search_pipeline'", + id="search-temporary-pipeline", + ), + pytest.param( + "_mget?source=%7B%22docs%22%3A%5B%7B%22_index%22%3A%22hr-private%22%2C%22_id%22%3A%221%22%7D%5D%7D" + "&source_content_type=application%2Fjson", + JSON_TYPE, + b"", + "source query parameter", + id="mget-source-query", + ), + pytest.param( + "_search", + "application/x-ndjson; multipart/form-data", + b'{"stream":true,"custom_body":{}}', + "application/json", + id="multipart-parameter", + ), + pytest.param( + "_doc", "application/json; boundary=x", b'{"title":"a"}', "application/json", id="content-type-parameter" + ), + pytest.param("_doc?pipeline=reroute", JSON_TYPE, b'{"title":"a"}', "pipeline", id="pipeline-query"), + pytest.param("_search", JSON_TYPE, b'{/*c*/"query":{"match_all":{}}}', "could not parse", id="comment"), + pytest.param("_search", "application/smile", b"\x3a\x29\x0a\x01", "application/json", id="smile"), + pytest.param("_search", "application/yaml", b"query:\n match_all: {}\n", "application/json", id="yaml"), + pytest.param( + "_search", JSON_TYPE, b'{"query":' + b"[" * 100_000 + b"]" * 100_000 + b"}", "nested", id="deep-nesting" + ), + ], + ) + def test_a_body_or_query_that_could_reach_another_index_is_a_400( + self, + opensearch_client: TestClient, + monkeypatch: pytest.MonkeyPatch, + api: str, + content_type: str, + body: bytes, + detail: str, + ) -> None: + _act_as(monkeypatch, _opensearch_caller("read", "write")) + with respx.mock(assert_all_called=False) as upstream: + route = upstream.route(host=OPENSEARCH_HOST).mock(return_value=httpx.Response(200, json={})) + response = opensearch_client.post( + f"/opensearch/docs/{api}", content=body, headers={"Content-Type": content_type} + ) + + assert (response.status_code, route.called) == (400, False) + assert detail in response.json()["detail"] + + @pytest.mark.parametrize( + "api, body", + [ + pytest.param( + "_bulk", + b'{"index":{"_id":"1"}}\n{"title":"a","index":"hr-private"}\n{"delete":{"_id":"2"}}\n\n' + b'{"update":{"_id":"3"}}\r\n{"doc":{"_index":"hr-private"}}\n', + id="bulk", + ), + ], + ) + @pytest.mark.parametrize("content_type", [NDJSON_TYPE, "application/json"]) + def test_an_ndjson_body_without_index_references_is_sent_byte_for_byte( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, api: str, body: bytes, content_type: str + ) -> None: + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", SHARED_MASTER_KEY) + with respx.mock(assert_all_called=True) as upstream: + route = upstream.post(f"{DOCS_UPSTREAM}/{api}").mock( + return_value=httpx.Response(200, json={"errors": False, "items": []}) + ) + response = opensearch_client.post( + f"/opensearch/docs/{api}", + content=body, + headers={"Content-Type": content_type, "Authorization": f"Bearer {SHARED_MASTER_KEY}"}, + ) + + assert response.status_code == 200 + sent: Final = route.calls.last.request + assert (sent.content, sent.headers["content-type"]) == (body, content_type) + assert _verified_sigv4_scope(sent) == ("aoss", "us-east-1") + + @pytest.mark.parametrize( + "document", + [ + b'{"title":"Remote work","model":"another-model"}', + b'{"title":"Remote work","api_base":"https://elsewhere.example"}', + b'{"title":"Remote work","base_url":"https://elsewhere.example"}', + ], + ids=["model", "api_base", "base_url"], + ) + def test_a_document_field_named_like_an_llm_param_is_written_for_a_model_limited_key( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, document: bytes + ) -> None: + from litellm.proxy import proxy_server + + key: Final = "sk-opensearch-model-limited-writer" + grant: Final = {"allowed_vector_store_indexes": [{"index_name": "docs", "index_permissions": ["write"]}]} + cache: Final = DualCache() + cache.set_cache( + key=hash_token(key), value=UserAPIKeyAuth(token=hash_token(key), models=["gpt-x"], metadata=grant) + ) + monkeypatch.setattr(proxy_server, "master_key", SHARED_MASTER_KEY) + monkeypatch.setattr(proxy_server, "prisma_client", "test-database") + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + with respx.mock(assert_all_called=True) as upstream: + route = upstream.post(f"{DOCS_UPSTREAM}/_doc").mock( + return_value=httpx.Response(201, json={"_id": "abc", "result": "created"}) + ) + response = opensearch_client.post( + "/opensearch/docs/_doc", + content=document, + headers={"Content-Type": JSON_TYPE, "Authorization": f"Bearer {key}"}, + ) + + assert (response.status_code, route.calls.last.request.content) == (201, document) + + def test_a_bulk_action_naming_the_path_index_is_sent_to_the_mapped_index( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch + ) -> None: + _act_as(monkeypatch, _opensearch_caller("write")) + with respx.mock(assert_all_called=True) as upstream: + route = upstream.post(f"{DOCS_UPSTREAM}/_bulk").mock( + return_value=httpx.Response(200, json={"errors": False, "items": []}) + ) + response = opensearch_client.post( + "/opensearch/docs/_bulk", + content=b'{"index":{"_index":"docs","_id":"1"}}\n{"title":"a","_index":"docs"}\n' + b'{"delete":{"_id":"2"}}\n{"update":{"_id":"3","_index":"docs"}}\n{"doc":{"title":"b"}}\n', + headers={"Content-Type": NDJSON_TYPE}, + ) + + assert response.status_code == 200 + sent: Final = route.calls.last.request + assert sent.content == ( + b'{"index":{"_index":"prod-docs-v2","_id":"1"}}\n{"title":"a","_index":"docs"}\n' + b'{"delete":{"_id":"2"}}\n{"update":{"_id":"3","_index":"prod-docs-v2"}}\n{"doc":{"title":"b"}}\n' + ) + assert _verified_sigv4_scope(sent) == ("aoss", "us-east-1") + + @pytest.mark.parametrize( + "method, path, body, sent_query", + [ + pytest.param( + "GET", + "/opensearch/docs/_search?q=title:remote+work&routing=a&routing=b", + b"", + b"q=title%3Aremote%20work&routing=a&routing=b", + id="space-and-repeated-key", + ), + pytest.param( + "POST", + "/opensearch/docs/_search", + b'{"size":1,"query_params":{"routing":"z"}}', + b"", + id="body-query-params-ignored", + ), + pytest.param( + "POST", + "/opensearch/docs/_search?search_pipeline=hybrid-min-max", + b'{"query":{"hybrid":{"queries":[{"match":{"title":"remote"}},' + b'{"knn":{"embedding":{"vector":[0.1,0.2],"k":2,"filter":{"term":{"lang":"en"}}}}}]}}}', + b"search_pipeline=hybrid-min-max", + id="hybrid-query-with-named-search-pipeline", + ), + ], + ) + def test_the_query_string_is_signed_as_sent( + self, + opensearch_client: TestClient, + monkeypatch: pytest.MonkeyPatch, + method: str, + path: str, + body: bytes, + sent_query: bytes, + ) -> None: + _act_as(monkeypatch, _opensearch_caller("read")) + with respx.mock(assert_all_called=True) as upstream: + route = upstream.route(method=method, host=OPENSEARCH_HOST, path="/prod-docs-v2/_search").mock( + return_value=httpx.Response(200, json=dict(self.HITS)) + ) + response = opensearch_client.request(method, path, content=body, headers={"Content-Type": JSON_TYPE}) + + assert response.status_code == 200 + sent: Final = route.calls.last.request + assert sent.url.query == sent_query + assert _verified_sigv4_scope(sent) == ("aoss", "us-east-1") + + def test_document_fields_named_like_request_params_leave_the_rate_limiter_alone( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch + ) -> None: + from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 + from litellm.proxy.utils import InternalUsageCache + + limiter: Final = PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(DualCache())) + + class RateLimit(CustomLogger): + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], + call_type: CallTypesLiteral, + ) -> dict[str, object]: + await limiter.async_pre_call_hook(user_api_key_dict, cache, data, call_type) + return data + + monkeypatch.setattr(litellm, "callbacks", [RateLimit()]) + grant: Final = {"index_name": "docs", "index_permissions": ["write"]} + _act_as( + monkeypatch, + UserAPIKeyAuth( + api_key="sk-opensearch-rate-limited-writer", + metadata={"allowed_vector_store_indexes": [grant], "model_rpm_limit": {"text-embedding-3-small": 1}}, + ), + ) + documents: Final = ( + b'{"text":"a","metadata":"handbook"}', + b'{"text":"b","metadata":["handbook"]}', + b'{"text":"c","model":"text-embedding-3-small"}', + b'{"text":"d","model":"text-embedding-3-small"}', + ) + with respx.mock(assert_all_called=True) as upstream: + route = upstream.post(f"{DOCS_UPSTREAM}/_doc").mock( + return_value=httpx.Response(201, json={"result": "created"}) + ) + statuses: Final = tuple( + opensearch_client.post( + "/opensearch/docs/_doc", content=document, headers={"Content-Type": JSON_TYPE} + ).status_code + for document in documents + ) + + assert statuses == (201, 201, 201, 201) + assert tuple(call.request.content for call in route.calls) == documents + + @pytest.mark.parametrize( + "document", [b'[{"text":"array body"}]', b'{"text":"\xff\xfe"}'], ids=["json-array", "invalid-utf-8"] + ) + def test_a_malformed_document_gets_the_opensearch_answer( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, document: bytes + ) -> None: + rejected: Final = {"error": {"type": "mapper_parsing_exception"}, "status": 400} + _act_as(monkeypatch, _opensearch_caller("write")) + with respx.mock(assert_all_called=True) as upstream: + route = upstream.post(f"{DOCS_UPSTREAM}/_doc").mock(return_value=httpx.Response(400, json=rejected)) + response = opensearch_client.post( + "/opensearch/docs/_doc", content=document, headers={"Content-Type": JSON_TYPE} + ) + + assert (response.status_code, response.json()) == (400, rejected) + assert route.calls.last.request.content == document + + def test_an_opensearch_error_keeps_its_status_and_body( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch + ) -> None: + not_found: Final = {"error": {"type": "index_not_found_exception", "index": "prod-docs-v2"}, "status": 404} + _act_as(monkeypatch, _opensearch_caller("read")) + with respx.mock(assert_all_called=True) as upstream: + upstream.post(f"{DOCS_UPSTREAM}/_search").mock(return_value=httpx.Response(404, json=not_found)) + response = opensearch_client.post( + "/opensearch/docs/_search", content=self.SEARCH_BODY, headers={"Content-Type": JSON_TYPE} + ) + + assert (response.status_code, response.json()) == (404, not_found) + + def test_missing_aws_credentials_are_a_401_before_anything_is_sent( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + for name in ("AWS_CONFIG_FILE", "AWS_SHARED_CREDENTIALS_FILE"): + empty_file: Path = tmp_path / name + empty_file.write_text("") + monkeypatch.setenv(name, str(empty_file)) + _act_as(monkeypatch, _opensearch_caller("read", index_name="no-credentials")) + with respx.mock(assert_all_called=False) as upstream: + route = upstream.route(host=OPENSEARCH_HOST).mock(return_value=httpx.Response(200, json={})) + response = opensearch_client.post( + "/opensearch/no-credentials/_search", content=self.SEARCH_BODY, headers={"Content-Type": JSON_TYPE} + ) + + assert (response.status_code, route.called) == (401, False) + + @pytest.mark.parametrize( + "method, api, action", + [ + ("GET", "_doc/1", "es:ESHttpGet"), + ("HEAD", "_doc/1", "es:ESHttpHead"), + ("PUT", "_doc/1", "es:ESHttpPut"), + ("DELETE", "_doc/1", "es:ESHttpDelete"), + ("POST", "_search", "es:ESHttpPost"), + ], + ) + def test_a_web_identity_session_allows_the_verb_sent_to_a_managed_domain( + self, opensearch_client: TestClient, monkeypatch: pytest.MonkeyPatch, method: str, api: str, action: str + ) -> None: + from litellm.types.vector_stores import IndexCreateLiteLLMParams, LiteLLM_ManagedVectorStoreIndex + from litellm.vector_stores.vector_store_registry import VectorStoreIndexRegistry, VectorStoreRegistry + + domain: Final = "search-docs-abc123.us-east-1.es.amazonaws.com" + monkeypatch.setattr( + litellm, + "vector_store_index_registry", + VectorStoreIndexRegistry( + [ + LiteLLM_ManagedVectorStoreIndex( + id="index-domain-docs", + index_name="docs", + litellm_params=IndexCreateLiteLLMParams( + vector_store_index="prod-docs-v2", vector_store_name="opensearch-domain" + ), + ) + ] + ), + ) + monkeypatch.setattr( + litellm, + "vector_store_registry", + VectorStoreRegistry( + [ + { + "vector_store_id": "opensearch-domain", + "vector_store_name": "opensearch-domain", + "custom_llm_provider": "opensearch", + "litellm_params": { + "api_base": f"https://{domain}", + "aws_web_identity_token": "oidc/google/litellm", + "aws_role_name": "arn:aws:iam::123456789012:role/litellm-role", + "aws_session_name": f"opensearch-index-{method}", + "aws_region_name": "us-east-1", + }, + } + ] + ), + ) + sts: Final = MagicMock() + sts.assume_role_with_web_identity.return_value = { + "Credentials": { + "AccessKeyId": "k", + "SecretAccessKey": "s", + "SessionToken": "t", + "Expiration": datetime(2099, 1, 1, tzinfo=timezone.utc), + }, + "PackedPolicySize": 0, + } + _act_as(monkeypatch, _opensearch_caller("read", "write")) + with ( + patch("boto3.client", return_value=sts), + patch("litellm.llms.bedrock.base_aws_llm.get_secret", return_value="oidc-jwt-token"), + respx.mock(assert_all_called=True) as upstream, + ): + upstream.route(method=method, host=domain).mock(return_value=httpx.Response(200, json={})) + response = opensearch_client.request( + method, f"/opensearch/docs/{api}", content=b"{}", headers={"Content-Type": JSON_TYPE} + ) + + policy: Final = json.loads(sts.assume_role_with_web_identity.call_args.kwargs["Policy"]) + granted: Final = frozenset(action for statement in policy["Statement"] for action in statement["Action"]) + assert (response.status_code, action in granted) == (200, True) diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 45abb880f92..464c39d0ba4 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -44,6 +44,8 @@ from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitErro from litellm.proxy.hooks.parallel_request_limiter import PROXY_MaxParallelRequestsHandler from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup +from litellm.types.proxy.auth.user_api_key_auth import COMMITTED_BUDGET_SPEND_PARAM +from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import router as llm_passthrough_router from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler import ( VertexPassthroughLoggingHandler, @@ -108,6 +110,7 @@ PROTOCOL_CONSTRAINED_PASS_THROUGH_ROUTES: Final = MappingProxyType( "/laya/v1/systemone": frozenset({"POST"}), "/bespoke/v1/systemone": frozenset({"POST"}), "/eu_jev/v1/systemone": frozenset({"POST"}), + "/opensearch/{endpoint:path}": frozenset({"GET", "HEAD", "POST", "PUT", "DELETE", "PATCH"}), } ) @@ -9692,6 +9695,31 @@ async def _track_cost_for_passthrough_kwargs(kwargs: dict) -> AsyncMock: return increment_spend_counters +@pytest.mark.asyncio +async def test_passthrough_charge_moves_the_committed_spend_of_the_entities_auth_named_budgeted(): + """A forged body value stays out, and loggers that dump the row never get the auth object or the spend it read.""" + counter_keys: Final = EntitySpendCounterKeys.of( + token="hashed-token", team_id="t1", user_id="u1", org_id=None, end_user_id=None + ) + user_api_key_dict: Final = UserAPIKeyAuth( + api_key="hashed-token", + user_id="u1", + committed_budget_spend=CommittedBudgetSpend(counter_keys=counter_keys, key=0.123456), + metadata={"logging": [{"callback_name": "langfuse"}], "kept": "yes"}, + ) + + kwargs: Final = _passthrough_kwargs_for_reservation( + user_api_key_dict, + parsed_body={"litellm_metadata": {"user_api_key_auth": {"committed_budget_spend": None}}}, + ) + increment_spend_counters: Final = await _track_cost_for_passthrough_kwargs(kwargs) + + assert increment_spend_counters.await_args.kwargs["committed_spend_scope"] == counter_keys + assert "user_api_key_auth" not in kwargs["litellm_params"]["metadata"] + assert "committed_budget_spend" not in repr(kwargs["litellm_params"]["metadata"]) + assert "0.123456" not in json.dumps(kwargs, default=str) + + @pytest.mark.asyncio async def test_passthrough_success_reconciles_budget_reservation(): """ @@ -11134,6 +11162,9 @@ def _split_pass_through_body( GEMINI_BODY: Final = '{"contents": [{"parts": [{"text": "hi"}]}], "generationConfig": {"temperature": 0}}' +PROXY_SET_PARAMS: Final = ("metadata", "proxy_server_request", COMMITTED_BUDGET_SPEND_PARAM) + + def _metadata_of(split: _PassThroughSplit) -> Mapping[str, object]: return MappingProxyType(_LITELLM_PARAMS.validate_python(split.litellm_params["metadata"])) @@ -11145,7 +11176,7 @@ def test_passthrough_moves_every_litellm_owned_key_from_the_forwarded_body_into_ ) assert frozenset(split.litellm_params) == frozenset( - ("ttl", "num_retries", "litellm_trace_id", "metadata", "proxy_server_request") + ("ttl", "num_retries", "litellm_trace_id", *PROXY_SET_PARAMS) ) assert tuple(split.litellm_params[k] for k in ("ttl", "num_retries", "litellm_trace_id")) == (30, 2, "trace-a") assert split.forwarded_body == json.loads(GEMINI_BODY) @@ -11179,7 +11210,7 @@ PROXY_STAMPED_NAMES: Final = frozenset( def test_passthrough_keeps_each_registered_litellm_owned_name_out_of_the_forwarded_body(name: str) -> None: split: Final = _split_pass_through_body(json.dumps({name: "owned", **json.loads(GEMINI_BODY)})) - assert frozenset(split.litellm_params) == frozenset((name, "metadata", "proxy_server_request")) + assert frozenset(split.litellm_params) == frozenset((name, *PROXY_SET_PARAMS)) assert split.litellm_params[name] == "owned" assert split.forwarded_body == json.loads(GEMINI_BODY) @@ -11223,7 +11254,7 @@ def test_passthrough_keeps_a_client_model_info_object_for_keys_that_allow_client def test_passthrough_drops_a_client_supplied_proxy_stamped_name(name: str) -> None: split: Final = _split_pass_through_body(json.dumps({name: {"forged": "by-client"}, **json.loads(GEMINI_BODY)})) - assert frozenset(split.litellm_params) == frozenset(("metadata", "proxy_server_request")) + assert frozenset(split.litellm_params) == frozenset(PROXY_SET_PARAMS) assert split.litellm_params["proxy_server_request"] != {"forged": "by-client"}, split.litellm_params assert split.forwarded_body == json.loads(GEMINI_BODY) @@ -11234,7 +11265,7 @@ def test_passthrough_merges_both_metadata_carriers_from_the_body_into_one_metada ' "litellm_metadata": {"lm": "b"}, "generationConfig": {"temperature": 0}}' ) - assert frozenset(split.litellm_params) == frozenset(("ttl", "metadata", "proxy_server_request")) + assert frozenset(split.litellm_params) == frozenset(("ttl", *PROXY_SET_PARAMS)) assert _metadata_of(split) == {**_metadata_of(_split_pass_through_body(GEMINI_BODY)), "client_tag": "a", "lm": "b"} assert split.forwarded_body == json.loads(GEMINI_BODY) @@ -11273,7 +11304,7 @@ def test_passthrough_sees_a_name_appended_to_the_public_list_after_import() -> N finally: litellm.all_litellm_params.remove("registered_later") - assert frozenset(split.litellm_params) == frozenset(("registered_later", "metadata", "proxy_server_request")) + assert frozenset(split.litellm_params) == frozenset(("registered_later", *PROXY_SET_PARAMS)) assert split.forwarded_body == {"contents": [{"parts": [{"text": "hi"}]}]} @@ -11282,7 +11313,7 @@ def test_passthrough_sees_the_public_list_rebound_after_import(monkeypatch: pyte split: Final = _split_pass_through_body(LATE_REGISTERED_BODY) - assert frozenset(split.litellm_params) == frozenset(("registered_later", "metadata", "proxy_server_request")) + assert frozenset(split.litellm_params) == frozenset(("registered_later", *PROXY_SET_PARAMS)) assert split.forwarded_body == {"contents": [{"parts": [{"text": "hi"}]}]} diff --git a/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py b/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py index 8127ffe558b..1817fd7d930 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py +++ b/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py @@ -445,6 +445,47 @@ async def test_chunk_processor_injects_cost_into_openai_passthrough_usage_frame( assert final_payload["usage"]["total_tokens"] == 15 +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("header_value", "cost_expected"), (("true", True), ("TRUE", True), ("1", False), ("false", False), (None, False)) +) +async def test_chunk_processor_honors_the_opt_in_header_while_global_flag_is_off( + monkeypatch: pytest.MonkeyPatch, header_value: str | None, cost_expected: bool +) -> None: + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False) + chunks: Final = _openai_passthrough_stream_chunks() + logging_obj: Final = _unarmed_logging_obj(model="gpt-4o-mini") + headers: Final = {} if header_value is None else {"x-litellm-include-cost-in-streaming-usage": header_value} + logging_obj.model_call_details = {"litellm_params": {"proxy_server_request": {"headers": headers}}} + + received: Final = [ + chunk + async for chunk in PassThroughStreamingHandler.chunk_processor( + response=_make_streaming_response(chunks), + request_body={"model": "gpt-4o-mini", "stream": True}, + litellm_logging_obj=logging_obj, + endpoint_type=EndpointType.OPENAI, + start_time=datetime(2026, 1, 1), + passthrough_success_handler_obj=MagicMock(), + url_route="/openai/v1/chat/completions", + route_streaming_logging=AsyncMock(), + ) + ] + await asyncio.sleep(0) + + usage_payloads: Final = [ + json.loads(chunk.decode("utf-8").split("data:", 1)[1].strip()) + for chunk in received + if b'"total_tokens"' in chunk + ] + assert len(usage_payloads) == 1 + assert ("cost" in usage_payloads[0]["usage"]) is cost_expected + if cost_expected: + assert usage_payloads[0]["usage"]["cost"] > 0 + else: + assert received == chunks + + @pytest.mark.asyncio async def test_chunk_processor_injects_cost_into_usage_frame_fragmented_across_chunks(monkeypatch): """Regression: an SSE usage frame split across transport chunks must still get @@ -826,7 +867,8 @@ async def test_dropped_attempt_logs_its_failure_without_claiming_the_shared_logg logging_obj.carry_failed_attempt_usage() retry: Final = _anthropic_messages_attempt( - logging_obj, _anthropic_stream_of((_anthropic_message_start_frame("msg_a2", 7), *_anthropic_completed_stream_tail())) + logging_obj, + _anthropic_stream_of((_anthropic_message_start_frame("msg_a2", 7), *_anthropic_completed_stream_tail())), ) async for _ in retry: pass @@ -858,6 +900,7 @@ async def test_every_dropped_attempt_logs_its_own_failure() -> None: stashed: Final = logging_obj.model_call_details.get("combined_usage_object") assert isinstance(stashed, litellm.Usage) and stashed.prompt_tokens == 9 + def _anthropic_message_start_frame(message_id: str = "msg_1", input_tokens: int = 52) -> bytes: return _anthropic_sse( "message_start", @@ -890,7 +933,11 @@ def _anthropic_completed_stream_tail() -> tuple[bytes, ...]: _anthropic_sse("content_block_stop", {"type": "content_block_stop", "index": 0}), _anthropic_sse( "message_delta", - {"type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": 3}}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 3}, + }, ), _anthropic_sse("message_stop", {"type": "message_stop"}), ) @@ -923,7 +970,9 @@ def _anthropic_stream_of(frames: tuple[bytes, ...]): "frames, billed", [ pytest.param( - (_anthropic_message_start_frame(), _anthropic_overloaded_error_frame()), True, id="error-after-message-start" + (_anthropic_message_start_frame(), _anthropic_overloaded_error_frame()), + True, + id="error-after-message-start", ), pytest.param((_anthropic_overloaded_error_frame(),), False, id="error-frame-only"), ], @@ -1102,7 +1151,8 @@ async def test_error_frame_then_exhausted_retry_bills_both_attempts() -> None: recorder: Final = _EventRecorder() logging_obj: Final = _anthropic_messages_logging_obj(recorder, "test-error-frame-terminal-billing") superseded: Final = _anthropic_messages_attempt( - logging_obj, _anthropic_stream_of((_anthropic_message_start_frame("msg_a1", 52), _anthropic_overloaded_error_frame())) + logging_obj, + _anthropic_stream_of((_anthropic_message_start_frame("msg_a1", 52), _anthropic_overloaded_error_frame())), ) async for chunk in superseded: if b"event: error" in chunk: @@ -1168,6 +1218,8 @@ async def test_retry_success_log_carries_its_own_usage_after_a_provider_error_fr success_payload = recorder.success_kwargs[0]["standard_logging_object"] assert success_payload["status"] == "success" assert (success_payload["prompt_tokens"], success_payload["completion_tokens"]) == (7, 3) + + def _google_sse(prompt_tokens: int, completion_tokens: int, text: str) -> bytes: payload = { "candidates": [{"content": {"parts": [{"text": text}], "role": "model"}, "index": 0}], 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 79f3e697a4f..f3fe1e2520a 100644 --- a/tests/unit/proxy/proxy_server/test_routes_model_info.py +++ b/tests/unit/proxy/proxy_server/test_routes_model_info.py @@ -12,7 +12,7 @@ from __future__ import annotations import asyncio import copy import json -from collections.abc import Callable +from collections.abc import Callable, Mapping from contextlib import AbstractContextManager from typing import Final from unittest.mock import AsyncMock, MagicMock @@ -26,7 +26,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.proxy._types import LiteLLM_ProxyModelTable, UserAPIKeyAuth from litellm.utils import _invalidate_model_cost_lowercase_map from .conftest import normalize, route_entry_points_to_stubs # type: ignore[import-not-found] @@ -315,7 +315,7 @@ def configured_router(monkeypatch): monkeypatch.setattr(proxy_server, "llm_router", router) monkeypatch.setattr(proxy_server, "llm_model_list", [{"model_name": "gpt-4"}]) monkeypatch.setattr(proxy_server, "user_model", None) - monkeypatch.setattr(proxy_server, "_get_proxy_model_info", lambda model: model) + monkeypatch.setattr(proxy_server, "_get_proxy_model_info", lambda model, caller_aliases: model) yield router @@ -585,6 +585,11 @@ def test_model_info_id_lookup_reports_the_same_cost_as_the_list( _invalidate_model_cost_lowercase_map() +def _discoverable_models(monkeypatch: pytest.MonkeyPatch, models: Mapping[str, list[str]]) -> None: + for provider, provider_models in models.items(): + monkeypatch.setitem(litellm.models_by_provider, provider, provider_models) + + async def test_v1_model_info_star_wildcard_filter_keeps_provider_expansion(monkeypatch): from litellm.proxy._types import SpecialModelNames, UserAPIKeyAuth from litellm.proxy.auth import model_checks @@ -671,6 +676,839 @@ def test_v1_model_info_asks_a_wildcard_server_only_for_keys_that_can_call_it( ) == (expected_names, expected_listings, 0) +def test_model_info_prices_expanded_wildcard_rows_as_the_model_their_calls_route_to( + client: TestClient, + auth_as: Callable[[], AbstractContextManager[object]], + monkeypatch: pytest.MonkeyPatch, + local_model_cost_map: None, +) -> None: + """A catch-all over Fireworks lists ``acme-chat``, which also names a DeepSeek entry, and a + prefixed OpenAI wildcard lists ``my-openai/gpt-acme``, which names no entry at all. Each row must + show the provider and price its calls are billed at.""" + + catalog: Final = { + "fireworks_ai/accounts/fireworks/models/acme-chat": { + "input_cost_per_token": 1e-07, + "output_cost_per_token": 2e-07, + "litellm_provider": "fireworks_ai", + "mode": "chat", + }, + "acme-chat": { + "input_cost_per_token": 9e-07, + "output_cost_per_token": 9e-06, + "litellm_provider": "deepseek", + "mode": "chat", + }, + "gpt-acme": { + "input_cost_per_token": 3e-07, + "output_cost_per_token": 4e-07, + "litellm_provider": "openai", + "mode": "chat", + }, + } + monkeypatch.setattr(litellm, "model_cost", {**litellm.model_cost, **catalog}) + _invalidate_model_cost_lowercase_map() + router: Final = litellm.Router( + model_list=[ + { + "model_name": "*", + "litellm_params": {"model": "fireworks_ai/accounts/fireworks/models/*", "api_key": "x"}, + "model_info": {"id": "fireworks-wildcard"}, + }, + { + "model_name": "my-openai/*", + "litellm_params": {"model": "openai/*", "api_key": "x"}, + "model_info": {"id": "openai-wildcard"}, + }, + ] + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + _discoverable_models( + monkeypatch, + { + "fireworks_ai": ["fireworks_ai/accounts/fireworks/models/acme-chat"], + "openai": ["gpt-acme"], + }, + ) + monkeypatch.setattr(proxy_server, "llm_model_list", router.get_model_list()) + monkeypatch.setattr(proxy_server, "user_model", None) + + with auth_as(): + response: Final = client.get("/model/info") + + assert response.status_code == 200, response.text + listed: Final = { + row["model_name"]: ( + row["litellm_params"]["model"], + row["model_info"].get("litellm_provider"), + row["model_info"].get("input_cost_per_token"), + row["model_info"].get("output_cost_per_token"), + ) + for row in response.json()["data"] + } + assert listed == { + "acme-chat": ("fireworks_ai/accounts/fireworks/models/acme-chat", "fireworks_ai", 1e-07, 2e-07), + "my-openai/gpt-acme": ("openai/gpt-acme", "openai", 3e-07, 4e-07), + } + calls: Final = { + model_name: router.completion( + model=model_name, messages=[{"role": "user", "content": "hi"}], mock_response="hello there" + ) + for model_name in listed + } + assert {model_name: call._hidden_params["response_cost"] for model_name, call in calls.items()} == pytest.approx( + { + model_name: listed[model_name][2] * call.usage.prompt_tokens + + listed[model_name][3] * call.usage.completion_tokens + for model_name, call in calls.items() + } + ) + _invalidate_model_cost_lowercase_map() + + +def test_model_info_prices_a_wildcard_row_as_the_other_deployment_its_calls_reach( + client: TestClient, + auth_as: Callable[[], AbstractContextManager[object]], + monkeypatch: pytest.MonkeyPatch, + local_model_cost_map: None, +) -> None: + """``acme-pro`` is an alias of the ``acme-lite`` deployment and ``acme-max`` is a routing group over it, so + calls to those names never reach the catch-all, which still lists them. Such a row shows the provider and + price of the deployment its calls reach, even where another provider's model goes by the listed name.""" + + catalog: Final = { + "fireworks_ai/accounts/fireworks/models/acme-lite": { + "input_cost_per_token": 5e-07, + "output_cost_per_token": 6e-07, + "litellm_provider": "fireworks_ai", + "mode": "chat", + }, + "fireworks_ai/accounts/fireworks/models/acme-pro": { + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 2.5e-07, + "litellm_provider": "fireworks_ai", + "mode": "chat", + }, + "fireworks_ai/accounts/fireworks/models/acme-max": { + "input_cost_per_token": 3.5e-07, + "output_cost_per_token": 4.5e-07, + "litellm_provider": "fireworks_ai", + "mode": "chat", + }, + "deepseek/acme-lite": { + "input_cost_per_token": 7e-07, + "output_cost_per_token": 8e-07, + "litellm_provider": "deepseek", + "mode": "chat", + }, + "acme-pro": { + "input_cost_per_token": 9e-07, + "output_cost_per_token": 1e-06, + "litellm_provider": "deepseek", + "mode": "chat", + }, + "acme-max": { + "input_cost_per_token": 1.1e-06, + "output_cost_per_token": 1.2e-06, + "litellm_provider": "deepseek", + "mode": "chat", + }, + } + monkeypatch.setattr(litellm, "model_cost", {**litellm.model_cost, **catalog}) + _invalidate_model_cost_lowercase_map() + router: Final = litellm.Router( + model_list=[ + { + "model_name": "*", + "litellm_params": {"model": "fireworks_ai/accounts/fireworks/models/*", "api_key": "x"}, + "model_info": {"id": "fireworks-wildcard"}, + }, + { + "model_name": "acme-lite", + "litellm_params": {"model": "deepseek/acme-lite", "api_key": "x"}, + "model_info": {"id": "deepseek-acme-lite"}, + }, + ], + model_group_alias={"acme-pro": "acme-lite"}, + routing_groups=[{"group_name": "acme-max", "models": ["acme-lite"], "routing_strategy": "simple-shuffle"}], + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + _discoverable_models( + monkeypatch, + { + "fireworks_ai": [ + "fireworks_ai/accounts/fireworks/models/acme-lite", + "fireworks_ai/accounts/fireworks/models/acme-pro", + "fireworks_ai/accounts/fireworks/models/acme-max", + ], + }, + ) + monkeypatch.setattr(proxy_server, "llm_model_list", router.get_model_list()) + monkeypatch.setattr(proxy_server, "user_model", None) + + with auth_as(): + response: Final = client.get("/model/info") + + assert response.status_code == 200, response.text + listed: Final = { + (row["model_info"]["id"], row["model_name"]): ( + row["model_info"].get("input_cost_per_token"), + row["model_info"].get("output_cost_per_token"), + ) + for row in response.json()["data"] + if "*" not in row["model_name"] + } + assert sorted(listed) == [ + ("deepseek-acme-lite", "acme-lite"), + ("deepseek-acme-lite", "acme-pro"), + ("fireworks-wildcard", "acme-max"), + ("fireworks-wildcard", "acme-pro"), + ] + shown_models: Final = { + (row["model_info"]["id"], row["model_name"]): ( + row["model_info"].get("litellm_provider"), + row["model_info"].get("key"), + ) + for row in response.json()["data"] + if "*" not in row["model_name"] and row["model_info"]["id"] != "deepseek-acme-lite" + } + assert {provider for provider, _ in shown_models.values()} == {"deepseek"} + assert listed[("deepseek-acme-lite", "acme-lite")] == (7e-07, 8e-07) + assert {key: prices for key, prices in listed.items() if key[0] != "deepseek-acme-lite"} == { + ("fireworks-wildcard", "acme-max"): (7e-07, 8e-07), + ("fireworks-wildcard", "acme-pro"): (7e-07, 8e-07), + } + calls: Final = { + model_name: router.completion( + model=model_name, messages=[{"role": "user", "content": "hi"}], mock_response="hello there" + ) + for _, model_name in listed + } + assert {key: calls[key[1]]._hidden_params["response_cost"] for key in listed} == pytest.approx( + { + key: input_cost * calls[key[1]].usage.prompt_tokens + output_cost * calls[key[1]].usage.completion_tokens + for key, (input_cost, output_cost) in listed.items() + } + ) + _invalidate_model_cost_lowercase_map() + + +def _catch_all_and_shadowed_wildcard_router(monkeypatch: pytest.MonkeyPatch) -> litellm.Router: + + catalog: Final = { + "gpt-5.6-acme": { + "input_cost_per_token": 3e-07, + "output_cost_per_token": 4e-07, + "litellm_provider": "openai", + "mode": "chat", + }, + "fireworks_ai/accounts/fireworks/models/acme-pro": { + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 2.5e-07, + "litellm_provider": "fireworks_ai", + "mode": "chat", + }, + "fireworks_ai/accounts/fireworks/models/acme-max": { + "input_cost_per_token": 9e-07, + "output_cost_per_token": 9.5e-07, + "litellm_provider": "fireworks_ai", + "mode": "chat", + }, + "fireworks_ai/accounts/fireworks/models/acme-mini": { + "input_cost_per_token": 5e-08, + "output_cost_per_token": 6e-08, + "litellm_provider": "fireworks_ai", + "mode": "chat", + }, + "deepseek/acme-lite": { + "input_cost_per_token": 7e-07, + "output_cost_per_token": 8e-07, + "litellm_provider": "deepseek", + "mode": "chat", + }, + } + monkeypatch.setattr(litellm, "model_cost", {**litellm.model_cost, **catalog}) + _invalidate_model_cost_lowercase_map() + router: Final = litellm.Router( + model_list=[ + { + "model_name": "*", + "litellm_params": {"model": "openai/*", "api_key": "x"}, + "model_info": {"id": "openai-catch-all"}, + }, + { + "model_name": "my-fw/*", + "litellm_params": {"model": "fireworks_ai/accounts/fireworks/models/*", "api_key": "x"}, + "model_info": {"id": "fireworks-wildcard"}, + }, + { + "model_name": "acme-lite", + "litellm_params": {"model": "deepseek/acme-lite", "api_key": "x"}, + "model_info": {"id": "deepseek-acme-lite"}, + }, + ], + model_group_alias={"my-fw/acme-pro": "acme-lite", "my-fw/acme-max": "my-fw/acme-mini"}, + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + _discoverable_models( + monkeypatch, + { + "openai": ["gpt-5.6-acme"], + "fireworks_ai": [ + "fireworks_ai/accounts/fireworks/models/acme-pro", + "fireworks_ai/accounts/fireworks/models/acme-max", + "fireworks_ai/accounts/fireworks/models/acme-mini", + ], + }, + ) + monkeypatch.setattr(proxy_server, "llm_model_list", router.get_model_list()) + monkeypatch.setattr(proxy_server, "user_model", None) + return router + + +def _expanded_wildcard_rows(client: TestClient) -> dict[str, dict[str, object]]: + response: Final = client.get("/model/info") + assert response.status_code == 200, response.text + return { + row["model_name"]: row + for row in response.json()["data"] + if row["model_info"]["id"] != "deepseek-acme-lite" and "*" not in row["model_name"] + } + + +def test_model_info_prices_a_wildcard_row_aliased_to_a_sibling_name_as_the_alias_target( + client: TestClient, + auth_as: Callable[[], AbstractContextManager[object]], + monkeypatch: pytest.MonkeyPatch, + local_model_cost_map: None, +) -> None: + """An alias from ``my-fw/acme-max`` to ``my-fw/acme-mini`` still reaches the same wildcard, but calls run + ``acme-mini``, so the ``my-fw/acme-max`` row must show the price those calls are billed at.""" + router: Final = _catch_all_and_shadowed_wildcard_router(monkeypatch) + + with auth_as(): + rows: Final = _expanded_wildcard_rows(client) + + shown: Final = rows["my-fw/acme-max"]["model_info"] + call: Final = router.completion( + model="my-fw/acme-max", messages=[{"role": "user", "content": "hi"}], mock_response="hello there" + ) + assert call._hidden_params["response_cost"] == pytest.approx( + shown["input_cost_per_token"] * call.usage.prompt_tokens + + shown["output_cost_per_token"] * call.usage.completion_tokens + ) + _invalidate_model_cost_lowercase_map() + + +def test_model_info_prices_a_wildcard_row_the_proxy_renames_as_the_model_its_calls_reach( + client: TestClient, + auth_as: Callable[..., AbstractContextManager[object]], + monkeypatch: pytest.MonkeyPatch, + local_model_cost_map: None, +) -> None: + """The proxy renames a requested model through the global ``model_alias_map``, the key's ``aliases`` and the + key's ``router_settings.model_group_alias`` before the router sees it. A row renamed to another deployment + shows that deployment's price and provider, and a row renamed to a sibling name of its own wildcard shows + that sibling's price.""" + _catch_all_and_shadowed_wildcard_router(monkeypatch) + monkeypatch.setattr(litellm, "model_alias_map", {"my-fw/acme-max": "acme-lite"}) + + with auth_as( + aliases={"gpt-5.6-acme": "acme-lite"}, + router_settings={"model_group_alias": {"my-fw/acme-pro": "my-fw/acme-mini"}}, + ): + rows: Final = _expanded_wildcard_rows(client) + + shown: Final = { + name: ( + row["model_info"].get("input_cost_per_token"), + row["model_info"].get("output_cost_per_token"), + row["model_info"].get("litellm_provider"), + ) + for name, row in rows.items() + } + assert shown == { + "my-fw/acme-max": (7e-07, 8e-07, "deepseek"), + "my-fw/acme-mini": (5e-08, 6e-08, "fireworks_ai"), + "my-fw/acme-pro": (5e-08, 6e-08, "fireworks_ai"), + "gpt-5.6-acme": (7e-07, 8e-07, "deepseek"), + } + _invalidate_model_cost_lowercase_map() + + +def test_model_info_prices_a_wildcard_row_whose_team_alias_targets_a_deleted_team_deployment( + client: TestClient, + auth_as: Callable[..., AbstractContextManager[object]], + monkeypatch: pytest.MonkeyPatch, + local_model_cost_map: None, +) -> None: + """Calls skip a team alias whose target team deployment was deleted, so the row that alias names keeps the + price its calls through the wildcard are billed at, while the team's live aliases still apply.""" + from litellm.proxy.litellm_pre_call_utils import _update_model_if_team_alias_exists + + _catch_all_and_shadowed_wildcard_router(monkeypatch) + caller: Final = { + "team_id": "team-a", + "team_model_aliases": {"my-fw/acme-mini": "model_name_team-a_deleted", "gpt-5.6-acme": "my-fw/acme-mini"}, + } + + with auth_as(**caller): + rows: Final = _expanded_wildcard_rows(client) + + stale_request: Final = {"model": "my-fw/acme-mini"} + live_request: Final = {"model": "gpt-5.6-acme"} + _update_model_if_team_alias_exists(data=stale_request, user_api_key_dict=UserAPIKeyAuth(**caller)) + _update_model_if_team_alias_exists(data=live_request, user_api_key_dict=UserAPIKeyAuth(**caller)) + shown: Final = rows["my-fw/acme-mini"]["model_info"] + live_alias_row: Final = rows["gpt-5.6-acme"]["model_info"] + assert (stale_request["model"], live_request["model"]) == ("my-fw/acme-mini", "my-fw/acme-mini") + assert ( + shown.get("input_cost_per_token"), + shown.get("output_cost_per_token"), + shown.get("litellm_provider"), + ) == (5e-08, 6e-08, "fireworks_ai") + assert (live_alias_row.get("input_cost_per_token"), live_alias_row.get("litellm_provider")) == ( + 5e-08, + "fireworks_ai", + ) + _invalidate_model_cost_lowercase_map() + + +def _fireworks_wildcard_rows(client: TestClient) -> dict[str, dict[str, object]]: + response: Final = client.get("/model/info") + assert response.status_code == 200, response.text + return { + row["model_name"]: row["model_info"] + for row in response.json()["data"] + if row["model_info"]["id"] == "fireworks-wildcard" and "*" not in row["model_name"] + } + + +_FIREWORKS_WILDCARD: Final = { + "model_name": "my-fw/*", + "litellm_params": {"model": "fireworks_ai/accounts/fireworks/models/*", "api_key": "x"}, + "model_info": {"id": "fireworks-wildcard"}, +} + + +def test_model_info_prices_a_wildcard_row_as_the_deployment_the_callers_team_publishes_under_its_name( + client: TestClient, + auth_as: Callable[..., AbstractContextManager[object]], + monkeypatch: pytest.MonkeyPatch, + local_model_cost_map: None, +) -> None: + """Team ``team-a`` publishes ``my-fw/acme-mini`` on its own DeepSeek deployment, so a team-a key's calls to that + name skip the global wildcard and the wildcard's row shows that key the DeepSeek price, while a key outside + the team still sees the Fireworks price its calls bill.""" + _catch_all_and_shadowed_wildcard_router(monkeypatch) + router: Final = litellm.Router( + model_list=[ + _FIREWORKS_WILDCARD, + { + "model_name": "model_name_team-a_acme-mini", + "litellm_params": {"model": "deepseek/acme-lite", "api_key": "x"}, + "model_info": {"id": "team-deepseek", "team_id": "team-a", "team_public_model_name": "my-fw/acme-mini"}, + }, + ] + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + _discoverable_models(monkeypatch, {"fireworks_ai": ["fireworks_ai/accounts/fireworks/models/acme-mini"]}) + monkeypatch.setattr(proxy_server, "llm_model_list", router.get_model_list()) + + with auth_as(team_id="team-a"): + team_rows: Final = _fireworks_wildcard_rows(client) + with auth_as(): + other_rows: Final = _fireworks_wildcard_rows(client) + + assert team_rows["my-fw/acme-mini"].get("input_cost_per_token") == 7e-07 + assert other_rows["my-fw/acme-mini"].get("input_cost_per_token") == 5e-08 + _invalidate_model_cost_lowercase_map() + + +def _call_cost_at_row_rates(call: litellm.ModelResponse, row: Mapping[str, float]) -> float: + return ( + row["input_cost_per_token"] * call.usage.prompt_tokens + + row["output_cost_per_token"] * call.usage.completion_tokens + ) + + +@pytest.mark.parametrize( + ("wildcards", "row_ids"), + [ + pytest.param( + (("gpt-*", "openai/gpt-*", "narrow"), ("*", "openai/*", "catch-all")), + {"narrow", "catch-all"}, + id="catch_all_row_a_narrower_wildcard_serves", + ), + pytest.param( + (("oa/*", "openai/*", "wide"), ("oa/gpt-*", "openai/gpt-*", "narrow")), + {"wide", "narrow"}, + id="overlapping_wildcards", + ), + ], +) +def test_model_info_keeps_the_price_and_limits_of_a_row_another_wildcard_serves( + client: TestClient, + auth_as: Callable[[], AbstractContextManager[object]], + monkeypatch: pytest.MonkeyPatch, + local_model_cost_map: None, + wildcards: tuple[tuple[str, str, str], ...], + row_ids: set[str], +) -> None: + """Two wildcards can list one name while calls reach only one of them. Each row keeps the price and limits of + the model those calls bill, so a client that reads the first row by name still gets them.""" + entry: Final = { + "input_cost_per_token": 3e-07, + "output_cost_per_token": 4e-07, + "max_input_tokens": 4321, + "litellm_provider": "openai", + "mode": "chat", + } + monkeypatch.setattr(litellm, "model_cost", {**litellm.model_cost, "gpt-5.6-acme": entry}) + _invalidate_model_cost_lowercase_map() + router: Final = litellm.Router( + model_list=[ + { + "model_name": model_name, + "litellm_params": {"model": litellm_model, "api_key": "x"}, + "model_info": {"id": model_id}, + } + for model_name, litellm_model, model_id in wildcards + ] + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + _discoverable_models(monkeypatch, {"openai": ["gpt-5.6-acme"]}) + monkeypatch.setattr(proxy_server, "llm_model_list", router.get_model_list()) + monkeypatch.setattr(proxy_server, "user_model", None) + + with auth_as(): + response: Final = client.get("/model/info") + + assert response.status_code == 200, response.text + rows: Final = [row for row in response.json()["data"] if "*" not in row["model_name"]] + assert {row["model_info"]["id"] for row in rows} == row_ids + assert {(row["model_info"]["input_cost_per_token"], row["model_info"]["max_input_tokens"]) for row in rows} == { + (3e-07, 4321) + } + call: Final = router.completion( + model=rows[0]["model_name"], messages=[{"role": "user", "content": "hi"}], mock_response="hello there" + ) + assert call._hidden_params["response_cost"] == pytest.approx(_call_cost_at_row_rates(call, rows[0]["model_info"])) + _invalidate_model_cost_lowercase_map() + + +def test_model_info_prices_an_aliased_wildcard_row_as_the_model_the_other_wildcard_bills( + client: TestClient, + auth_as: Callable[[], AbstractContextManager[object]], + monkeypatch: pytest.MonkeyPatch, + local_model_cost_map: None, +) -> None: + """``based/gpt-5.6-acme`` is aliased to a name another wildcard serves, so its row shows the price of the model + the other wildcard fills in from the name, which is what its calls bill.""" + catalog: Final = { + "gpt-5.6-acme": { + "input_cost_per_token": 3e-07, + "output_cost_per_token": 4e-07, + "litellm_provider": "openai", + "mode": "chat", + }, + "gpt-5.6-acme-mini": { + "input_cost_per_token": 3e-08, + "output_cost_per_token": 4e-08, + "litellm_provider": "openai", + "mode": "chat", + }, + } + monkeypatch.setattr(litellm, "model_cost", {**litellm.model_cost, **catalog}) + _invalidate_model_cost_lowercase_map() + router: Final = litellm.Router( + model_list=[ + { + "model_name": name_pattern, + "litellm_params": {"model": "openai/*", "api_key": "x"}, + "model_info": {"id": name_pattern}, + } + for name_pattern in ("based/*", "other/*") + ], + model_group_alias={"based/gpt-5.6-acme": "other/gpt-5.6-acme-mini"}, + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + _discoverable_models(monkeypatch, {"openai": ["gpt-5.6-acme"]}) + monkeypatch.setattr(proxy_server, "llm_model_list", router.get_model_list()) + monkeypatch.setattr(proxy_server, "user_model", None) + + with auth_as(): + response: Final = client.get("/model/info") + + assert response.status_code == 200, response.text + aliased: Final = next(row for row in response.json()["data"] if row["model_name"] == "based/gpt-5.6-acme") + call: Final = router.completion( + model="based/gpt-5.6-acme", messages=[{"role": "user", "content": "hi"}], mock_response="hello there" + ) + assert aliased["model_info"]["input_cost_per_token"] == 3e-08 + assert call._hidden_params["response_cost"] == pytest.approx(_call_cost_at_row_rates(call, aliased["model_info"])) + _invalidate_model_cost_lowercase_map() + + +def test_model_info_keeps_the_price_of_a_wildcard_row_whose_listed_name_no_route_serves( + client: TestClient, + auth_as: Callable[[], AbstractContextManager[object]], + monkeypatch: pytest.MonkeyPatch, + local_model_cost_map: None, +) -> None: + """``my-*`` expands to provider-prefixed names no route matches, so nothing says the row's price is wrong and + the row keeps the price of the model it lists.""" + entry: Final = { + "input_cost_per_token": 3e-07, + "output_cost_per_token": 4e-07, + "litellm_provider": "openai", + "mode": "chat", + } + monkeypatch.setattr(litellm, "model_cost", {**litellm.model_cost, "gpt-5.6-acme": entry}) + _invalidate_model_cost_lowercase_map() + router: Final = litellm.Router( + model_list=[ + { + "model_name": "my-*", + "litellm_params": {"model": "openai/*", "api_key": "x"}, + "model_info": {"id": "my-wildcard"}, + } + ] + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + _discoverable_models(monkeypatch, {"openai": ["gpt-5.6-acme"]}) + monkeypatch.setattr(proxy_server, "llm_model_list", router.get_model_list()) + monkeypatch.setattr(proxy_server, "user_model", None) + + with auth_as(): + response: Final = client.get("/model/info") + + assert response.status_code == 200, response.text + assert { + row["model_info"].get("input_cost_per_token") for row in response.json()["data"] if "*" not in row["model_name"] + } == {3e-07} + _invalidate_model_cost_lowercase_map() + + +@pytest.mark.parametrize( + "targets", + [ + pytest.param( + ( + { + "litellm_params": {"model": "deepseek/acme-lite", "api_key": "x", "input_cost_per_token": 9e-06}, + "model_info": {"id": "target"}, + }, + ), + id="target_with_its_own_rate", + ), + pytest.param( + ( + {"litellm_params": {"model": "deepseek/acme-lite", "api_key": "x"}, "model_info": {"id": "target-a"}}, + {"litellm_params": {"model": "deepseek/acme-big", "api_key": "x"}, "model_info": {"id": "target-b"}}, + ), + id="targets_that_bill_different_models", + ), + ], +) +def test_model_info_shows_no_wildcard_price_when_the_cost_map_cannot_say_what_calls_bill( + client: TestClient, + auth_as: Callable[[], AbstractContextManager[object]], + monkeypatch: pytest.MonkeyPatch, + local_model_cost_map: None, + targets: tuple[Mapping[str, object], ...], +) -> None: + """An alias sends ``my-fw/acme-mini`` to a group whose bill the cost map cannot state as one price, so the + wildcard's row shows no price rather than one calls are not billed at.""" + catalog: Final = { + "deepseek/acme-lite": { + "input_cost_per_token": 7e-07, + "output_cost_per_token": 8e-07, + "litellm_provider": "deepseek", + "mode": "chat", + }, + "deepseek/acme-big": { + "input_cost_per_token": 1.7e-06, + "output_cost_per_token": 1.8e-06, + "litellm_provider": "deepseek", + "mode": "chat", + }, + } + monkeypatch.setattr(litellm, "model_cost", {**litellm.model_cost, **catalog}) + _invalidate_model_cost_lowercase_map() + router: Final = litellm.Router( + model_list=[_FIREWORKS_WILDCARD, *({"model_name": "acme-target", **target} for target in targets)], + model_group_alias={"my-fw/acme-mini": "acme-target"}, + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + _discoverable_models(monkeypatch, {"fireworks_ai": ["fireworks_ai/accounts/fireworks/models/acme-mini"]}) + monkeypatch.setattr(proxy_server, "llm_model_list", router.get_model_list()) + monkeypatch.setattr(proxy_server, "user_model", None) + + with auth_as(): + rows: Final = _fireworks_wildcard_rows(client) + + assert rows["my-fw/acme-mini"].get("input_cost_per_token") is None + _invalidate_model_cost_lowercase_map() + + +def test_model_info_prices_a_wildcard_row_aliased_to_an_aliased_group_as_the_group_its_calls_reach( + client: TestClient, + auth_as: Callable[[], AbstractContextManager[object]], + monkeypatch: pytest.MonkeyPatch, + local_model_cost_map: None, +) -> None: + """``my-fw/acme-mini`` is aliased to ``acme-target``, a group that is itself aliased onward. Calls follow one + alias and reach ``acme-target``'s deployment, so the row shows that deployment's price, not the onward one's.""" + catalog: Final = { + "deepseek/acme-lite": { + "input_cost_per_token": 7e-07, + "output_cost_per_token": 8e-07, + "litellm_provider": "deepseek", + "mode": "chat", + }, + "deepseek/acme-big": { + "input_cost_per_token": 1.7e-06, + "output_cost_per_token": 1.8e-06, + "litellm_provider": "deepseek", + "mode": "chat", + }, + } + monkeypatch.setattr(litellm, "model_cost", {**litellm.model_cost, **catalog}) + _invalidate_model_cost_lowercase_map() + router: Final = litellm.Router( + model_list=[ + _FIREWORKS_WILDCARD, + { + "model_name": "acme-target", + "litellm_params": {"model": "deepseek/acme-lite", "api_key": "x"}, + "model_info": {"id": "target"}, + }, + { + "model_name": "acme-onward", + "litellm_params": {"model": "deepseek/acme-big", "api_key": "x"}, + "model_info": {"id": "onward"}, + }, + ], + model_group_alias={"my-fw/acme-mini": "acme-target", "acme-target": "acme-onward"}, + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + _discoverable_models(monkeypatch, {"fireworks_ai": ["fireworks_ai/accounts/fireworks/models/acme-mini"]}) + monkeypatch.setattr(proxy_server, "llm_model_list", router.get_model_list()) + monkeypatch.setattr(proxy_server, "user_model", None) + + with auth_as(): + rows: Final = _fireworks_wildcard_rows(client) + + call: Final = router.completion( + model="my-fw/acme-mini", messages=[{"role": "user", "content": "hi"}], mock_response="hello there" + ) + assert rows["my-fw/acme-mini"].get("input_cost_per_token") is not None + assert call._hidden_params["response_cost"] == pytest.approx(_call_cost_at_row_rates(call, rows["my-fw/acme-mini"])) + _invalidate_model_cost_lowercase_map() + + +def test_model_info_prices_a_wildcard_row_aliased_to_another_deployment_as_that_deployment_not_the_wildcards_base_model( + client: TestClient, + auth_as: Callable[[], AbstractContextManager[object]], + monkeypatch: pytest.MonkeyPatch, + local_model_cost_map: None, +) -> None: + """The wildcard declares a ``base_model`` for its own calls, but ``my-fw/acme-mini`` is aliased to another + deployment, so the row shows the price that deployment bills.""" + catalog: Final = { + "deepseek/acme-lite": { + "input_cost_per_token": 7e-07, + "output_cost_per_token": 8e-07, + "litellm_provider": "deepseek", + "mode": "chat", + }, + "deepseek/acme-base": { + "input_cost_per_token": 9e-06, + "output_cost_per_token": 9.5e-06, + "litellm_provider": "deepseek", + "mode": "chat", + }, + } + monkeypatch.setattr(litellm, "model_cost", {**litellm.model_cost, **catalog}) + _invalidate_model_cost_lowercase_map() + router: Final = litellm.Router( + model_list=[ + {**_FIREWORKS_WILDCARD, "model_info": {"id": "fireworks-wildcard", "base_model": "deepseek/acme-base"}}, + { + "model_name": "acme-target", + "litellm_params": {"model": "deepseek/acme-lite", "api_key": "x"}, + "model_info": {"id": "target"}, + }, + ], + model_group_alias={"my-fw/acme-mini": "acme-target"}, + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + _discoverable_models(monkeypatch, {"fireworks_ai": ["fireworks_ai/accounts/fireworks/models/acme-mini"]}) + monkeypatch.setattr(proxy_server, "llm_model_list", router.get_model_list()) + monkeypatch.setattr(proxy_server, "user_model", None) + + with auth_as(): + rows: Final = _fireworks_wildcard_rows(client) + + call: Final = router.completion( + model="my-fw/acme-mini", messages=[{"role": "user", "content": "hi"}], mock_response="hello there" + ) + assert rows["my-fw/acme-mini"].get("input_cost_per_token") is not None + assert call._hidden_params["response_cost"] == pytest.approx(_call_cost_at_row_rates(call, rows["my-fw/acme-mini"])) + _invalidate_model_cost_lowercase_map() + + +def test_saving_a_replayed_model_info_wildcard_row_stores_none_of_its_shown_metadata( + client: TestClient, + auth_as: Callable[[], AbstractContextManager[object]], + monkeypatch: pytest.MonkeyPatch, + local_model_cost_map: None, +) -> None: + """An API client that saves a ``/model/info`` row back unchanged, apart from one edit, must not pin the + cost-map metadata the row showed onto the wildcard, whether the row is a catch-all name or a name an alias + sends to a sibling name or to another deployment.""" + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo, updateDeployment + + router: Final = _catch_all_and_shadowed_wildcard_router(monkeypatch) + with auth_as(): + rows: Final = _expanded_wildcard_rows(client) + + assert sorted(rows) == ["gpt-5.6-acme", "my-fw/acme-max", "my-fw/acme-mini", "my-fw/acme-pro"] + sources: Final = { + deployment["model_info"]["id"]: deployment for deployment in router.get_model_list() or [] + } + stored: Final = { + name: json.loads( + update_db_model( + db_model=Deployment( + model_name=sources[row["model_info"]["id"]]["model_name"], + litellm_params=LiteLLM_Params(model=sources[row["model_info"]["id"]]["litellm_params"]["model"]), + model_info=ModelInfo(id=row["model_info"]["id"]), + ), + updated_patch=updateDeployment( + model_name=name, model_info=ModelInfo(**{**row["model_info"], "access_groups": ["prod"]}) + ), + llm_router=router, + )["model_info"] + ) + for name, row in rows.items() + } + shown_metadata: Final = { + name: set(row["model_info"]) - set(sources[row["model_info"]["id"]]["model_info"]) for name, row in rows.items() + } + assert sorted(name for name, row in rows.items() if row["model_info"].get("key")) == [ + "gpt-5.6-acme", + "my-fw/acme-max", + "my-fw/acme-mini", + "my-fw/acme-pro", + ] + assert {name: info["access_groups"] for name, info in stored.items()} == dict.fromkeys(rows, ["prod"]) + assert {name: shown_metadata[name] & set(info) for name, info in stored.items()} == dict.fromkeys(rows, set()) + _invalidate_model_cost_lowercase_map() + + # --------------------------------------------------------------------------- # GET /model/info — team BYOK scoping (issue #30983) # --------------------------------------------------------------------------- diff --git a/tests/unit/proxy/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index 19db80b8930..d04021758ef 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/tests/unit/proxy/proxy_server/test_spend_counters.py @@ -38,10 +38,13 @@ from litellm.models.end_user import LiteLLM_EndUserTable from litellm.models.tag import LiteLLM_TagTable from litellm.models.team import LiteLLM_TeamTable, LiteLLM_TeamTableCachedObj from litellm.models.user import LiteLLM_UserTable +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_local from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, end_user_cache_key, tag_cache_key from litellm.proxy.management_endpoints.team_endpoints import refresh_cached_team +from litellm.proxy.spend_tracking.committed_spend import committed_spend_key +from litellm.types.proxy.carried_budget_state import EntitySpendCounterKeys from .conftest import normalize from litellm.rust_bridge._native import NativeRedisStoreError @@ -679,7 +682,7 @@ async def test_increment_spend_counters_pipelines_all_scopes_in_one_redis_call( "spend:org:org-1", } spend = _spend_counters(monkeypatch, redis=dict.fromkeys(expected_keys, 1.0)) - monkeypatch.setattr(ps, "user_api_key_cache", ps.DualCache()) + monkeypatch.setattr(ps, "user_api_key_cache", UserApiKeyCache()) monkeypatch.setattr(ps, "prisma_client", None) await ps.increment_spend_counters( @@ -697,6 +700,97 @@ async def test_increment_spend_counters_pipelines_all_scopes_in_one_redis_call( assert spend.memory == dict.fromkeys(expected_keys, 1.5) +class _VerificationTokenTable: + def __init__(self, spend: float) -> None: + self._spend: Final = spend + self.lookups: list[dict[str, str]] = [] + + async def find_unique(self, where: dict[str, str]) -> MagicMock: + self.lookups.append(where) + return MagicMock(token=where["token"], spend=self._spend) + + +def _prisma_with_key_spend(spend: float) -> MagicMock: + prisma: Final = MagicMock() + prisma.db.litellm_verificationtoken = _VerificationTokenTable(spend) + return prisma + + +@pytest.mark.asyncio +async def test_committed_twin_takes_the_charge_but_never_a_budget_hold(monkeypatch): + """The key counter holds the DB spend 10.0, this request's 2.0 reservation and another request's 1.0 hold.""" + spend: Final = _spend_counters(monkeypatch, redis=None, memory={"spend:key:hashed": 13.0}) + monkeypatch.setattr(ps, "user_api_key_cache", UserApiKeyCache()) + monkeypatch.setattr(ps, "prisma_client", _prisma_with_key_spend(10.0)) + reservation: Final = {"reserved_cost": 2.0, "entries": [{"counter_key": "spend:key:hashed"}], "finalized": False} + scope: Final = EntitySpendCounterKeys(key="spend:key:hashed") + + await ps.increment_spend_counters( + token="hashed", + team_id=None, + user_id=None, + response_cost=0.5, + budget_reservation=reservation, + committed_spend_scope=scope, + ) + await ps.increment_spend_counters( + token="hashed", team_id=None, user_id=None, response_cost=0.25, committed_spend_scope=scope + ) + + assert spend.memory == pytest.approx({"spend:key:hashed": 11.75, committed_spend_key("spend:key:hashed"): 10.75}) + + +@pytest.mark.asyncio +async def test_a_charge_outside_the_committed_scope_keeps_no_twin_and_reads_no_stored_spend(monkeypatch): + spend: Final = _spend_counters(monkeypatch, redis=None, memory={"spend:key:hashed": 10.0}) + prisma: Final = _prisma_with_key_spend(10.0) + key_cache: Final = UserApiKeyCache() + key_cache.set_cache(key="hashed", value=UserAPIKeyAuth(token="hashed")) + monkeypatch.setattr(ps, "user_api_key_cache", key_cache) + monkeypatch.setattr(ps, "prisma_client", prisma) + + await ps.increment_spend_counters(token="hashed", team_id=None, user_id=None, response_cost=0.5) + + assert spend.memory == pytest.approx({"spend:key:hashed": 10.5}) + assert prisma.db.litellm_verificationtoken.lookups == [] + + +@pytest.mark.asyncio +async def test_only_the_entities_in_the_committed_scope_take_the_charge_on_their_twin(monkeypatch): + spend: Final = _spend_counters(monkeypatch, redis=None, memory={"spend:key:hashed": 10.0, "spend:team:t1": 10.0}) + monkeypatch.setattr(ps, "user_api_key_cache", UserApiKeyCache()) + monkeypatch.setattr(ps, "prisma_client", _prisma_with_key_spend(10.0)) + + await ps.increment_spend_counters( + token="hashed", + team_id="t1", + user_id=None, + response_cost=0.5, + committed_spend_scope=EntitySpendCounterKeys(key="spend:key:hashed"), + ) + + assert spend.memory.get(committed_spend_key("spend:key:hashed")) == pytest.approx(10.5) + assert committed_spend_key("spend:team:t1") not in spend.memory + + +@pytest.mark.asyncio +async def test_a_budgeted_team_whose_id_names_a_window_keeps_its_charge_and_gets_no_twin(monkeypatch): + team_counter_key: Final = "spend:team:t1:window:1d" + spend: Final = _spend_counters(monkeypatch, redis=None, memory={team_counter_key: 1.0}) + monkeypatch.setattr(ps, "user_api_key_cache", UserApiKeyCache()) + monkeypatch.setattr(ps, "prisma_client", MagicMock()) + + await ps.increment_spend_counters( + token=None, + team_id="t1:window:1d", + user_id=None, + response_cost=0.5, + committed_spend_scope=EntitySpendCounterKeys(team=team_counter_key), + ) + + assert spend.memory == pytest.approx({team_counter_key: 1.5}) + + _PENDING_SPEND_COUNTER_KEYS: tuple[str, ...] = ( "spend:key:hashed", "spend:team:team-1", @@ -717,7 +811,7 @@ def _warm_counters_behind_a_failing_pipeline(monkeypatch, failure: str) -> FakeS fail_pipelines=True, failure=failure, ) - monkeypatch.setattr(ps, "user_api_key_cache", ps.DualCache()) + monkeypatch.setattr(ps, "user_api_key_cache", UserApiKeyCache()) monkeypatch.setattr(ps, "prisma_client", None) return spend diff --git a/tests/unit/proxy/proxy_server/test_team_model_name_translation.py b/tests/unit/proxy/proxy_server/test_team_model_name_translation.py index b647fbdd02a..e2b4acc0c09 100644 --- a/tests/unit/proxy/proxy_server/test_team_model_name_translation.py +++ b/tests/unit/proxy/proxy_server/test_team_model_name_translation.py @@ -808,7 +808,7 @@ async def test_model_info_v1_litellm_model_id_include_team_models_filters_inacce monkeypatch.setattr(ps, "llm_model_list", [team_row]) monkeypatch.setattr(ps, "llm_router", router) monkeypatch.setattr(ps, "prisma_client", MagicMock()) - monkeypatch.setattr(ps, "_get_proxy_model_info", lambda model: team_row) + monkeypatch.setattr(ps, "_get_proxy_model_info", lambda model, caller_aliases: team_row) monkeypatch.setattr(ps, "_populate_team_access_on_models", _fake_populate) caller = UserAPIKeyAuth( @@ -843,7 +843,7 @@ async def test_model_info_v1_litellm_model_id_team_id_applies_team_filter(monkey monkeypatch.setattr(ps, "llm_model_list", [team_row]) monkeypatch.setattr(ps, "llm_router", router) monkeypatch.setattr(ps, "prisma_client", MagicMock()) - monkeypatch.setattr(ps, "_get_proxy_model_info", lambda model: team_row) + monkeypatch.setattr(ps, "_get_proxy_model_info", lambda model, caller_aliases: team_row) monkeypatch.setattr(ps, "_populate_team_access_on_models", _fake_populate) monkeypatch.setattr(ps, "_filter_models_by_team_id", team_filter) diff --git a/tests/unit/proxy/search_endpoints/test_endpoints.py b/tests/unit/proxy/search_endpoints/test_endpoints.py index 7bb4df37314..92e9076678f 100644 --- a/tests/unit/proxy/search_endpoints/test_endpoints.py +++ b/tests/unit/proxy/search_endpoints/test_endpoints.py @@ -142,6 +142,7 @@ def test_direct_search_body_tool_name_is_denied_under_search_tool_deny_by_defaul def _json_request(body: dict[str, object]) -> MagicMock: request = MagicMock() request.body = AsyncMock(return_value=orjson.dumps(body)) + request.headers = {} return request diff --git a/tests/unit/proxy/spend_tracking/test_carried_budget_state.py b/tests/unit/proxy/spend_tracking/test_carried_budget_state.py index 0bdf43b396c..75b9e422330 100644 --- a/tests/unit/proxy/spend_tracking/test_carried_budget_state.py +++ b/tests/unit/proxy/spend_tracking/test_carried_budget_state.py @@ -2,6 +2,9 @@ from datetime import datetime, timezone +import pytest + +from litellm.models.end_user import LiteLLM_EndUserTable from litellm.proxy._types import ( LiteLLM_BudgetTable, LiteLLM_OrganizationTable, @@ -11,10 +14,13 @@ from litellm.proxy._types import ( ) from litellm.proxy.spend_tracking.carried_budget_state import ( carried_budget_metadata, + carry_committed_budget_spend, carry_organization_budget_state, carry_team_and_user_budget_state, ) from litellm.types.proxy.carried_budget_state import ( + CommittedBudgetSpend, + EntitySpendCounterKeys, KeyBudgetSnapshot, OrgBudgetSnapshot, TeamBudgetSnapshot, @@ -123,10 +129,105 @@ def test_snapshots_never_reach_the_serialized_token(): user_object=LiteLLM_UserTable(user_id="u1", user_alias="Alice"), ) token.org_budget_snapshot = OrgBudgetSnapshot(spend=1.0, max_budget=2.0) + metadata_without_committed_spend = dict(carried_budget_metadata(token)) + token.committed_budget_spend = CommittedBudgetSpend(key=1.0, team=2.0, user=3.0, org=4.0, end_user=5.0) dumped = token.model_dump() assert "team_budget_snapshot" not in dumped assert "user_budget_snapshot" not in dumped assert "org_budget_snapshot" not in dumped + assert "committed_budget_spend" not in dumped + assert dict(carried_budget_metadata(token)) == metadata_without_committed_spend assert UserAPIKeyAuth(**dumped).team_budget_snapshot is None + + +def _committed_reader(committed: CommittedBudgetSpend): + reads: list[EntitySpendCounterKeys] = [] + + async def read_committed(counter_keys: EntitySpendCounterKeys) -> CommittedBudgetSpend: + reads.append(counter_keys) + return committed + + return read_committed, reads + + +@pytest.mark.asyncio +async def test_committed_spend_is_read_once_for_every_budgeted_entity(): + token = UserAPIKeyAuth(token="hashed", team_id="t1", user_id="u1", org_id="o1", spend=1.0, max_budget=10.0) + token.org_budget_snapshot = OrgBudgetSnapshot(spend=40.0, max_budget=500.0) + committed = CommittedBudgetSpend(key=1.3, team=20.3, user=5.3, org=40.3, end_user=0.8) + read_committed, reads = _committed_reader(committed) + + await carry_committed_budget_spend( + valid_token=token, + team_object=LiteLLM_TeamTable(team_id="t1", max_budget=300.0, spend=20.0), + user_object=LiteLLM_UserTable(user_id="u1", max_budget=50.0, spend=5.0), + end_user_object=LiteLLM_EndUserTable( + user_id="cust-1", blocked=False, spend=0.5, litellm_budget_table=LiteLLM_BudgetTable(max_budget=3.0) + ), + read_committed=read_committed, + ) + + assert reads == [ + EntitySpendCounterKeys( + key="spend:key:hashed", + team="spend:team:t1", + user="spend:user:u1", + org="spend:org:o1", + end_user="spend:end_user:cust-1", + ) + ] + assert token.committed_budget_spend == committed + assert (token.spend, token.team_spend) == (1.0, None) + + +@pytest.mark.asyncio +async def test_route_without_a_budget_check_carries_the_counters_it_charges_and_reads_nothing(): + token = UserAPIKeyAuth(token="hashed", max_budget=10.0) + + await carry_committed_budget_spend( + valid_token=token, team_object=None, user_object=None, end_user_object=None, read_committed=None + ) + + assert token.committed_budget_spend == CommittedBudgetSpend( + counter_keys=EntitySpendCounterKeys(key="spend:key:hashed") + ) + + +@pytest.mark.asyncio +async def test_committed_spend_skips_entities_without_a_max_budget(): + token = UserAPIKeyAuth(token="hashed", team_id="t1", user_id="u1", org_id="o1", spend=1.0) + token.org_budget_snapshot = OrgBudgetSnapshot(spend=40.0, max_budget=None) + read_committed, reads = _committed_reader(CommittedBudgetSpend()) + + await carry_committed_budget_spend( + valid_token=token, + team_object=LiteLLM_TeamTable(team_id="t1", spend=20.0), + user_object=LiteLLM_UserTable(user_id="u1", spend=5.0), + end_user_object=LiteLLM_EndUserTable( + user_id="cust-1", blocked=False, spend=0.5, litellm_budget_table=LiteLLM_BudgetTable(max_budget=None) + ), + read_committed=read_committed, + ) + + assert reads == [EntitySpendCounterKeys()] + assert token.committed_budget_spend == CommittedBudgetSpend() + + +@pytest.mark.asyncio +async def test_committed_spend_reads_the_team_org_when_the_key_has_none(): + token = UserAPIKeyAuth(token="hashed", team_id="t1") + token.org_budget_snapshot = OrgBudgetSnapshot(spend=40.0, max_budget=500.0) + read_committed, reads = _committed_reader(CommittedBudgetSpend(org=41.0)) + + await carry_committed_budget_spend( + valid_token=token, + team_object=LiteLLM_TeamTable(team_id="t1", organization_id="team-org"), + user_object=None, + end_user_object=None, + read_committed=read_committed, + ) + + assert reads == [EntitySpendCounterKeys(org="spend:org:team-org")] + assert token.committed_budget_spend == CommittedBudgetSpend(org=41.0) diff --git a/tests/unit/proxy/spend_tracking/test_committed_spend.py b/tests/unit/proxy/spend_tracking/test_committed_spend.py new file mode 100644 index 00000000000..1ad5a2e7ed4 --- /dev/null +++ b/tests/unit/proxy/spend_tracking/test_committed_spend.py @@ -0,0 +1,274 @@ +"""Committed spend twins hold only charged cost, so the remaining budget gauges never count a budget hold.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Callable, Mapping +from types import MappingProxyType, SimpleNamespace +from typing import Final +from unittest.mock import MagicMock +from uuid import uuid4 + +import pytest + +import litellm +from litellm.integrations.prometheus import PrometheusLogger +from litellm.proxy.spend_tracking.committed_spend import ( + bind_committed_admission_counters, + drop_committed_spend, + read_admission_committed_spend, + read_committed_budget_spend, + seed_committed_spend, +) +from litellm.proxy.spend_tracking.spend_counter_batch import ( + active_spend_counter_batch, + read_batched_spend_counter, + spend_counter_batch_scope, +) +from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys +from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters + + +def _no_side_effect() -> None: + return None + + +class _FindUniqueTable: + def __init__(self, row: SimpleNamespace | None, on_read: Callable[[], None]) -> None: + self._row: Final = row + self._on_read: Final = on_read + self.reads = 0 + + async def find_unique(self, where: dict[str, str]) -> SimpleNamespace | None: + self.reads += 1 + await asyncio.sleep(0) + self._on_read() + return self._row + + +def _prisma_with_spend(spend: float | None, on_read: Callable[[], None] = _no_side_effect) -> SimpleNamespace: + row: Final = None if spend is None else SimpleNamespace(spend=spend) + return SimpleNamespace( + db=SimpleNamespace( + litellm_verificationtoken=_FindUniqueTable(row, on_read), + litellm_teamtable=_FindUniqueTable(row, on_read), + litellm_endusertable=_FindUniqueTable(row, on_read), + ) + ) + + +@pytest.fixture(params=[False, True], ids=["in_memory", "redis"]) +def has_redis(request: pytest.FixtureRequest) -> bool: + return request.param + + +def _spend_counters(has_redis: bool, counters: Mapping[str, float] = MappingProxyType({})) -> FakeSpendCounters: + if has_redis: + return FakeSpendCounters(redis=dict(counters)) + return FakeSpendCounters(redis=None, memory=dict(counters)) + + +@pytest.mark.asyncio +async def test_read_returns_each_committed_twin_and_none_for_one_never_charged(has_redis: bool): + spend: Final = _spend_counters( + has_redis, + {"spend:committed:key:hashed": 1.5, "spend:committed:end_user:cust-1": 0.25, "spend:team:t1": 99.0}, + ) + counter_keys: Final = EntitySpendCounterKeys.of( + token="hashed", team_id="t1", user_id=None, org_id=None, end_user_id="cust-1" + ) + + committed: Final = await read_committed_budget_spend(spend.store(), counter_keys) + + assert committed == CommittedBudgetSpend(counter_keys=counter_keys, key=1.5, end_user=0.25) + + +@pytest.mark.asyncio +async def test_a_twin_holding_garbage_reads_as_unknown_without_hiding_the_other_twins(): + spend: Final = FakeSpendCounters( + redis={"spend:committed:key:hashed": "not-a-number", "spend:committed:team:t1": 4.0} + ) + counter_keys: Final = EntitySpendCounterKeys.of( + token="hashed", team_id="t1", user_id=None, org_id=None, end_user_id=None + ) + + committed: Final = await read_committed_budget_spend(spend.store(), counter_keys) + + assert committed == CommittedBudgetSpend(counter_keys=counter_keys, team=4.0) + + +@pytest.mark.parametrize( + ("counter_keys", "expected_spend"), + [ + (EntitySpendCounterKeys(key="spend:key:hashed"), {"key": 12.0}), + (EntitySpendCounterKeys(end_user="spend:end_user:cust-1"), {"end_user": 12.0}), + ], +) +@pytest.mark.asyncio +async def test_seed_writes_the_db_spend_when_the_twin_is_absent( + has_redis: bool, counter_keys: EntitySpendCounterKeys, expected_spend: dict[str, float] +): + store: Final = _spend_counters(has_redis).store() + (counter_key,) = counter_keys.tracked() + + seeded: Final = await seed_committed_spend(store, _prisma_with_spend(12.0), counter_key) + + assert seeded is True + assert await read_committed_budget_spend(store, counter_keys) == CommittedBudgetSpend( + counter_keys=counter_keys, **expected_spend + ) + + +@pytest.mark.asyncio +async def test_concurrent_charges_on_a_cold_twin_read_the_db_once(has_redis: bool): + store: Final = _spend_counters(has_redis).store() + prisma_client: Final = _prisma_with_spend(12.0) + counter_key: Final = f"spend:key:{uuid4()}" + + seeded: Final = await asyncio.gather(*(seed_committed_spend(store, prisma_client, counter_key) for _ in range(5))) + + assert seeded == [True] * 5 + assert prisma_client.db.litellm_verificationtoken.reads == 1 + + +@pytest.mark.asyncio +async def test_seed_reports_no_twin_when_the_twin_cannot_be_read(): + async def unreachable(counter: str) -> None: + raise ConnectionError("redis unreachable") + + spend: Final = FakeSpendCounters(read_hook=unreachable) + + assert await seed_committed_spend(spend.store(), _prisma_with_spend(12.0), "spend:team:t1") is False + + +@pytest.mark.asyncio +async def test_seed_never_lowers_a_twin_another_replica_already_advanced(has_redis: bool): + store: Final = _spend_counters(has_redis, {"spend:committed:team:t1": 20.0}).store() + + seeded: Final = await seed_committed_spend(store, _prisma_with_spend(12.0), "spend:team:t1") + + committed: Final = await read_committed_budget_spend(store, EntitySpendCounterKeys(team="spend:team:t1")) + assert seeded is True + assert committed.team == 20.0 + + +@pytest.mark.asyncio +async def test_seed_keeps_the_twin_another_worker_wrote_while_this_one_read_the_db(): + spend: Final = FakeSpendCounters(redis={}) + + def another_worker_seeds_and_charges() -> None: + assert spend.redis is not None + spend.redis["spend:committed:team:t1"] = 20.0 + + seeded: Final = await seed_committed_spend( + spend.store(), _prisma_with_spend(12.0, on_read=another_worker_seeds_and_charges), "spend:team:t1" + ) + + committed: Final = await read_committed_budget_spend(spend.store(), EntitySpendCounterKeys(team="spend:team:t1")) + assert seeded is True + assert committed.team == 20.0 + + +@pytest.mark.asyncio +async def test_read_without_any_counter_key_sends_redis_nothing(): + spend: Final = FakeSpendCounters(redis={}) + + committed: Final = await read_committed_budget_spend(spend.store(), EntitySpendCounterKeys()) + + assert committed == CommittedBudgetSpend() + assert spend.commands == [] + + +@pytest.mark.asyncio +async def test_seed_reports_no_twin_when_redis_rejects_the_write(): + spend: Final = FakeSpendCounters(redis={}, fail=True) + + seeded: Final = await seed_committed_spend(spend.store(), _prisma_with_spend(12.0), "spend:team:t1") + + committed: Final = await read_committed_budget_spend(spend.store(), EntitySpendCounterKeys(team="spend:team:t1")) + assert seeded is False + assert committed.team is None + + +@pytest.mark.asyncio +async def test_seed_creates_no_twin_when_the_db_has_no_spend_for_the_entity(has_redis: bool): + store: Final = _spend_counters(has_redis).store() + + seeded: Final = await seed_committed_spend(store, _prisma_with_spend(None), "spend:key:hashed") + + committed: Final = await read_committed_budget_spend(store, EntitySpendCounterKeys(key="spend:key:hashed")) + assert seeded is False + assert committed.key is None + + +@pytest.mark.asyncio +async def test_admission_read_rides_the_admission_mget_with_the_enforcement_counters( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "callbacks", [MagicMock(spec=PrometheusLogger)]) + spend: Final = FakeSpendCounters(redis={"spend:committed:key:hashed": 1.5}) + store: Final = spend.store() + token: Final = UserAPIKeyAuth(token="hashed", team_id="t1") + counter_keys: Final = EntitySpendCounterKeys(key="spend:key:hashed", team="spend:team:t1") + + with spend_counter_batch_scope(store, counter_keys=frozenset(("spend:key:hashed",))): + bind_committed_admission_counters(token, end_user_id=None) + await read_batched_spend_counter("spend:key:hashed") + committed: Final = await read_admission_committed_spend(store, counter_keys) + + assert committed == CommittedBudgetSpend(counter_keys=counter_keys, key=1.5) + assert len(spend.commands) == 1 + assert {"spend:committed:key:hashed", "spend:committed:team:t1"} <= set(spend.commands[0].split()[1:]) + + +@pytest.mark.asyncio +async def test_admission_read_outside_a_batch_leaves_redis_alone_and_the_spend_unknown(): + spend: Final = FakeSpendCounters(redis={"spend:committed:key:hashed": 1.5}) + counter_keys: Final = EntitySpendCounterKeys(key="spend:key:hashed") + + committed: Final = await read_admission_committed_spend(spend.store(), counter_keys) + + assert committed == CommittedBudgetSpend(counter_keys=counter_keys) + assert spend.commands == [] + + +@pytest.mark.asyncio +async def test_admission_read_without_redis_reads_the_in_memory_twin(): + spend: Final = FakeSpendCounters(redis=None, memory={"spend:committed:key:hashed": 1.5}) + + committed: Final = await read_admission_committed_spend( + spend.store(), EntitySpendCounterKeys(key="spend:key:hashed") + ) + + assert committed.key == 1.5 + + +def test_no_committed_counter_is_bound_without_a_prometheus_logger(): + with spend_counter_batch_scope(FakeSpendCounters(redis={}).store()): + bind_committed_admission_counters(UserAPIKeyAuth(token="hashed"), end_user_id=None) + batch: Final = active_spend_counter_batch() + assert batch is not None + bound: Final = batch.counter_keys + + assert bound == frozenset() + + +@pytest.mark.asyncio +async def test_drop_removes_the_twin_and_leaves_the_enforcement_counter(has_redis: bool): + store: Final = _spend_counters(has_redis, {"spend:committed:org:o1": 5.0, "spend:org:o1": 5.0}).store() + counter_keys: Final = EntitySpendCounterKeys(org="spend:org:o1") + + await drop_committed_spend(store, ("spend:org:o1",)) + + assert (await read_committed_budget_spend(store, counter_keys)).org is None + assert (await store.read("spend:org:o1"))[0] == 5.0 + + +@pytest.mark.asyncio +async def test_a_failed_drop_never_fails_the_write_that_already_committed(): + spend: Final = FakeSpendCounters(redis={"spend:committed:org:o1": 5.0}, failing=frozenset({"delete"})) + + await drop_committed_spend(spend.store(), ("spend:org:o1",)) + + assert spend.deleted() == frozenset({"spend:committed:org:o1"}) diff --git a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py index c3e916d581d..05667629305 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -1,12 +1,14 @@ import asyncio import collections +import contextlib import datetime import hashlib import json import re import sqlite3 -from collections.abc import Mapping, Sequence +from collections.abc import AsyncIterator, Mapping, Sequence from datetime import timezone +from types import SimpleNamespace from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch @@ -17,14 +19,17 @@ from fastapi.testclient import TestClient import litellm import litellm.proxy.proxy_server as ps +from litellm.integrations.prometheus import PrometheusLogger from litellm.proxy.auth.authorization import OwnedRows from litellm.proxy.auth.authorization_dependencies import get_log_team_lookup, load_permitted_log_team_ids from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.spend_tracking.spend_management_endpoints import global_spend_reset from tests.unit.proxy.db.fake_prisma_engine import ( ReadTargetOf, statement_timeout_error, statement_timeout_reads, ) +from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters def _default_date_range(): @@ -8336,6 +8341,109 @@ async def test_global_spend_refresh_runs_the_refresh_on_the_writer_when_a_reader ] +class _ResetTableRows: + def __init__(self) -> None: + self.update_many: Final = AsyncMock() + + +class _ResetDb: + def __init__(self) -> None: + self.litellm_verificationtoken: Final = _ResetTableRows() + self.litellm_teamtable: Final = _ResetTableRows() + + async def query_raw(self, sql_query: str, after_id: str, limit: int) -> list[dict[str, str]]: + rows: Final = ["hashed"] if "VerificationToken" in sql_query else ["t1"] + return [{"id": row} for row in rows if row > after_id][:limit] + + +class _ResetTransaction: + def __init__(self, db: _ResetDb) -> None: + self.litellm_verificationtoken: Final = db.litellm_verificationtoken + self.litellm_teamtable: Final = db.litellm_teamtable + + async def query_raw(self, sql_query: str) -> list[dict[str, str]]: + return [] + + +def _reset_prisma(db: _ResetDb) -> SimpleNamespace: + @contextlib.asynccontextmanager + async def tx() -> AsyncIterator[_ResetTransaction]: + yield _ResetTransaction(db) + + return SimpleNamespace(db=db, tx=tx) + + +def _committed_spend(monkeypatch: pytest.MonkeyPatch, twins: Sequence[str]) -> FakeSpendCounters: + spend: Final = FakeSpendCounters(redis=None, memory=dict.fromkeys(twins, 7.0)) + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_store", spend.store()) + return spend + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("registered_callbacks", "twin_after"), + [ + ([MagicMock(spec=PrometheusLogger)], None), + ([], 7.0), + ], +) +async def test_global_spend_reset_drops_the_committed_key_and_team_spend_the_gauges_read( + monkeypatch, registered_callbacks, twin_after +): + twins: Final = ("spend:committed:key:hashed", "spend:committed:team:t1") + spend: Final = _committed_spend(monkeypatch, twins) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _reset_prisma(_ResetDb())) + monkeypatch.setattr(litellm, "callbacks", registered_callbacks) + + response: Final = await global_spend_reset() + + assert response["status"] == "success" + assert [spend.memory.get(twin) for twin in twins] == [twin_after, twin_after] + + +class _PagedResetDb(_ResetDb): + def __init__(self, token_count: int) -> None: + super().__init__() + self._tokens: Final = sorted(f"token-{index}" for index in range(token_count)) + self.largest_page = 0 + + async def query_raw(self, sql_query: str, after_id: str, limit: int) -> list[dict[str, str]]: + rows: Final = self._tokens if "VerificationToken" in sql_query else [] + page: Final = [{"id": row} for row in rows if row > after_id][:limit] + self.largest_page = max(self.largest_page, len(page)) + return page + + +@pytest.mark.asyncio +@pytest.mark.parametrize("token_count", [999, 1000, 1001, 2500]) +async def test_global_spend_reset_drops_every_key_twin_reading_the_keys_in_bounded_pages(monkeypatch, token_count): + db: Final = _PagedResetDb(token_count) + twins: Final = tuple(f"spend:committed:key:token-{index}" for index in range(token_count)) + spend: Final = _committed_spend(monkeypatch, twins) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _reset_prisma(db)) + monkeypatch.setattr(litellm, "callbacks", [MagicMock(spec=PrometheusLogger)]) + + response: Final = await global_spend_reset() + + assert response["status"] == "success" + assert spend.memory == {} + assert db.largest_page <= 1000 + + +@pytest.mark.asyncio +async def test_global_spend_reset_still_succeeds_when_listing_the_committed_counters_fails(monkeypatch): + db: Final = _ResetDb() + db.query_raw = AsyncMock(side_effect=ConnectionError("db unreachable")) + _committed_spend(monkeypatch, ("spend:committed:team:t1",)) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _reset_prisma(db)) + monkeypatch.setattr(litellm, "callbacks", [MagicMock(spec=PrometheusLogger)]) + + response: Final = await global_spend_reset() + + assert response["status"] == "success" + db.litellm_teamtable.update_many.assert_awaited_once() + + def test_spend_calculate_rejection_body_is_openai_shaped(client): app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user", api_key="hashed-k" diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index c59cde6d630..901dbb6ea97 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -4,9 +4,9 @@ import copy import datetime import gzip import json -from collections.abc import AsyncIterable, Awaitable, Mapping +from collections.abc import AsyncIterable, Awaitable, Coroutine, Mapping from types import MappingProxyType, SimpleNamespace -from typing import AsyncGenerator, Callable, Final, Iterator, Literal, Optional, Sequence +from typing import AsyncGenerator, Callable, Final, Iterator, Literal, Optional, Sequence, TypeVar from urllib.parse import unquote_plus from unittest.mock import AsyncMock, MagicMock, patch @@ -82,6 +82,7 @@ from litellm.proxy.dd_span_tagger import DDSpanTagger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.anthropic.chat.transformation import AnthropicConfig from litellm.proxy._types import Litellm_EntityType, ProxyErrorTypes, ProxyException +from litellm.litellm_core_utils.total_deadline import current_total_deadline from litellm.proxy._types import UserAPIKeyAuth as ProxyUserAPIKeyAuth from litellm.proxy.auth.fallback_budget import RouterFallbackBudgetCheck from litellm.proxy.hooks.model_max_budget_limiter import ( @@ -5284,9 +5285,7 @@ class TestHandleLLMApiExceptionProviderOverload: assert proxy_exc.type == "internal_server_error" assert proxy_exc.to_dict()["code"] == "500" - async def test_openai_route_answers_wrapped_provider_overload_as_503( - self, monkeypatch: pytest.MonkeyPatch - ) -> None: + async def test_openai_route_answers_wrapped_provider_overload_as_503(self, monkeypatch: pytest.MonkeyPatch) -> None: original: Final = InternalServerError( message="Overloaded", llm_provider="anthropic", @@ -5517,6 +5516,43 @@ class TestAsyncStreamingDataGeneratorFastPath: ProxyLogging._callback_capabilities_cache.clear() + async def test_opt_in_header_injects_cost_while_global_flag_is_off(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False) + monkeypatch.setattr(litellm, "callbacks", []) + ProxyLogging._callback_capabilities_cache.clear() + usage_frame: Final = ( + b'data: {"object":"chat.completion.chunk","choices":[],' + b'"usage":{"prompt_tokens":11,"completion_tokens":4,"total_tokens":15}}\n\n' + ) + + async def frames() -> AsyncGenerator[bytes, None]: + yield usage_frame + + async def stream(headers: dict[str, str]) -> list[object]: + return [ + c + async for c in ProxyBaseLLMRequestProcessing.async_streaming_data_generator( + response=frames(), + user_api_key_dict=MagicMock(spec=UserAPIKeyAuth), + request_data={"model": "gpt-4o-mini", "proxy_server_request": {"headers": headers}}, + proxy_logging_obj=ProxyLogging(user_api_key_cache=MagicMock()), + serialize_chunk=ProxyBaseLLMRequestProcessing.return_sse_chunk, + serialize_error=lambda e: "data: error\n\n", + ) + ] + + plain: Final = await stream({}) + opted_in: Final = await stream({"x-litellm-include-cost-in-streaming-usage": "true"}) + declined: Final = await stream({"x-litellm-include-cost-in-streaming-usage": "false"}) + ProxyLogging._callback_capabilities_cache.clear() + + assert plain == [usage_frame] + assert declined == [usage_frame] + assert isinstance(opted_in[0], bytes) + payload: Final = json.loads(opted_in[0].decode("utf-8").split("data:", 1)[1].strip()) + assert payload["usage"]["cost"] > 0 + assert payload["usage"]["total_tokens"] == 15 + class TestDisconnectGatherCleanup: def _disconnect_request(self) -> Request: @@ -10823,6 +10859,23 @@ class TestProcessChunkWithCostInjection: payload = json.loads(result.decode("utf-8").split("data:", 1)[1].strip()) assert payload["usage"]["cost"] > 0 + def test_request_header_opts_a_frame_into_cost_while_global_flag_is_off( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False) + chunk: Final = ( + b'data: {"object":"chat.completion.chunk","choices":[],' + b'"usage":{"prompt_tokens":11,"completion_tokens":4,"total_tokens":15}}\n\n' + ) + + assert ProxyBaseLLMRequestProcessing.process_chunk_with_cost_injection(chunk, "gpt-4o-mini") == chunk + injected: Final = ProxyBaseLLMRequestProcessing.process_chunk_with_cost_injection( + chunk, "gpt-4o-mini", request_headers={"X-LiteLLM-Include-Cost-In-Streaming-Usage": "True"} + ) + + payload: Final = json.loads(injected.decode("utf-8").split("data:", 1)[1].strip()) + assert payload["usage"]["cost"] > 0 + def test_chunk_ending_in_partial_frame_passes_through_byte_identical(self, monkeypatch): monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) chunk = ( @@ -11161,7 +11214,7 @@ async def test_base_process_llm_request_pings_while_the_upstream_call_is_still_r with patch.object(litellm, "sse_keepalive_ping_interval_seconds", 0.05): with patch.object(ProxyBaseLLMRequestProcessing, "_process_llm_request", slow_inner): response = await processor.base_process_llm_request( - request=MagicMock(spec=Request), + request=MagicMock(spec=Request, headers={}), fastapi_response=Response(), user_api_key_dict=MagicMock(spec=UserAPIKeyAuth), route_type="acompletion", @@ -11339,7 +11392,7 @@ async def test_base_process_llm_request_audits_a_failure_that_lands_after_its_ke with patch.object(litellm, "sse_keepalive_ping_interval_seconds", 0.05): with patch.object(ProxyBaseLLMRequestProcessing, "_process_llm_request", slow_failure): response = await processor.base_process_llm_request( - request=MagicMock(spec=Request), + request=MagicMock(spec=Request, headers={}), fastapi_response=Response(), user_api_key_dict=user_api_key_dict, route_type="acompletion", @@ -11385,7 +11438,7 @@ async def test_base_process_llm_request_honours_a_deployment_hard_disable(deploy with patch.object(litellm, "sse_keepalive_ping_interval_seconds", 0.05): with patch.object(ProxyBaseLLMRequestProcessing, "_process_llm_request", slow_inner): response = await processor.base_process_llm_request( - request=MagicMock(spec=Request), + request=MagicMock(spec=Request, headers={}), fastapi_response=Response(), user_api_key_dict=MagicMock(spec=UserAPIKeyAuth), route_type="acompletion", @@ -13269,3 +13322,343 @@ async def test_a_chained_proxy_redacts_secrets_and_internal_details_from_relayed assert {key: answered[key] for key in _DOWNSTREAM_GUARDRAIL_FIELDS} == dict(_DOWNSTREAM_GUARDRAIL_FIELDS) assert answered["scanner"]["attempts"] == 2 assert [value for value in _DOWNSTREAM_LEAKABLE_VALUES if value in json.dumps(answered)] == [] + + +_ScenarioT = TypeVar("_ScenarioT") + + +def _run_on_autojump_clock(scenario: Coroutine[object, object, _ScenarioT]) -> _ScenarioT: + loop: Final = AutojumpClockLoop() + try: + return loop.run_until_complete(scenario) + finally: + loop.run_until_complete(loop.shutdown_asyncgens()) + loop.close() + + +class TestTotalTimeout: + """``total_timeout`` / ``x-litellm-total-timeout`` bound a whole non-streaming request.""" + + def _request(self, headers: Mapping[str, str]) -> Request: + async def receive(): + await asyncio.Event().wait() + + return Request( + scope={"type": "http", "headers": [(k.encode(), v.encode()) for k, v in headers.items()]}, + receive=receive, + ) + + async def _process( + self, + monkeypatch: pytest.MonkeyPatch, + llm_call: Callable[[Mapping[str, object]], object], + *, + routing: Callable[[], Awaitable[None]] | None = None, + body: Mapping[str, object] = MappingProxyType({}), + headers: Mapping[str, str] = MappingProxyType({}), + moderation: Callable[..., object] | None = None, + post_call: Callable[..., object] | None = None, + route_type: str = "acompletion", + proxy_logging_obj: MagicMock | None = None, + preparation: Callable[..., Awaitable[tuple[dict[str, object], MagicMock]]] | None = None, + ) -> tuple[object, MagicMock]: + from litellm.proxy._types import UserAPIKeyAuth + + logging_obj: Final = MagicMock() + logging_obj.litellm_call_id = "total-timeout-call" + logging_obj.defer_async_logging = False + logging_obj._on_deferred_stream_complete = None + logging_obj.cost_breakdown = None + logging_obj.finalized_cost = None + logging_obj.prepared_cost = None + + def prepared(response: object, now: object) -> object: + logging_obj.prepared_cost = response + return response + + logging_obj.prepare_response_cost.side_effect = prepared + processor: Final = ProxyBaseLLMRequestProcessing( + data={"model": "fake-model", "litellm_logging_obj": logging_obj, **body} + ) + hooks: Final = proxy_logging_obj if proxy_logging_obj is not None else MagicMock(spec=ProxyLogging) + hooks.during_call_hook = AsyncMock(side_effect=moderation, return_value=None) + hooks.update_request_status = AsyncMock(return_value=None) + hooks.post_call_success_hook = AsyncMock( + side_effect=post_call or (lambda data, user_api_key_dict, response: response) + ) + hooks.post_call_response_headers_hook = AsyncMock(return_value=None) + hooks.post_call_failure_hook = AsyncMock(return_value=None) + + async def fake_route_request(**kwargs): + if routing is not None: + await routing() + return llm_call(kwargs["data"]) + + monkeypatch.setattr(litellm.proxy.common_request_processing, "route_request", fake_route_request) + if preparation is not None: + monkeypatch.setattr(processor, "_pre_call_with_fallbacks", preparation) + result: Final = await processor.base_process_llm_request( + request=self._request(headers), + fastapi_response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + route_type=route_type, + proxy_logging_obj=hooks, + general_settings={}, + proxy_config=MagicMock(spec=ProxyConfig), + skip_pre_call_logic=preparation is None, + ) + return result, logging_obj + + @pytest.mark.parametrize( + ("body", "headers"), + [({"total_timeout": 60}, {}), ({}, {"x-litellm-total-timeout": "60"})], + ids=["body", "header"], + ) + def test_a_request_past_its_total_timeout_is_answered_with_a_408_and_its_call_is_cancelled( + self, monkeypatch: pytest.MonkeyPatch, body: Mapping[str, object], headers: Mapping[str, str] + ) -> None: + async def provider_hangs_past_the_budget() -> tuple[ProxyException, bool]: + provider_called: Final = asyncio.Event() + upstream_cancelled: Final = asyncio.Event() + + async def hanging(_: Mapping[str, object]) -> object: + provider_called.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + upstream_cancelled.set() + raise + raise AssertionError + + processing: Final = asyncio.create_task(self._process(monkeypatch, hanging, body=body, headers=headers)) + await provider_called.wait() + with pytest.raises(ProxyException) as raised: + await processing + return raised.value, upstream_cancelled.is_set() + + expired, upstream_cancelled = _run_on_autojump_clock(provider_hangs_past_the_budget()) + + assert expired.code == "408" + assert expired.param == "total_timeout" + assert upstream_cancelled + + def test_a_moderation_check_outlasting_the_budget_after_the_provider_answered_prepares_its_cost( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + answer: Final = litellm.ModelResponse() + + async def moderation_outlasts_the_budget() -> tuple[ProxyException, MagicMock]: + answered: Final = asyncio.Event() + moderating: Final = asyncio.Event() + logging_objs: Final[list[MagicMock]] = [] + + async def provider(data: Mapping[str, object]) -> object: + logging_objs.append(data["litellm_logging_obj"]) # pyright: ignore[reportArgumentType] # the harness's logging double + answered.set() + return answer + + async def slow_moderation(**_: object) -> None: + moderating.set() + await asyncio.Event().wait() + + processing: Final = asyncio.create_task( + self._process(monkeypatch, provider, body={"total_timeout": 60}, moderation=slow_moderation) + ) + await answered.wait() + await moderating.wait() + with pytest.raises(ProxyException) as raised: + await processing + return raised.value, logging_objs[0] + + expired, logging_obj = _run_on_autojump_clock(moderation_outlasts_the_budget()) + + costed: Final = [call.args[0] for call in logging_obj.prepare_response_cost.call_args_list] + assert expired.code == "408" + assert costed + assert all(response is answer for response in costed) + + async def test_a_provider_timeout_inside_the_budget_is_not_turned_into_a_408( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + async def provider_timed_out(_: Mapping[str, object]) -> object: + raise TimeoutError("provider read timeout") + + with pytest.raises(TimeoutError, match="provider read timeout"): + await self._process(monkeypatch, provider_timed_out, body={"total_timeout": 60}) + + async def test_cancelling_a_request_inside_its_budget_still_cancels_it( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + provider_called: Final = asyncio.Event() + + async def hanging(_: Mapping[str, object]) -> object: + provider_called.set() + await asyncio.Event().wait() + raise AssertionError + + processing: Final = asyncio.create_task(self._process(monkeypatch, hanging, body={"total_timeout": 60})) + await provider_called.wait() + processing.cancel() + + with pytest.raises(asyncio.CancelledError): + await processing + + async def test_a_bounded_request_reaches_the_router_bounded_and_without_its_total_timeout_field( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + seen: Final[list[tuple[bool, bool]]] = [] + + async def answered(data: Mapping[str, object]) -> object: + seen.append(("total_timeout" in data, current_total_deadline() is not None)) + return litellm.ModelResponse() + + await self._process(monkeypatch, answered, body={"total_timeout": 60}) + + assert seen == [(False, True)] + + @pytest.mark.parametrize( + ("body", "route_type"), + [ + ({"total_timeout": 60, "stream": True}, "acompletion"), + ({"total_timeout": 60, "background": True}, "aresponses"), + ({"total_timeout": 60}, "aembedding"), + ({"total_timeout": 0}, "acompletion"), + ({"total_timeout": None}, "acompletion"), + ], + ids=["streaming", "background_responses", "other_route", "zero", "null"], + ) + async def test_a_body_total_timeout_it_cannot_apply_is_refused_before_the_call( + self, monkeypatch: pytest.MonkeyPatch, body: Mapping[str, object], route_type: str + ) -> None: + called: Final[list[object]] = [] + + async def answered(data: Mapping[str, object]) -> object: + called.append(data) + return litellm.ModelResponse() + + with pytest.raises(ProxyException) as raised: + await self._process(monkeypatch, answered, body=body, route_type=route_type) + + assert raised.value.code == "400" + assert raised.value.param == "total_timeout" + assert called == [] + + @pytest.mark.parametrize("header", ["60", "10s", "", "0"], ids=["valid", "with_a_unit", "empty", "zero"]) + @pytest.mark.parametrize( + ("body", "route_type"), + [({"stream": True}, "acompletion"), ({"background": True}, "aresponses"), ({}, "aembedding")], + ids=["streaming", "background_responses", "other_route"], + ) + async def test_a_header_total_timeout_leaves_a_request_it_cannot_apply_unbounded_whatever_its_value( + self, monkeypatch: pytest.MonkeyPatch, body: Mapping[str, object], route_type: str, header: str + ) -> None: + bounded: Final[list[bool]] = [] + + async def answered(_: Mapping[str, object]) -> object: + bounded.append(current_total_deadline() is not None) + return litellm.ModelResponse() + + await self._process( + monkeypatch, answered, body=body, headers={"x-litellm-total-timeout": header}, route_type=route_type + ) + + assert bounded == [False] + + def test_a_post_call_hook_outlasting_the_budget_is_a_408_not_a_client_disconnect( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + hooks: Final = MagicMock(spec=ProxyLogging) + + async def guardrail_outlasts_the_budget() -> ProxyException: + guarding: Final = asyncio.Event() + + async def answered(_: Mapping[str, object]) -> object: + return litellm.ModelResponse() + + async def slow_guardrail(**_: object) -> object: + guarding.set() + await asyncio.Event().wait() + raise AssertionError + + processing: Final = asyncio.create_task( + self._process( + monkeypatch, answered, body={"total_timeout": 60}, post_call=slow_guardrail, proxy_logging_obj=hooks + ) + ) + await guarding.wait() + with pytest.raises(ProxyException) as raised: + await processing + return raised.value + + expired: Final = _run_on_autojump_clock(guardrail_outlasts_the_budget()) + + assert expired.code == "408" + hooks.post_call_failure_hook.assert_not_awaited() + + def test_a_budget_that_runs_out_while_routing_also_stops_the_moderation_check( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + async def routing_outlasts_the_budget() -> tuple[ProxyException, bool]: + moderating: Final = asyncio.Event() + routing: Final = asyncio.Event() + moderation_stopped: Final = asyncio.Event() + + async def moderation(**_: object) -> None: + moderating.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + moderation_stopped.set() + raise + + async def stuck_routing() -> None: + routing.set() + await asyncio.Event().wait() + + async def never_called(_: Mapping[str, object]) -> object: + raise AssertionError + + processing: Final = asyncio.create_task( + self._process( + monkeypatch, never_called, routing=stuck_routing, body={"total_timeout": 60}, moderation=moderation + ) + ) + await moderating.wait() + await routing.wait() + with pytest.raises(ProxyException) as raised: + await processing + return raised.value, moderation_stopped.is_set() + + expired, moderation_stopped = _run_on_autojump_clock(routing_outlasts_the_budget()) + + assert expired.code == "408" + assert moderation_stopped + + def test_a_budget_that_runs_out_while_preparing_cuts_preparation_short_and_never_calls_the_provider( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + async def preparation_outlasts_the_budget() -> tuple[ProxyException, bool, bool, float]: + loop: Final = asyncio.get_running_loop() + started: Final = loop.time() + prepared: Final = asyncio.Event() + called: Final = asyncio.Event() + + async def slow_guardrail(**_: object) -> tuple[dict[str, object], MagicMock]: + await asyncio.sleep(120) + prepared.set() + return {"model": "fake-model"}, MagicMock() + + async def provider(_: Mapping[str, object]) -> object: + called.set() + return litellm.ModelResponse() + + with pytest.raises(ProxyException) as raised: + await self._process(monkeypatch, provider, body={"total_timeout": 60}, preparation=slow_guardrail) + return raised.value, prepared.is_set(), called.is_set(), loop.time() - started + + expired, prepared, called, elapsed = _run_on_autojump_clock(preparation_outlasts_the_budget()) + + assert expired.code == "408" + assert not prepared + assert not called + assert elapsed < 120 diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index c4069cb1318..b9dcbff8c29 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -114,7 +114,8 @@ class TestGetMetadataVariableName: assert get_metadata_variable_name(request) == "litellm_metadata" @pytest.mark.parametrize( - "path", ["/v1beta/interactions", "/interactions", "/v1beta/interactions/v1_abc123", "/interactions/v1_abc123/cancel"] + "path", + ["/v1beta/interactions", "/interactions", "/v1beta/interactions/v1_abc123", "/interactions/v1_abc123/cancel"], ) def test_returns_litellm_metadata_for_interactions_routes(self, path): request = self._make_request(path) @@ -810,7 +811,9 @@ async def test_interactions_spend_logs_metadata_survives_the_router_metadata_buc # The router then adds its deployment fields to the same bucket, and the spend log # resolves that bucket. - data["litellm_metadata"].update({"model_group": "gemini-3.1-flash-lite", "deployment": "gemini/gemini-3.1-flash-lite"}) + data["litellm_metadata"].update( + {"model_group": "gemini-3.1-flash-lite", "deployment": "gemini/gemini-3.1-flash-lite"} + ) resolved = get_litellm_metadata_from_kwargs( {"litellm_params": {"metadata": data.get("metadata") or {}, "litellm_metadata": data["litellm_metadata"]}} ) @@ -3761,53 +3764,87 @@ async def test_spend_log_caller_session_stays_out_of_requester_metadata(monkeypa "surface,headers,body,child_id,trace_id", [ ( - "chat/completions", {"X-Session-Id": "child-01"}, {"metadata": {"session_id": "body-001"}}, - "child-01", "parent-01", + "chat/completions", + {"X-Session-Id": "child-01"}, + {"metadata": {"session_id": "body-001"}}, + "child-01", + "parent-01", ), ("responses", {"x-vendor-session-id": "child-01"}, {}, "child-01", "parent-01"), ("messages", {"session-id": "child-01", "user-agent": "codex_cli_rs/1.0"}, {}, "child-01", "parent-01"), ("chat/completions", {"x-litellm-session-id": "short", "x-session-id": "child-01"}, {}, "short", "short"), ("responses", {"x-litellm-trace-id": "trace-01", "x-session-id": "child-01"}, {}, "child-01", "trace-01"), ( - "messages", {"x-litellm-trace-id": "trace-01", "x-litellm-session-id": "child-01"}, - {}, "child-01", "trace-01", + "messages", + {"x-litellm-trace-id": "trace-01", "x-litellm-session-id": "child-01"}, + {}, + "child-01", + "trace-01", ), ("responses", {}, {"metadata": {"session_id": "body-001"}}, "body-001", "parent-01"), ("messages", {}, {"litellm_metadata": {"session_id": "body-001"}}, "body-001", "parent-01"), ("chat/completions", {}, {"litellm_metadata": {"session_id": "body-001"}}, "body-001", "parent-01"), ("messages", {"baggage": "session.id=child-01"}, {}, "child-01", "parent-01"), ( - "responses", {"baggage": "session.id=child-01"}, {"metadata": {"session_id": "body-001"}}, - "body-001", "parent-01", + "responses", + {"baggage": "session.id=child-01"}, + {"metadata": {"session_id": "body-001"}}, + "body-001", + "parent-01", ), ( - "chat/completions", {}, {"litellm_session_id": "root-001", "metadata": {"session_id": "body-001"}}, - "root-001", "parent-01", + "chat/completions", + {}, + {"litellm_session_id": "root-001", "metadata": {"session_id": "body-001"}}, + "root-001", + "parent-01", ), ( - "responses", {}, {"metadata": {"session_id": "body-001"}, "litellm_metadata": {"session_id": "active-1"}}, - "active-1", "parent-01", + "responses", + {}, + {"metadata": {"session_id": "body-001"}, "litellm_metadata": {"session_id": "active-1"}}, + "active-1", + "parent-01", ), ( - "chat/completions", {}, + "chat/completions", + {}, {"metadata": {"session_id": "active-1"}, "litellm_metadata": {"session_id": "body-001"}}, - "active-1", "parent-01", + "active-1", + "parent-01", ), ( - "chat/completions", {}, {"metadata": {"session_id": ""}, "litellm_metadata": {"session_id": "body-001"}}, - None, "parent-01", + "chat/completions", + {}, + {"metadata": {"session_id": ""}, "litellm_metadata": {"session_id": "body-001"}}, + None, + "parent-01", ), ( - "messages", {}, {"metadata": {"session_id": "body-001"}, "litellm_metadata": {"session_id": ""}}, - None, "parent-01", + "messages", + {}, + {"metadata": {"session_id": "body-001"}, "litellm_metadata": {"session_id": ""}}, + None, + "parent-01", ), ( - "messages", {}, {"litellm_session_id": "root-001", "litellm_metadata": {"session_id": "active-1"}, - "metadata": {"session_id": "body-001", "user_id": "user_abc_account__session_anthro-1"}}, - "anthro-1", "parent-01", + "messages", + {}, + { + "litellm_session_id": "root-001", + "litellm_metadata": {"session_id": "active-1"}, + "metadata": {"session_id": "body-001", "user_id": "user_abc_account__session_anthro-1"}, + }, + "anthro-1", + "parent-01", + ), + ( + "messages", + {"x-session-id": "child-01"}, + {"metadata": {"user_id": {"session_id": "body-001"}}}, + "child-01", + "parent-01", ), - ("messages", {"x-session-id": "child-01"}, {"metadata": {"user_id": {"session_id": "body-001"}}}, - "child-01", "parent-01"), ("chat/completions", {}, {}, None, "parent-01"), ("responses", {"x-session-id": "bad"}, {}, None, "parent-01"), ("messages", {}, {"metadata": {"session_id": 123}}, None, "parent-01"), @@ -4349,9 +4386,7 @@ def _bedrock_passthrough_data( ), ], ) -def test_bedrock_request_metadata_thread_id_becomes_session_id( - headers: dict[str, str], client_body: dict[str, object] -): +def test_bedrock_request_metadata_thread_id_becomes_session_id(headers: dict[str, str], client_body: dict[str, object]): sent_body: Final = copy.deepcopy(client_body) data: Final = _bedrock_passthrough_data(client_body) @@ -4431,7 +4466,9 @@ def test_bedrock_request_metadata_without_usable_session_id_sets_no_session( @pytest.mark.parametrize( ("body_tags", "header_tags", "expected"), [ - pytest.param({"team": "search"}, {"threadId": "thread-abc-123"}, "thread-abc-123", id="body_without_session_key"), + pytest.param( + {"team": "search"}, {"threadId": "thread-abc-123"}, "thread-abc-123", id="body_without_session_key" + ), pytest.param({"sessionId": "body-session"}, {"sessionId": "header-session"}, "body-session", id="body_wins"), ], ) @@ -8048,6 +8085,7 @@ async def test_add_litellm_data_to_request_refuses_a_model_config_credential_hol ("API-Key", True), ("Cookie", True), ("Proxy-Authorization", True), + ("X-Amz-Security-Token", True), ("anthropic-version", False), ("user-agent", False), ], @@ -11271,9 +11309,7 @@ async def test_header_mapped_user_id_cannot_load_another_users_connection(monkey from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request _per_user_env(monkeypatch) - load = AsyncMock( - return_value={"copilot-cred": {"access_token": "gho_user_a_token", "github_host": "github.com"}} - ) + load = AsyncMock(return_value={"copilot-cred": {"access_token": "gho_user_a_token", "github_host": "github.com"}}) monkeypatch.setattr(upc, "aget_user_provider_credentials", load) router = Router( model_list=[ @@ -11326,9 +11362,7 @@ async def test_per_user_token_resolution_uses_key_identity_without_header_mappin from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request _per_user_env(monkeypatch) - load = AsyncMock( - return_value={"copilot-cred": {"access_token": "gho_user_a_token", "github_host": "github.com"}} - ) + load = AsyncMock(return_value={"copilot-cred": {"access_token": "gho_user_a_token", "github_host": "github.com"}}) monkeypatch.setattr(upc, "aget_user_provider_credentials", load) router = Router( model_list=[ diff --git a/tests/unit/proxy/test_plugin_routes.py b/tests/unit/proxy/test_plugin_routes.py index 7d3fbda14fd..23145f13156 100644 --- a/tests/unit/proxy/test_plugin_routes.py +++ b/tests/unit/proxy/test_plugin_routes.py @@ -12,10 +12,19 @@ Covers three bugs: """ import asyncio +import base64 +import hashlib +import hmac +import json +import os +import time +from typing import Final, NoReturn from unittest.mock import MagicMock +import _hashlib import pytest import respx +from cryptography.fernet import Fernet from starlette.requests import Request import litellm @@ -25,7 +34,12 @@ from litellm.proxy._types import ( PluginConfig, UserAPIKeyAuth, ) -from litellm.proxy.plugin_routes import list_plugins, register_plugins_from_config +from litellm.proxy.plugin_routes import ( + issue_plugin_session_claim, + list_plugins, + register_plugins_from_config, + verify_plugin_session_claim, +) def _admin() -> UserAPIKeyAuth: @@ -320,3 +334,50 @@ def test_plugin_proxy_refuses_non_admin_callers_without_contacting_the_plugin( assert response.status_code == 403 assert response.body == b"Plugin proxy access requires proxy_admin role." assert not upstream.called + + +def test_plugin_session_claim_opens_under_the_documented_hmac_plugin_key(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-salt") + plugin_key: Final = base64.urlsafe_b64encode(hmac.new(b"sk-salt", b"chat-ui", hashlib.sha256).digest()) + + claim: Final[dict[str, object]] = json.loads( + Fernet(plugin_key).decrypt(issue_plugin_session_claim("chat-ui", "user-1", "proxy_admin")) + ) + + assert (claim["plugin"], claim["user_id"], claim["user_role"]) == ("chat-ui", "user-1", "proxy_admin") + + +@pytest.mark.skipif( + not hasattr(_hashlib, "__file__"), reason="this Python links OpenSSL statically, out of ctypes reach" +) +@pytest.mark.usefixtures("fips_mode") +def test_fips_mode_plugin_session_claim_opens_under_the_kbkdf_key_without_os_urandom_or_hmac_new( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def outside_the_fips_provider(*_: object) -> NoReturn: + raise AssertionError("fips_mode must not draw on os.urandom or hmac.new") + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-" + bytes(range(32)).hex()) + monkeypatch.setattr(os, "urandom", outside_the_fips_provider) + monkeypatch.setattr(hmac, "new", outside_the_fips_provider) + monkeypatch.setattr(time, "time", lambda: 1_800_000_000.0) + fixed_input: Final = ( + (1).to_bytes(4, "big") + b"litellm-plugin:session-claim:" + b"\x00" + b"chat-ui" + (256).to_bytes(4, "big") + ) + plugin_key: Final = base64.urlsafe_b64encode(hmac.digest(bytes(range(32)), fixed_input, "sha256")) + token: Final = issue_plugin_session_claim("chat-ui", "user-1", "proxy_admin") + + claim: Final[dict[str, object]] = json.loads(Fernet(plugin_key).decrypt(token)) + + assert (claim["plugin"], claim["user_id"], claim["user_role"]) == ("chat-ui", "user-1", "proxy_admin") + assert verify_plugin_session_claim("chat-ui", token) == claim + + +@pytest.mark.usefixtures("fips_mode") +def test_fips_mode_refuses_a_plugin_session_claim_under_a_salt_key_that_is_not_hex( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-salt") + + with pytest.raises(ValueError, match="LITELLM_SALT_KEY must be at least 32 bytes of hex"): + issue_plugin_session_claim("chat-ui", "user-1", "proxy_admin") 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 1cfb14e6bf9..c62e73232d0 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -918,6 +918,49 @@ async def test_initialize_scheduled_jobs_loads_credentials_only_through_add_depl scheduler.shutdown(wait=False) +@pytest.mark.asyncio +@pytest.mark.parametrize("proxy_batch_write_at", [10, 60]) +async def test_committed_spend_twins_outlive_the_slowest_spend_write_to_the_db( + monkeypatch: pytest.MonkeyPatch, proxy_batch_write_at: int +) -> None: + """A twin that expires before its charges reach the DB seeds again from a total still missing them.""" + monkeypatch.delenv("DISABLE_PRISMA_SCHEMA_UPDATE", raising=False) + monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False) + from apscheduler.schedulers.asyncio import AsyncIOScheduler + + from litellm.constants import SPEND_LOG_QUEUE_MAX_BACKOFF_SECONDS + from litellm.proxy.proxy_server import _BATCH_WRITE_JITTER_MAX_SECONDS, ProxyStartupEvent + from litellm.proxy.utils import ProxyLogging + from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters + + spend: Final = FakeSpendCounters() + mock_proxy_logging: Final = MagicMock(spec=ProxyLogging) + mock_proxy_logging.slack_alerting_instance = MagicMock() + mock_proxy_logging.db_spend_update_writer = MagicMock() + scheduler: Final = AsyncIOScheduler() + try: + with ( + patch("litellm.proxy.proxy_server.proxy_config", _mock_scheduled_proxy_config()), + patch("litellm.proxy.proxy_server.store_model_in_db", False), + patch("litellm.proxy.proxy_server.AsyncIOScheduler", return_value=scheduler), + patch("litellm.proxy.proxy_server.spend_counter_store", spend.store()), + ): + await ProxyStartupEvent.initialize_scheduled_background_jobs( + general_settings={}, + prisma_client=MagicMock(), + proxy_budget_rescheduler_min_time=1, + proxy_budget_rescheduler_max_time=2, + proxy_batch_write_at=proxy_batch_write_at, + proxy_logging_obj=mock_proxy_logging, + ) + finally: + scheduler.shutdown(wait=False) + + slowest_db_write: Final = SPEND_LOG_QUEUE_MAX_BACKOFF_SECONDS + proxy_batch_write_at + _BATCH_WRITE_JITTER_MAX_SECONDS + assert spend.committed_ttl is not None + assert spend.committed_ttl >= slowest_db_write + + @pytest.mark.asyncio async def test_periodic_reload_job_scheduled_without_store_model_in_db(monkeypatch): """ diff --git a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py index f74ff9b1fca..867259a9b3b 100644 --- a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py +++ b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py @@ -561,7 +561,11 @@ class TestPostCallFailureHookDoesNotInventAProviderHandoff: litellm_call_id="pre-handoff-auth-failure", function_id="pre-handoff", ) - request_data: Final[dict[str, object]] = {"litellm_logging_obj": logging_obj, "model": "gpt-4.1", "metadata": {}} + request_data: Final[dict[str, object]] = { + "litellm_logging_obj": logging_obj, + "model": "gpt-4.1", + "metadata": {}, + } proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) proxy_logging_obj.alert_types = [] with ( @@ -1116,6 +1120,44 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens: assert expected_input_cost > 0 assert request_data["response_cost"] == pytest.approx(expected_input_cost) + @pytest.mark.asyncio + async def test_a_total_timeout_that_ran_out_after_dispatch_is_priced_at_input_rate(self): + from datetime import datetime + + from litellm.litellm_core_utils.total_deadline import TotalTimeoutExpired + from litellm.proxy._types import ProxyException + from litellm.types.utils import Usage + + def ran_out() -> ProxyException: + try: + raise TotalTimeoutExpired("the request did not finish within its total_timeout of 2.0s") + except TotalTimeoutExpired as expired: + try: + raise ProxyException( + message=str(expired), type="timeout", param="total_timeout", code=408 + ) from expired + except ProxyException as answered: + return answered + + request_data = { + "litellm_logging_obj": self._logging_obj( + { + "first_api_call_start_time": datetime.now(), + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "count these input tokens please"}], + "call_type": "acompletion", + } + ), + "metadata": {}, + } + await self._run(request_data, original_exception=ran_out()) + + estimated = request_data["combined_usage_object"] + assert isinstance(estimated, Usage) + expected_input_cost, _ = litellm.cost_per_token(model="gpt-3.5-turbo", prompt_tokens=estimated.prompt_tokens) + assert expected_input_cost > 0 + assert request_data["response_cost"] == pytest.approx(expected_input_cost) + @pytest.mark.asyncio async def test_upstream_499_error_is_a_plain_failure_not_a_priced_disconnect(self): """A provider answering HTTP 499 surfaces as litellm.APIError with status_code 499. Only the @@ -2693,7 +2735,12 @@ def test_an_earlier_attempt_on_another_model_is_priced_at_that_models_rate_not_t monkeypatch.setitem( litellm.model_cost, pricing_model, - {"litellm_provider": provider, "mode": "chat", "input_cost_per_token": input_rate, "output_cost_per_token": 1.0}, + { + "litellm_provider": provider, + "mode": "chat", + "input_cost_per_token": input_rate, + "output_cost_per_token": 1.0, + }, ) usage = litellm.Usage(prompt_tokens=24, completion_tokens=0, total_tokens=24) @@ -2926,9 +2973,7 @@ async def test_dropped_stream_failure_attribution_reaches_terminal_spend_row() - request_data: Final = {"litellm_logging_obj": logging_obj} lifted: Final = _failure_fields_to_lift(request_data, provider_billed=True) terminal_kwargs: Final = { - key: value - for key, value in {**request_data, **lifted}.items() - if key != "litellm_logging_obj" + key: value for key, value in {**request_data, **lifted}.items() if key != "litellm_logging_obj" } row: Final = get_logging_payload( kwargs=terminal_kwargs, @@ -3101,9 +3146,7 @@ async def test_latest_dropped_attempt_publishes_its_payload_after_earlier_dispat request_data: Final = {"litellm_logging_obj": logging_obj} lifted: Final = _failure_fields_to_lift(request_data, provider_billed=True) terminal_kwargs: Final = { - key: value - for key, value in {**request_data, **lifted}.items() - if key != "litellm_logging_obj" + key: value for key, value in {**request_data, **lifted}.items() if key != "litellm_logging_obj" } row: Final = get_logging_payload( kwargs=terminal_kwargs, @@ -3158,9 +3201,7 @@ async def test_dropped_attempt_attribution_is_published_while_failure_callback_i request_data: Final = {"litellm_logging_obj": logging_obj} lifted: Final = _failure_fields_to_lift(request_data, provider_billed=True) terminal_kwargs: Final = { - key: value - for key, value in {**request_data, **lifted}.items() - if key != "litellm_logging_obj" + key: value for key, value in {**request_data, **lifted}.items() if key != "litellm_logging_obj" } row: Final = get_logging_payload( kwargs=terminal_kwargs, @@ -3237,9 +3278,7 @@ async def test_failed_attempt_publishes_baseline_reason_after_invalidation( request_data: Final = {"litellm_logging_obj": logging_obj} lifted: Final = _failure_fields_to_lift(request_data, provider_billed=True) terminal_kwargs: Final = { - key: value - for key, value in {**request_data, **lifted}.items() - if key != "litellm_logging_obj" + key: value for key, value in {**request_data, **lifted}.items() if key != "litellm_logging_obj" } row: Final = get_logging_payload( kwargs=terminal_kwargs, @@ -3254,9 +3293,7 @@ async def test_failed_attempt_publishes_baseline_reason_after_invalidation( assert row["prompt_tokens"] == 4 assert row["completion_tokens"] == 1 lifted_payload: Final = TypeAdapter(StandardLoggingPayload).validate_python(lifted["standard_logging_object"]) - callback_payload: Final = TypeAdapter(StandardLoggingPayload).validate_python( - recorder.standard_logging_object - ) + callback_payload: Final = TypeAdapter(StandardLoggingPayload).validate_python(recorder.standard_logging_object) def baseline_reason(payload: StandardLoggingPayload) -> str: encoded: Final = payload["autorouter_baseline_observation"] @@ -3335,9 +3372,7 @@ async def test_queued_stale_attempt_dispatch_does_not_publish_over_current_attem request_data: Final = {"litellm_logging_obj": logging_obj} lifted: Final = _failure_fields_to_lift(request_data, provider_billed=True) terminal_kwargs: Final = { - key: value - for key, value in {**request_data, **lifted}.items() - if key != "litellm_logging_obj" + key: value for key, value in {**request_data, **lifted}.items() if key != "litellm_logging_obj" } row: Final = get_logging_payload( kwargs=terminal_kwargs, @@ -3402,9 +3437,7 @@ async def test_dropped_stream_failure_does_not_overwrite_newer_terminal_payload( request_data: Final = {"litellm_logging_obj": logging_obj} lifted: Final = _failure_fields_to_lift(request_data, provider_billed=True) terminal_kwargs: Final = { - key: value - for key, value in {**request_data, **lifted}.items() - if key != "litellm_logging_obj" + key: value for key, value in {**request_data, **lifted}.items() if key != "litellm_logging_obj" } row: Final = get_logging_payload( kwargs=terminal_kwargs, @@ -3658,7 +3691,6 @@ def test_a_request_that_timed_out_at_a_provider_bills_its_input_at_the_input_rat assert response_cost == pytest.approx(expected_input_cost) - COMPACTION_INSTRUCTIONS = {"role": "system", "content": "answer briefly"} COMPACTED_HISTORY = [ {"role": "user", "content": "a long turn " * 200}, @@ -3735,8 +3767,6 @@ def test_an_earlier_timed_out_responses_attempt_counts_its_input_from_the_newest assert sent == earlier_attempt_prompt_tokens([COMPACTION_INSTRUCTIONS, COMPACTION_ITEM, TURN_AFTER_COMPACTION]) - - @pytest.mark.parametrize( "current_attempt_api_call_start_time", [None, 1700000001.0], @@ -4076,9 +4106,7 @@ def test_create_model_info_response_resolves_mode_through_deployment_model(): ] ) - response = create_model_info_response( - model_id="my-embeddings", provider="openai", llm_router=router - ) + response = create_model_info_response(model_id="my-embeddings", provider="openai", llm_router=router) finally: litellm.model_cost.clear() litellm.model_cost.update(saved_model_cost) @@ -4108,12 +4136,8 @@ def test_create_model_info_response_resolves_model_group_alias_to_target(model_g model_group_alias=model_group_alias, ) - alias_response = create_model_info_response( - model_id="team-embeddings", provider="openai", llm_router=router - ) - target_response = create_model_info_response( - model_id="my-embeddings", provider="openai", llm_router=router - ) + alias_response = create_model_info_response(model_id="team-embeddings", provider="openai", llm_router=router) + target_response = create_model_info_response(model_id="my-embeddings", provider="openai", llm_router=router) assert alias_response["id"] == "team-embeddings" for field in ("mode", "max_input_tokens", "max_output_tokens"): @@ -4196,7 +4220,9 @@ async def test_post_call_failure_hook_redacts_traceback_before_callbacks(monkeyp with patch.object(proxy_logging_obj, "update_request_status", new=AsyncMock()): await proxy_logging_obj.post_call_failure_hook( request_data={"metadata": {}}, - original_exception=HTTPException(status_code=400, detail="Upstream passthrough request failed with status 400"), + original_exception=HTTPException( + status_code=400, detail="Upstream passthrough request failed with status 400" + ), user_api_key_dict=UserAPIKeyAuth(), traceback_str=upstream_traceback, ) @@ -4260,9 +4286,13 @@ def test_mcp_conversion_preserves_request_policy_and_isolates_guardrail_data(buc parent = { "model": "parent-model", bucket: { - "guardrails": ["policy-rule"], "guardrail_config": {"language": "en"}, - "applied_policies": ["parent-policy"], "policy_sources": {"parent-policy": "model"}, - "_guardrail_pipelines": [], "_pipeline_managed_guardrails": ["pipeline-rule"], "tags": ["review"], + "guardrails": ["policy-rule"], + "guardrail_config": {"language": "en"}, + "applied_policies": ["parent-policy"], + "policy_sources": {"parent-policy": "model"}, + "_guardrail_pipelines": [], + "_pipeline_managed_guardrails": ["pipeline-rule"], + "tags": ["review"], }, "guardrails": [{"request-rule": {"extra_body": {"threshold": 0.9}}}], "guardrail_config": {"entities": ["EMAIL_ADDRESS"]}, @@ -4292,13 +4322,26 @@ def test_mcp_conversion_honors_only_authenticated_global_guardrail_opt_outs(opt_ from litellm.responses.mcp.request_context import MCPRequestContext auth = UserAPIKeyAuth(metadata={"opted_out_global_guardrails": ["global-rule"] if opt_out else []}) - context = MCPRequestContext.resolve(kwargs={"metadata": { - "user_api_key_auth": auth, "disable_global_guardrails": True, - "user_api_key_metadata": {"disable_global_guardrails": True}, - }}, tools=None) + context = MCPRequestContext.resolve( + kwargs={ + "metadata": { + "user_api_key_auth": auth, + "disable_global_guardrails": True, + "user_api_key_metadata": {"disable_global_guardrails": True}, + } + }, + tools=None, + ) proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) - kwargs = {"name": "execute", "arguments": {}, "user_api_key_auth": auth, "guardrail_context": context.guardrail_context} - synthetic = proxy_logging._convert_mcp_to_llm_format(proxy_logging._create_mcp_request_object_from_kwargs(kwargs), kwargs) + kwargs = { + "name": "execute", + "arguments": {}, + "user_api_key_auth": auth, + "guardrail_context": context.guardrail_context, + } + synthetic = proxy_logging._convert_mcp_to_llm_format( + proxy_logging._create_mcp_request_object_from_kwargs(kwargs), kwargs + ) guardrail = CustomGuardrail(guardrail_name="global-rule", event_hook="pre_mcp_call", default_on=True) assert guardrail.should_run_guardrail(synthetic, GuardrailEventHooks.pre_mcp_call) is (not opt_out) synthetic["metadata"]["user_api_key_metadata"]["opted_out_global_guardrails"].append("unrelated") @@ -4312,19 +4355,26 @@ def test_mcp_auth_policy_uses_original_request_model(monkeypatch, model, expecte from litellm.types.proxy.policy_engine import Policy, PolicyCondition, PolicyGuardrails registry = policy_registry.PolicyRegistry() - registry._policies = {"model-policy": Policy( - condition=PolicyCondition(model="parent-model"), guardrails=PolicyGuardrails(add=["model-rule"]) - )} + registry._policies = { + "model-policy": Policy( + condition=PolicyCondition(model="parent-model"), guardrails=PolicyGuardrails(add=["model-rule"]) + ) + } registry._initialized = True monkeypatch.setattr(policy_registry, "_policy_registry", registry) monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) kwargs = { - "name": "execute", "arguments": {}, + "name": "execute", + "arguments": {}, "user_api_key_auth": UserAPIKeyAuth(metadata={"policies": ["model-policy"]}), - "guardrail_context": MCPRequestContext.resolve_guardrail_context({"model": model, "guardrails": ["request-rule"]}), + "guardrail_context": MCPRequestContext.resolve_guardrail_context( + {"model": model, "guardrails": ["request-rule"]} + ), } - synthetic = proxy_logging._convert_mcp_to_llm_format(proxy_logging._create_mcp_request_object_from_kwargs(kwargs), kwargs) + synthetic = proxy_logging._convert_mcp_to_llm_format( + proxy_logging._create_mcp_request_object_from_kwargs(kwargs), kwargs + ) assert ("model-rule" in synthetic["metadata"]["guardrails"]) is expected assert "request-rule" in synthetic["metadata"]["guardrails"] @@ -4365,7 +4415,5 @@ def test_is_client_error_exception_existing_behaviour_unchanged() -> None: assert _is_client_error_exception(HTTPException(status_code=400, detail="x")) is True assert _is_client_error_exception(HTTPException(status_code=500, detail="x")) is False - assert _is_client_error_exception( - ProxyException(message="x", type="x", param=None, code="429") - ) is True + assert _is_client_error_exception(ProxyException(message="x", type="x", param=None, code="429")) is True assert _is_client_error_exception(ValueError("x")) is False diff --git a/tests/unit/proxy/test_response_polling_pre_call_checks.py b/tests/unit/proxy/test_response_polling_pre_call_checks.py index c6e4301a501..4228c1ff8f8 100644 --- a/tests/unit/proxy/test_response_polling_pre_call_checks.py +++ b/tests/unit/proxy/test_response_polling_pre_call_checks.py @@ -6,6 +6,8 @@ BEFORE a polling ID is created, so rate-limited requests get a synchronous error instead of a polling ID that immediately fails. """ +from collections.abc import Mapping +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -83,7 +85,7 @@ class TestSkipPreCallLogic: ): try: await processor.base_process_llm_request( - request=MagicMock(spec=Request), + request=MagicMock(spec=Request, headers={}), fastapi_response=MagicMock(spec=Response), user_api_key_dict=UserAPIKeyAuth(), route_type="aresponses", @@ -176,3 +178,77 @@ class TestPollingEndpointPreCallGuard: assert exc_info.value.status_code == 429 generate_polling_id_mock.assert_not_called() + + +async def _polled_background_response( + body: Mapping[str, object], headers: Mapping[str, str] +) -> tuple[MagicMock, Mapping[str, object]]: + from litellm.proxy.response_api_endpoints.endpoints import responses_api + from litellm.proxy.response_polling.polling_handler import ResponsePollingHandler + + generate_polling_id: Final = MagicMock(return_value="litellm_poll_test") + pre_call_saw: Final[list[Mapping[str, object]]] = [] + + async def pre_call(processor: ProxyBaseLLMRequestProcessing, **_: object) -> tuple[Mapping[str, object], MagicMock]: + pre_call_saw.append(dict(processor.data)) + return processor.data, MagicMock() + + request: Final = MagicMock(spec=Request) + request.headers = headers + proxy_server_patches: Final[Mapping[str, object]] = { + "read_request_body": AsyncMock(return_value=dict(body)), + "general_settings": {}, + "llm_router": MagicMock(), + "native_background_mode": None, + "polling_cache_ttl": 3600, + "polling_via_cache_enabled": True, + "proxy_config": MagicMock(), + "proxy_logging_obj": AsyncMock(), + "redis_usage_cache": AsyncMock(), + "user_api_base": None, + "user_max_tokens": None, + "user_model": None, + "user_request_timeout": None, + "user_temperature": None, + "version": "1.0.0", + } + with ( + patch.multiple("litellm.proxy.proxy_server", **proxy_server_patches), + patch("litellm.proxy.response_polling.polling_handler.should_use_polling_for_request", return_value=True), + patch.object(ProxyBaseLLMRequestProcessing, "common_processing_pre_call_logic", pre_call), + patch.object( + ProxyBaseLLMRequestProcessing, + "handle_llm_api_exception", + new_callable=AsyncMock, + side_effect=lambda e, **_: e, + ), + patch("litellm.proxy.response_api_endpoints.endpoints.raise_if_required_body_param_missing"), + patch.object(ResponsePollingHandler, "generate_polling_id", generate_polling_id), + patch("asyncio.create_task"), + patch.object(ResponsePollingHandler, "create_initial_state", new_callable=AsyncMock, return_value=MagicMock()), + ): + await responses_api( + request=request, fastapi_response=MagicMock(spec=Response), user_api_key_dict=MagicMock(spec=UserAPIKeyAuth) + ) + return generate_polling_id, pre_call_saw[0] if pre_call_saw else {} + + +@pytest.mark.asyncio +async def test_a_body_total_timeout_on_a_polled_background_response_is_refused_before_a_polling_id_is_issued() -> None: + from litellm.proxy._types import ProxyException + + with pytest.raises(ProxyException) as refused: + _ = await _polled_background_response({"model": "gpt-4", "background": True, "total_timeout": 10}, {}) + + assert refused.value.code == "400" + assert refused.value.param == "total_timeout" + + +@pytest.mark.asyncio +async def test_a_header_total_timeout_on_a_polled_background_response_is_dropped_and_polling_proceeds() -> None: + generate_polling_id, pre_call_saw = await _polled_background_response( + {"model": "gpt-4", "background": True, "total_timeout": "ignored"}, {"x-litellm-total-timeout": "10"} + ) + + generate_polling_id.assert_called_once() + assert "total_timeout" not in pre_call_saw diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index cb903fb134a..7f384d8e6bf 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -1,30 +1,26 @@ import asyncio import contextlib import gc +import json import threading import weakref -import httpx -import json -import litellm -import pytest -import redis -import respx - - from collections.abc import Awaitable, Callable, Coroutine, Iterator, Mapping from datetime import datetime from types import SimpleNamespace -from uuid import uuid4 from typing import Final from unittest.mock import MagicMock +from uuid import uuid4 -from litellm.router_api import RouterSettings -from tests.unit.proxy.router_entry_point_stubs import route_entry_points_to_stubs - +import httpx +import pytest +import redis +import respx from fakeredis import TcpFakeServer from fastapi import HTTPException +import litellm from litellm.caching.caching import DualCache +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.route_llm_request import ( @@ -36,11 +32,13 @@ from litellm.proxy.route_llm_request import ( user_config_routers, ) from litellm.proxy.utils import ProxyLogging +from litellm.router_api import RouterSettings from litellm.router_strategy.adaptive_router.config import ( ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY, ADAPTIVE_ROUTER_RESPONSE_HEADER, ) from litellm.types.router import RouterErrors, RouterGeneralSettings, RouterRateLimitError +from tests.unit.proxy.router_entry_point_stubs import route_entry_points_to_stubs @pytest.mark.parametrize( @@ -78,7 +76,6 @@ async def test_route_request_dynamic_credentials(route_type, required_body_param @pytest.mark.asyncio async def test_route_request_proxy_admin_can_call_all_team_scoped_deployments_without_team_id(): import litellm - from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth router = litellm.Router( @@ -670,7 +667,6 @@ async def test_route_request_forwards_mock_params_when_opted_in(monkeypatch): """End-to-end through ``route_request``: with the opt-in set the param reaches the router, which is what makes a fallback drill possible.""" import litellm.proxy.proxy_server as proxy_server - from litellm.proxy.route_llm_request import MOCK_TESTING_CONFIG_KEY monkeypatch.setattr( @@ -775,10 +771,12 @@ async def test_route_request_realtime_wildcard_model_resolves_credentials( POST /realtime/client_secrets with a request model like openai/gpt-realtime must match an openai/* deployment and forward its api_key upstream. """ - import httpx - import litellm from unittest.mock import AsyncMock, patch + import httpx + + import litellm + monkeypatch.delenv("OPENAI_API_KEY", raising=False) router = litellm.Router( @@ -813,10 +811,12 @@ async def test_route_request_realtime_team_scoped_model_resolves_credentials( Team-scoped deployments (team_public_model_name) must be selected when user_api_key_team_id is present, same as /chat/completions. """ - import httpx - import litellm from unittest.mock import AsyncMock, patch + import httpx + + import litellm + monkeypatch.delenv("OPENAI_API_KEY", raising=False) router = litellm.Router( @@ -859,10 +859,12 @@ async def test_route_request_realtime_litellm_credential_name_resolves_api_key( litellm_credential_name on a wildcard deployment must resolve to the stored api_key when routing acreate_realtime_client_secret through the router. """ - import httpx - import litellm from unittest.mock import AsyncMock, patch + import httpx + + import litellm + monkeypatch.delenv("OPENAI_API_KEY", raising=False) router = litellm.Router( @@ -894,9 +896,9 @@ async def test_route_request_realtime_unresolvable_model_raises_not_found( An unknown model must not silently fall through to litellm with an empty OPENAI_API_KEY env var. """ - import litellm from unittest.mock import AsyncMock, patch + import litellm from litellm.proxy.route_llm_request import ProxyModelNotFoundError monkeypatch.delenv("OPENAI_API_KEY", raising=False) @@ -926,10 +928,12 @@ async def test_route_request_realtime_calls_resolves_api_base(monkeypatch): non-default (self-hosted / proxied) OpenAI endpoint is honored, instead of defaulting to https://api.openai.com. """ - import httpx - import litellm from unittest.mock import AsyncMock, patch + import httpx + + import litellm + monkeypatch.delenv("OPENAI_API_KEY", raising=False) monkeypatch.delenv("OPENAI_API_BASE", raising=False) @@ -969,10 +973,12 @@ async def test_route_request_realtime_transcription_session_resolves_credentials /realtime/transcription_sessions must resolve credentials through the router (wildcard deployment) rather than falling back to an empty OPENAI_API_KEY. """ - import httpx - import litellm from unittest.mock import AsyncMock, patch + import httpx + + import litellm + monkeypatch.delenv("OPENAI_API_KEY", raising=False) router = litellm.Router( @@ -1399,7 +1405,9 @@ async def test_route_request_applies_override_fallbacks_when_request_carries_cre @pytest.mark.parametrize( "router_only_setting", [ - pytest.param({"model_group_retry_policy": {"gpt-4": {"TimeoutErrorRetries": 1}}}, id="model-group-retry-policy"), + pytest.param( + {"model_group_retry_policy": {"gpt-4": {"TimeoutErrorRetries": 1}}}, id="model-group-retry-policy" + ), pytest.param({"context_window_fallbacks": [{"gpt-4": ["gpt-3.5-turbo"]}]}, id="context-window-fallbacks"), pytest.param({"content_policy_fallbacks": [{"gpt-4": ["gpt-3.5-turbo"]}]}, id="content-policy-fallbacks"), ], @@ -2102,11 +2110,12 @@ async def test_route_request_read_through_disabled_without_store_model_in_db(mon @pytest.mark.asyncio async def test_route_request_read_through_supplies_db_model_default_for_missing_param(monkeypatch): - import litellm - import litellm.proxy.proxy_server as proxy_server from types import SimpleNamespace from unittest.mock import AsyncMock, patch + import litellm + import litellm.proxy.proxy_server as proxy_server + model_name = "e2e-db-only-max-tokens-default" router = litellm.Router( model_list=[{"model_name": "some-other-model", "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"}}] @@ -2545,11 +2554,38 @@ def _mock_deployment(model_name: str, **deployment_limits: int) -> dict[str, obj } +async def _settle_router_accounting() -> None: + try: + await GLOBAL_LOGGING_WORKER.flush() + finally: + await GLOBAL_LOGGING_WORKER.stop() + while pending := {task for task in asyncio.all_tasks() - {asyncio.current_task()} if not _is_periodic_sync(task)}: + _, still_running = await asyncio.wait(pending, timeout=5) + assert not still_running, still_running + + +def _is_periodic_sync(task: asyncio.Task[object]) -> bool: + return task.get_coro().__qualname__.endswith(".periodic_sync_in_memory_spend_with_redis") + + async def _user_config_request_outcome( redis_port: int | None, routing_settings: Mapping[str, object], deployment_limits: Mapping[str, int], logging_obj: LiteLLMLoggingObj | None, +) -> str: + outcome: Final = await _unsettled_user_config_request_outcome( + redis_port, routing_settings, deployment_limits, logging_obj + ) + await _settle_router_accounting() + return outcome + + +async def _unsettled_user_config_request_outcome( + redis_port: int | None, + routing_settings: Mapping[str, object], + deployment_limits: Mapping[str, int], + logging_obj: LiteLLMLoggingObj | None, ) -> str: data: Final = { "model": "rate-limited", @@ -2566,8 +2602,6 @@ async def _user_config_request_outcome( await (await route_request(data, None, None, "acompletion")) except (litellm.RateLimitError, RouterRateLimitError): return "429" - finally: - await asyncio.sleep(0.3) return "200" @@ -2606,7 +2640,7 @@ async def _proxy_router_outcomes( except (litellm.RateLimitError, RouterRateLimitError): return "429" finally: - await asyncio.sleep(0.3) + await _settle_router_accounting() return "200" try: @@ -2636,7 +2670,9 @@ async def test_a_user_config_adaptive_request_answers_with_the_model_it_picked() }, { **_mock_deployment("smart"), - "model_info": {"adaptive_router_preferences": {"quality_tier": 3, "strengths": ["code_generation"]}}, + "model_info": { + "adaptive_router_preferences": {"quality_tier": 3, "strengths": ["code_generation"]} + }, }, ] }, @@ -2700,16 +2736,21 @@ async def test_concurrent_user_config_requests_each_count_once_against_the_tpm_l tpm_for_three_requests: Final = await _tpm_for_three_requests() def request() -> Coroutine[None, None, str]: - return _user_config_request_outcome( + return _unsettled_user_config_request_outcome( fake_redis_port, {"routing_strategy": "usage-based-routing-v2"}, {"tpm": tpm_for_three_requests}, _proxy_logging_obj() if arrives_with_logging_obj else None, ) + async def settled(*requests: Coroutine[None, None, str]) -> list[str]: + outcomes: Final = await asyncio.gather(*requests) + await _settle_router_accounting() + return outcomes + await _wait_out_a_rate_limit_minute_rollover() - concurrent_outcomes: Final = await asyncio.gather(request(), request()) - later_outcomes: Final = [await request() for _ in range(2)] + concurrent_outcomes: Final = await settled(request(), request()) + later_outcomes: Final = [*(await settled(request())), *(await settled(request()))] assert [*concurrent_outcomes, *later_outcomes] == ["200", "200", "200", "429"] @@ -2717,12 +2758,13 @@ async def test_concurrent_user_config_requests_each_count_once_against_the_tpm_l @pytest.mark.asyncio async def test_overlapping_user_config_requests_share_one_rpm_limit(fake_redis_port: int) -> None: def request() -> Coroutine[None, None, str]: - return _user_config_request_outcome( + return _unsettled_user_config_request_outcome( fake_redis_port, {"routing_strategy": "usage-based-routing-v2"}, {"rpm": 2}, _proxy_logging_obj() ) await _wait_out_a_rate_limit_minute_rollover() outcomes: Final = await asyncio.gather(*(request() for _ in range(4))) + await _settle_router_accounting() assert outcomes.count("200") <= 2, outcomes assert "429" in outcomes, outcomes @@ -2747,13 +2789,14 @@ async def test_a_proxy_router_on_the_same_strategy_leaves_the_user_config_limit_ async def outcomes() -> list[str]: first_instance: Final = [await request() for _ in range(2)] user_config_routers.clear() + await _settle_router_accounting() second_instance: Final = [await request() for _ in range(3)] return [*first_instance, *second_instance] await _wait_out_a_rate_limit_minute_rollover() alone: Final = await outcomes() user_config_routers.clear() - await asyncio.sleep(0.3) + await _settle_router_accounting() redis.Redis(host="127.0.0.1", port=fake_redis_port).flushall() proxy_router: Final = litellm.Router( routing_strategy=routing_strategy, @@ -2781,7 +2824,7 @@ async def test_user_config_requests_stay_under_the_proxy_provider_budget( ) try: await proxy_router.acompletion(model="proxy-model", messages=[{"role": "user", "content": "hi"}]) - await asyncio.sleep(0.3) + await _settle_router_accounting() with pytest.raises(ValueError, match=RouterErrors.no_deployments_with_provider_budget_routing.value): await _user_config_request_outcome(None, user_config_settings, {}, _proxy_logging_obj()) finally: @@ -3092,3 +3135,47 @@ def test_fallback_probe_skips_only_a_loaded_model_that_refuses_the_route() -> No assert skipped("embed-only") is True assert skipped("chat-fallback") is False assert skipped("added-on-another-worker") is False + + +def _router_answering_hi() -> "litellm.Router": + import litellm + + return litellm.Router( + model_list=[{"model_name": "g", "litellm_params": {"model": "openai/gpt-5.4-nano", "mock_response": "hi"}}] + ) + + +@pytest.mark.parametrize("flag", ["stream", "background"]) +@pytest.mark.asyncio +async def test_a_request_a_hook_made_stream_or_run_in_the_background_is_refused_under_a_body_total_timeout( + flag: str, +) -> None: + from litellm.litellm_core_utils.total_deadline import TotalTimeoutNotApplicable, total_deadline + + async with total_deadline(60, from_body=True): + with pytest.raises(TotalTimeoutNotApplicable): + await route_request( + data={"model": "g", "messages": [{"role": "user", "content": "hi"}], flag: True}, + llm_router=_router_answering_hi(), + user_model=None, + route_type="acompletion", + ) + + +@pytest.mark.asyncio +async def test_a_request_a_hook_turned_streaming_runs_unbounded_under_a_header_total_timeout() -> None: + from litellm.litellm_core_utils.total_deadline import current_total_deadline, total_deadline + + async with total_deadline(60, from_body=False): + llm_call: Final = await route_request( + data={"model": "g", "messages": [{"role": "user", "content": "hi"}], "stream": True}, + llm_router=_router_answering_hi(), + user_model=None, + route_type="acompletion", + ) + stream: Final = await llm_call + chunks: Final = [chunk async for chunk in stream] + deadline: Final = current_total_deadline() + + assert chunks + assert deadline is not None and deadline.allows_wait(10**9) diff --git a/tests/unit/proxy/video_endpoints/test_utils.py b/tests/unit/proxy/video_endpoints/test_utils.py index 01a2b7c2840..fff88ab462f 100644 --- a/tests/unit/proxy/video_endpoints/test_utils.py +++ b/tests/unit/proxy/video_endpoints/test_utils.py @@ -20,8 +20,10 @@ import pytest from litellm.router import Router +from litellm.proxy._types import UserAPIKeyAuth from litellm.types.videos.main import DecodedVideoId from litellm.proxy.video_endpoints.utils import ( + authorized_video_model, encode_character_id_in_response, encode_video_id_in_response, extract_model_from_target_model_names, @@ -548,6 +550,33 @@ def _wildcard_router() -> Router: ) +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("enforce_fallback_model_access", "allowed_model"), + [(False, "sora-pair"), (True, "sora-backup")], +) +async def test_authorized_video_model_keeps_fallback_routing_after_the_key_check( + monkeypatch: pytest.MonkeyPatch, enforce_fallback_model_access: bool, allowed_model: str +) -> None: + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", + {"enforce_fallback_model_access": enforce_fallback_model_access}, + ) + decoded: Final = decode_video_id_with_provider( + encode_video_id_with_provider( + "video_raw", "azure", "sora-2", deployment_id="dep-sora-backup", requested_model="sora-pair" + ) + ) + token: Final = UserAPIKeyAuth(api_key="hashed", models=[allowed_model]) + + assert await authorized_video_model(decoded, _fallback_router(), token) == "sora-backup" + + +@pytest.mark.asyncio +async def test_authorized_video_model_without_a_router_has_no_routing_model() -> None: + assert await authorized_video_model(_decoded("sora-2", "dep-sora-backup"), None, UserAPIKeyAuth()) is None + + @pytest.mark.parametrize( ("requested_model", "answering_deployment", "expected"), [ diff --git a/tests/unit/redis/stores/spend_counters_double.py b/tests/unit/redis/stores/spend_counters_double.py index 1f3c4b94d45..c50133f7076 100644 --- a/tests/unit/redis/stores/spend_counters_double.py +++ b/tests/unit/redis/stores/spend_counters_double.py @@ -35,9 +35,22 @@ def _is_number(value: object) -> bool: return value is None or (isinstance(value, (int, float)) and not isinstance(value, bool)) +_TWIN_SCOPES: Final = frozenset({"key", "team", "user", "org", "end_user"}) +_WINDOWED_SCOPES: Final = frozenset({"key", "team"}) + + +def _native_store_takes_twin(twin: str) -> bool: + scope, _, counter_id = twin.partition(":") + _, separator, window = counter_id.rpartition(":window:") + return scope in _TWIN_SCOPES and not (scope in _WINDOWED_SCOPES and separator and window) + + def _check(counter: str) -> None: if not counter.startswith("spend:"): raise ValueError(f"{counter!r} is not a spend counter") + twin: Final = counter.removeprefix("spend:committed:") + if twin != counter and not _native_store_takes_twin(twin): + raise ValueError(f"{counter!r} is not a spend counter") @dataclass @@ -53,6 +66,11 @@ class DeclaredRead: return self.owner.mget(self.counters) +@dataclass +class FakeSettlement: + settled: bool = False + + @dataclass class FakeSpendCounters: """``redis`` is None for a worker without Redis, and ``redis_default`` answers reads of a key it lacks. ``fail`` @@ -68,11 +86,12 @@ class FakeSpendCounters: failure: str = "transport" redis_default: float | None = None redis_ttl: float = 60 + committed_ttl: float | None = None batch_reads: bool = True read_hook: Callable[[str], Awaitable[None]] | None = None on_delete: Callable[[tuple[str, ...]], None] | None = None deletes: list[tuple[str, ...]] = field(default_factory=list) - post_call: list[tuple[object, tuple[SpendIncrement, ...]]] = field(default_factory=list) + post_call: list[tuple[object, tuple[SpendIncrement, ...], FakeSettlement]] = field(default_factory=list) declared: list[DeclaredRead] = field(default_factory=list) def store(self) -> SpendCounterStore: @@ -86,6 +105,9 @@ class FakeSpendCounters: def redis_ttl_seconds(self) -> float | None: return None if self.redis is None else self.redis_ttl + def set_committed_ttl(self, ttl_seconds: float, /) -> None: + self.committed_ttl = ttl_seconds + def attach(self, client: NativeRedisClient, /) -> bool: if self.redis is not None: return False @@ -139,17 +161,19 @@ class FakeSpendCounters: self.commands.append(f"PIPELINE {' '.join(counter for counter, _ in increments)}") return _done([self._incr(counter, amount) for counter, amount in increments]) - def increment_many_post_call(self, request: object, increments: Sequence[SpendIncrement]) -> bool: + def increment_many_post_call(self, request: object, increments: Sequence[SpendIncrement]) -> FakeSettlement | None: if self.redis is None: - return False - self.post_call.append((request, tuple(increments))) - return True + return None + settlement: Final = FakeSettlement() + self.post_call.append((request, tuple(increments), settlement)) + return settlement def settle_post_call(self) -> None: """What the native task does once the post-call batch goes out.""" pending: Final = tuple(self.post_call) self.post_call.clear() - for _, increments in pending: + for _, increments, settlement in pending: + settlement.settled = True if self.fail: for counter, _ in increments: self.memory.pop(counter, None) diff --git a/tests/unit/redis/stores/test_spend_counters.py b/tests/unit/redis/stores/test_spend_counters.py index c54dc6cc81e..1d39cffc933 100644 --- a/tests/unit/redis/stores/test_spend_counters.py +++ b/tests/unit/redis/stores/test_spend_counters.py @@ -10,6 +10,7 @@ from litellm.redis.stores.connection import RedisStores, connect from litellm.redis.stores.spend_counters import ( SPEND_COUNTERS_TARGET, SpendCounterStore, + committed_counter, end_user_counter, is_unavailable, key_counter, @@ -28,7 +29,7 @@ from litellm.redis.store_settings import RedisStoreSettings from litellm.rust_bridge._native import NativeRedisClient, NativeRedisStoreError, NativeSpendCounters, SpendRead from tests.unit.caching.test_redis_batch import FakeNatives, fake_send -from .spend_counters_double import FakeSpendCounters +from .spend_counters_double import FakeSettlement, FakeSpendCounters _KEYS: Final = ( ( @@ -46,6 +47,11 @@ _KEYS: Final = ( (project_counter("proj-1"), "spend:project:proj-1"), (model_access_group_counter("gpt-group"), "spend:model_access_group:gpt-group"), (shadow_eval_counter("job-7"), "spend:shadow_eval:job-7"), + (committed_counter(key_counter("88dc28d0")), "spend:committed:key:88dc28d0"), + (committed_counter(team_counter("team-a")), "spend:committed:team:team-a"), + (committed_counter(user_counter("user@example.com")), "spend:committed:user:user@example.com"), + (committed_counter(org_counter("org-1")), "spend:committed:org:org-1"), + (committed_counter(end_user_counter("customer é 1")), "spend:committed:end_user:customer é 1"), ) @@ -81,7 +87,17 @@ def test_a_read_decodes_what_redis_cache_decodes(wire: str, spend: float) -> Non assert type(decoded) is type(spend) -@pytest.mark.parametrize("key", ["spend_db_floor:spend:key:k", "virtual_key_spend:k:gpt:1d", "spend:unknown:k"]) +@pytest.mark.parametrize( + "key", + [ + "spend_db_floor:spend:key:k", + "virtual_key_spend:k:gpt:1d", + "spend:unknown:k", + "spend:committed:key:k:window:1d", + "spend:committed:tag:prod", + "spend:committed:committed:key:k", + ], +) def test_the_native_store_refuses_keys_outside_the_spend_counter_keyspace(key: str) -> None: with pytest.raises(ValueError, match="a spend counter is a spend:\* key"): NativeSpendCounters(None).get_local(key) @@ -153,7 +169,7 @@ async def test_post_call_increments_go_on_the_requests_native_post_call_batch_un targets: Final[list[str | None]] = [] class Recording(FakeSpendCounters): - def increment_many_post_call(self, request: object, increments: object) -> bool: + def increment_many_post_call(self, request: object, increments: object) -> FakeSettlement | None: targets.append(current_service_target()) return super().increment_many_post_call(request, increments) # pyright: ignore[reportArgumentType] # recorded as given @@ -165,8 +181,8 @@ async def test_post_call_increments_go_on_the_requests_native_post_call_batch_un inside: Final = store.increment_many_post_call((("spend:key:k", 1.0),)) empty: Final = store.increment_many_post_call(()) - assert (outside, inside, empty) == (False, True, False) - assert spend.post_call == [(natives.built[0], (("spend:key:k", 1.0),))] + assert (outside, empty) == (None, None) + assert spend.post_call == [(natives.built[0], (("spend:key:k", 1.0),), inside)] assert targets == [SPEND_COUNTERS_TARGET] diff --git a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index d5860b7519b..31ab4c6ac8c 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -5949,12 +5949,12 @@ async def test_bridge_rejects_untranslatable_tool_choice_with_a_400(stream: bool model="anthropic/claude-haiku-4-5", input="Which fruit is red?", tools=[{"type": "function", "name": "lookup_fruit", "parameters": {"type": "object"}}], - tool_choice={"type": "file_search"}, + tool_choice={"type": "web_search_tool"}, stream=stream, api_key="sk-unused", ) assert exc_info.value.status_code == 400 - assert "tool_choice={'type': 'file_search'}" in str(exc_info.value) + assert "tool_choice={'type': 'web_search_tool'}" in str(exc_info.value) def test_cache_control_on_user_message(): diff --git a/tests/unit/responses/test_streaming_iterator_error_events.py b/tests/unit/responses/test_streaming_iterator_error_events.py index 0acf5d9ce62..29cb0412422 100644 --- a/tests/unit/responses/test_streaming_iterator_error_events.py +++ b/tests/unit/responses/test_streaming_iterator_error_events.py @@ -959,3 +959,73 @@ async def test_responses_websocket_turn_failure_reports_the_deployment_it_was_se assert failure is not None _assert_names_the_azure_deployment(failure, reported_api_base) + + +_BEDROCK_SERVICE_UNAVAILABLE_CHUNK: Final = { + "error": { + "message": "The server had an error while processing your request. Sorry about that!", + "type": "server_error", + "code": "service_unavailable", + } +} + + +def _real_config_iterator(provider: str, model: str, events: tuple[Mapping[str, object], ...]): + payload: Final = b"".join(b"data: " + json.dumps(event).encode() + b"\n\n" for event in events) + + async def aiter_bytes(): + yield payload + + response: Final = Mock() + response.headers = {} + response.aiter_bytes = aiter_bytes + logging_obj: Final = Mock(spec=LiteLLMLoggingObj) + logging_obj.model_call_details = {"litellm_params": {}} + logging_obj.completion_start_time = None + config: Final = ProviderConfigManager.get_provider_responses_api_config( + provider=litellm.LlmProviders(provider), model=model + ) + assert config is not None + return ResponsesAPIStreamingIterator( + response=response, + model=model, + responses_api_provider_config=config, + logging_obj=logging_obj, + custom_llm_provider=provider, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("provider", "model"), + [ + ("bedrock", "us.openai.gpt-5.6-sol"), + ("openai", "gpt-5.5"), + ("volcengine", "doubao-seed-1-6"), + ("bigmodel", "glm-5"), + ], +) +async def test_untyped_mid_stream_error_chunk_raises_the_upstream_error(provider: str, model: str): + iterator: Final = _real_config_iterator( + provider, + model, + ( + { + "type": "response.output_text.delta", + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": "partial ", + "sequence_number": 1, + }, + _BEDROCK_SERVICE_UNAVAILABLE_CHUNK, + ), + ) + stream: Final = aiter(iterator) + assert (await anext(stream)).type == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA + + with pytest.raises(MidStreamFallbackError) as exc_info: + await anext(stream) + assert "The server had an error while processing your request" in str(exc_info.value.original_exception) + assert exc_info.value.status_code >= 500 + assert exc_info.value.generated_content == "partial " diff --git a/tests/unit/router_strategy/test_least_busy.py b/tests/unit/router_strategy/test_least_busy.py index f156a5ee405..98dadb8c40e 100644 --- a/tests/unit/router_strategy/test_least_busy.py +++ b/tests/unit/router_strategy/test_least_busy.py @@ -16,11 +16,11 @@ import respx import litellm from litellm import Router - from litellm.caching.caching import DualCache from litellm.caching.in_memory_cache import DEFAULT_MAX_SIZE_IN_MEMORY, InMemoryCache from litellm.caching.redis_cache import RedisCircuitBreakerOpenError from litellm.constants import LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL +from litellm.litellm_core_utils.attempt_settlement import routed_attempt from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.router_strategy.least_busy import ( IN_FLIGHT_COUNT_TTL_SECONDS, @@ -28,6 +28,7 @@ from litellm.router_strategy.least_busy import ( LeastBusyLoggingHandler, owning_picks, picks_reserve, + release_started_attempt, release_unstarted_picks, ) from litellm.types.utils import CallTypes, ModelResponse @@ -287,7 +288,9 @@ class OpenBreakerRedis(SharedRedisCounters): @pytest.mark.asyncio -async def test_an_open_circuit_breaker_falls_back_without_a_warning_per_request(caplog: pytest.LogCaptureFixture) -> None: +async def test_an_open_circuit_breaker_falls_back_without_a_warning_per_request( + caplog: pytest.LogCaptureFixture, +) -> None: worker: Final = _worker(OpenBreakerRedis()) with caplog.at_level(logging.DEBUG, logger="LiteLLM Router"): @@ -463,9 +466,7 @@ async def test_router_picks_least_busy_deployment_and_completion_restores_counts if use_async else router.get_available_deployment(model=ROUTER_GROUP, messages=None) ) - response: Final = await router.acompletion( - model=ROUTER_GROUP, messages=[{"role": "user", "content": "hi"}] - ) + response: Final = await router.acompletion(model=ROUTER_GROUP, messages=[{"role": "user", "content": "hi"}]) for _ in range(10): await asyncio.sleep(0) await GLOBAL_LOGGING_WORKER.flush() @@ -604,7 +605,9 @@ def test_two_threads_picking_at_once_reserve_different_deployments() -> None: worker: Final = LeastBusyLoggingHandler(router_cache=DualCache(), reservation_store=store) with store.first_reads_meet(), ThreadPoolExecutor(max_workers=2) as pool: - picks: Final = tuple(pool.map(lambda _: worker.get_available_deployments(GROUP, HEALTHY, reserve=True), range(2))) + picks: Final = tuple( + pool.map(lambda _: worker.get_available_deployments(GROUP, HEALTHY, reserve=True), range(2)) + ) assert sorted(pick["model_info"]["id"] for pick in picks if pick is not None) == ["dep-a", "dep-b"] @@ -626,7 +629,6 @@ def test_a_reservation_and_a_release_on_one_deployment_at_once_both_count() -> N assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B - class _SlowLocalRead(InMemoryCache): """Answers a count read with the counts from when the read started, once another thread has begun its request.""" @@ -706,7 +708,9 @@ async def test_looking_up_a_pass_through_deployment_does_not_hold_it_for_later_p for _ in range(2) ] sync_picks: Final = [ - router.get_available_deployment_for_pass_through(ROUTER_GROUP, request_kwargs={"metadata": {}})["model_info"]["id"] + router.get_available_deployment_for_pass_through(ROUTER_GROUP, request_kwargs={"metadata": {}})["model_info"][ + "id" + ] for _ in range(2) ] @@ -733,7 +737,9 @@ def test_a_cache_hit_leaves_requests_in_flight_on_its_deployment_counted() -> No assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B -@pytest.mark.parametrize("outcome", [LeastBusyLoggingHandler.log_success_event, LeastBusyLoggingHandler.log_failure_event]) +@pytest.mark.parametrize( + "outcome", [LeastBusyLoggingHandler.log_success_event, LeastBusyLoggingHandler.log_failure_event] +) def test_a_call_that_ends_before_reaching_the_provider_frees_its_pick( outcome: Callable[[LeastBusyLoggingHandler, dict[str, object], None, None, None], None], ) -> None: @@ -755,6 +761,16 @@ def test_a_call_the_proxy_rejected_after_logging_its_start_leaves_no_count_behin assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A +def test_a_proxy_error_logged_on_a_request_whose_failure_was_already_logged_leaves_no_count_behind() -> None: + worker: Final = _worker(None) + + worker.log_pre_api_call( + model="m", messages=[], kwargs={**_call_kwargs("dep-a"), LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL: True} + ) + + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + + def test_a_call_logged_under_its_callers_group_name_frees_the_pick_made_for_another_group() -> None: worker: Final = _worker(None) assert worker.get_available_deployments("tier-group", HEALTHY, reserve=True) is DEPLOYMENT_A @@ -916,7 +932,9 @@ _REJECTED_BATCH: Final = {"error": {"message": "Invalid input file", "type": "in @pytest.mark.asyncio -@pytest.mark.parametrize("provider_response", [httpx.Response(200, json=_BATCH), httpx.Response(400, json=_REJECTED_BATCH)]) +@pytest.mark.parametrize( + "provider_response", [httpx.Response(200, json=_BATCH), httpx.Response(400, json=_REJECTED_BATCH)] +) async def test_a_batch_create_does_not_keep_its_deployment_busy( provider_response: httpx.Response, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ) -> None: @@ -926,7 +944,10 @@ async def test_a_batch_create_does_not_keep_its_deployment_busy( with contextlib.suppress(openai.BadRequestError): await router.acreate_batch( - model=ROUTER_GROUP, completion_window="24h", endpoint="/v1/chat/completions", input_file_id="file-least-busy" + model=ROUTER_GROUP, + completion_window="24h", + endpoint="/v1/chat/completions", + input_file_id="file-least-busy", ) await _drain_logging() @@ -1013,7 +1034,6 @@ def test_a_sync_text_completion_does_not_keep_its_deployment_busy( assert _router_pick(router, TEXT_GROUP) == _called_deployment(route) - @pytest.mark.parametrize("call_identity", _CALL_IDENTITIES) def test_a_sync_responses_call_that_fails_before_reaching_its_provider_frees_its_pick( call_identity: dict[str, str], @@ -1025,6 +1045,7 @@ def test_a_sync_responses_call_that_fails_before_reaching_its_provider_frees_its assert _router_pick(router) == ROUTER_DEPLOYMENT_IDS[0] + @pytest.mark.asyncio async def test_a_batch_cancel_does_not_keep_its_deployment_busy( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch @@ -1047,3 +1068,161 @@ async def test_picking_a_deployment_without_calling_it_does_not_keep_it_busy() - picked: Final = await router.select(ROUTER_GROUP, {"metadata": {}}) assert _router_pick(router) == picked["model_info"]["id"] + + +def _started_in_its_own_attempt(worker: LeastBusyLoggingHandler, kwargs: dict[str, object]) -> contextvars.Context: + def start() -> contextvars.Context: + with routed_attempt(): + worker.log_pre_api_call(model="m", messages=[], kwargs=kwargs) + return contextvars.copy_context() + + return contextvars.copy_context().run(start) + + +def test_a_late_duplicate_failure_of_one_attempt_leaves_the_next_attempt_counted() -> None: + worker: Final = _worker(None) + attempt: Final = _call_kwargs("dep-a", "retried-call") + first: Final = _started_in_its_own_attempt(worker, attempt) + first.run(worker.log_failure_event, attempt, None, None, None) + _ = _started_in_its_own_attempt(worker, attempt) + + first.run(worker.log_failure_event, attempt, None, None, None) + + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B + + +class _HeldSharedDecrement(SharedRedisCounters): + def __init__(self) -> None: + super().__init__() + self.decrementing: Final = asyncio.Event() + self.release: Final = asyncio.Event() + + async def async_increment_with_floor(self, key: str, value: int, ttl: int) -> int: + if value < 0: + self.decrementing.set() + await self.release.wait() + return self.increment_with_floor(key, value, ttl) + + +@pytest.mark.asyncio +async def test_a_failure_callback_cancelled_mid_decrement_still_frees_the_shared_count() -> None: + shared: Final = _HeldSharedDecrement() + worker: Final = _worker(shared) + attempt: Final = _call_kwargs("dep-a") + worker.log_pre_api_call(model="m", messages=[], kwargs=attempt) + settling: Final = asyncio.create_task(worker.async_log_failure_event(attempt, None, None, None)) + await shared.decrementing.wait() + + settling.cancel() + settling.cancel() + with pytest.raises(asyncio.CancelledError): + await settling + shared.release.set() + await asyncio.sleep(0) + + assert shared.count(f"{GROUP}_request_count:dep-a") == 0 + + +_CHAT_REPLY: Final = { + "id": "chatcmpl-least-busy", + "object": "chat.completion", + "created": 1, + "model": "gpt-4.1-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "done"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} + + +@pytest.mark.parametrize( + "attempts", + [("hang",), ("503", "hang"), ("503", "503")], + ids=["cancelled_while_waiting_on_the_provider", "cancelled_on_its_retry", "every_attempt_failed"], +) +@pytest.mark.asyncio +async def test_a_request_sharing_one_logging_object_that_ends_without_an_answer_leaves_its_deployment_idle( + attempts: tuple[str, ...], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + waiting_on_provider: Final = asyncio.Event() + script: Final = iter((*attempts, "ok")) + + async def upstream(_: httpx.Request) -> httpx.Response: + match next(script): + case "503": + return httpx.Response(503, json={"error": {"message": "overloaded"}}) + case "hang": + waiting_on_provider.set() + await asyncio.Event().wait() + return httpx.Response(200, json=_CHAT_REPLY) + + respx_mock.post("https://api.openai.com/v1/chat/completions").mock(side_effect=upstream) + router: Final = Router( + model_list=[ + { + "model_name": ROUTER_GROUP, + "litellm_params": {"model": "openai/gpt-4.1-mini", "api_key": "key-1"}, + "model_info": {"id": "1"}, + } + ], + routing_strategy="least-busy", + num_retries=1, + retry_after=0, + ) + logging_obj, request = litellm.utils.function_setup( + original_function="acompletion", + rules_obj=litellm.utils.Rules(), + start_time=datetime.datetime(2026, 1, 1), + model=ROUTER_GROUP, + messages=[{"role": "user", "content": "hi"}], + litellm_call_id="unanswered-call", + ) + call: Final["asyncio.Task[object]"] = asyncio.create_task( + router.acompletion(**request, litellm_logging_obj=logging_obj) + ) + if "hang" in attempts: + await waiting_on_provider.wait() + call.cancel() + + with pytest.raises((asyncio.CancelledError, litellm.ServiceUnavailableError)): + await call + await _drain_logging() + + assert router.cache.get_cache(key=f"{ROUTER_GROUP}_request_count:1") == 0 + assert isinstance( + await router.acompletion(model=ROUTER_GROUP, messages=[{"role": "user", "content": "hi"}]), ModelResponse + ) + + +@pytest.mark.asyncio +async def test_an_attempt_that_retried_inside_itself_frees_every_start_once_whatever_order_they_end_in() -> None: + shared: Final = SharedRedisCounters() + store: Final = InMemoryCache() + cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=shared) # pyright: ignore[reportArgumentType] # duck-typed Redis double + worker: Final = LeastBusyLoggingHandler(router_cache=cache, reservation_store=store) + attempt: Final = _call_kwargs("dep-a") + + async def retried_then_cancelled() -> None: + with routed_attempt() as routed: + worker.log_pre_api_call(model="m", messages=[], kwargs=attempt) + worker.log_pre_api_call(model="m", messages=[], kwargs=attempt) + await worker.async_log_failure_event(attempt, None, None, None) + release: Final = release_started_attempt(store, cache, routed.attempt_id) + if release is not None: + await release + + await retried_then_cancelled() + + assert shared.count(f"{GROUP}_request_count:dep-a") == 0 + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + + +def test_every_call_in_flight_at_once_frees_its_count_with_the_default_store_size() -> None: + worker: Final = _worker(None) + in_flight: Final = 7_000 + calls: Final = tuple(_call_kwargs("dep-a", f"call-{n}") for n in range(in_flight)) + attempts: Final = tuple(_started_in_its_own_attempt(worker, call) for call in calls) + + for attempt, call in zip(attempts, calls, strict=True): + attempt.run(worker.log_failure_event, call, None, None, None) + + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A diff --git a/tests/unit/rust_bridge/host/test_model_capabilities.py b/tests/unit/rust_bridge/host/test_model_capabilities.py index 4d73a29d078..501730431a7 100644 --- a/tests/unit/rust_bridge/host/test_model_capabilities.py +++ b/tests/unit/rust_bridge/host/test_model_capabilities.py @@ -73,3 +73,85 @@ def test_capability_source_observes_runtime_registration_changes(monkeypatch: py assert before["supports_output_config"] is True assert after["supports_reasoning"] is False assert after["supports_output_config"] is False + + +@pytest.mark.parametrize( + "model, expected", + [ + ("claude-sonnet-5", True), + ("claude-haiku-4-5", False), + ("claude-sonnet-5-5", True), + ], +) +def test_default_thinking_capability_is_projected(model: str, expected: bool) -> None: + capabilities: Final = anthropic_model_capabilities(model, "anthropic") + assert capabilities["thinking_on_by_default"] is expected + + +def test_capabilities_are_probed_once_per_model_map_state(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + model: Final = "claude-test-probe-count" + _flag_model(monkeypatch, model, supports_reasoning=True) + probes: list[str] = [] + probe: Final = AnthropicModelInfo._supports_model_capability + + def counted(name: str, key: str, provider: str) -> bool: + probes.append(key) + return probe(name, key, provider) + + monkeypatch.setattr(AnthropicModelInfo, "_supports_model_capability", staticmethod(counted)) + + first: Final = anthropic_model_capabilities(model, "anthropic") + probed: Final = len(probes) + second: Final = anthropic_model_capabilities(model, "anthropic") + + assert probed > 0 + assert len(probes) == probed + assert second == first + + +def test_a_cached_answer_is_a_fresh_copy(monkeypatch: pytest.MonkeyPatch) -> None: + model: Final = "claude-test-fresh-copy" + _flag_model(monkeypatch, model, supports_xhigh_reasoning_effort=True) + first: Final = anthropic_model_capabilities(model, "anthropic") + tiers: Final = first["effort_tiers"] + assert isinstance(tiers, dict) + first["supports_reasoning"] = "changed" + tiers["xhigh"] = "changed" + + second: Final = anthropic_model_capabilities(model, "anthropic") + + assert second["supports_reasoning"] is False + assert second["effort_tiers"] == { + **dict.fromkeys(("minimal", "low", "medium", "high", "max"), False), + "xhigh": True, + } + + +def test_register_model_is_observed_after_a_cached_answer() -> None: + model: Final = "claude-test-registered-later" + try: + litellm.register_model({model: {"litellm_provider": "anthropic", "mode": "chat", "supports_reasoning": False}}) + before: Final = anthropic_model_capabilities(model, "anthropic") + litellm.register_model({model: {"litellm_provider": "anthropic", "mode": "chat", "supports_reasoning": True}}) + after: Final = anthropic_model_capabilities(model, "anthropic") + finally: + litellm.model_cost.pop(model, None) + + assert before["supports_reasoning"] is False + assert after["supports_reasoning"] is True + + +def test_the_cache_key_holds_the_entries_it_names(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.rust_bridge.host.model_capabilities import _fingerprint # pyright: ignore[reportPrivateUsage] + + model: Final = "claude-test-held-entry" + _flag_model(monkeypatch, model, supports_reasoning=True) + entry: Final = litellm.model_cost[model] # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the model map is an untyped dict of rows + held: Final = _fingerprint(model, "anthropic", model, "anthropic") + _flag_model(monkeypatch, model, supports_reasoning=True) + same_content: Final = _fingerprint(model, "anthropic", model, "anthropic") + + assert any(h.value is entry for h in held[5]) + assert held != same_content diff --git a/tests/unit/rust_bridge/test_detectors.py b/tests/unit/rust_bridge/test_detectors.py index 3001d23e57e..a6e6b2c0dca 100644 --- a/tests/unit/rust_bridge/test_detectors.py +++ b/tests/unit/rust_bridge/test_detectors.py @@ -51,8 +51,10 @@ def call(**kwargs: object) -> NativeCall: ("openai/gpt-luna-latest", None, True), ("gpt-luna-latest", "openai", True), ("gemini/gemini-pro-latest", None, False), - ("gpt-luna-latest", None, False), + ("gpt-luna-latest", None, True), ("anthropic/claude-opus-latest", None, True), + ("moonshot/kimi-latest", None, True), + ("kimi-latest", None, False), ("claude-opus-latest", "anthropic", True), ("", None, False), ), @@ -350,6 +352,14 @@ def test_stream_cost_injection_needs_a_stream(monkeypatch: pytest.MonkeyPatch) - assert not messages_detectors.stream_cost_injection(call()) +def test_stream_cost_injection_reads_the_request_header_the_proxy_reads(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False) + opted_in: Final = {"proxy_server_request": {"headers": {"x-litellm-include-cost-in-streaming-usage": "true"}}} + assert messages_detectors.stream_cost_injection(call(stream=True, **opted_in)) + assert not messages_detectors.stream_cost_injection(call(**opted_in)) + assert not messages_detectors.stream_cost_injection(call(stream=True, proxy_server_request={"headers": {}})) + + def test_foundry_beta_row_reads_the_api_base_the_anthropic_handler_reads(monkeypatch: pytest.MonkeyPatch) -> None: for name in ("ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"): monkeypatch.delenv(name, raising=False) diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 548f78d148b..0509b8dcf30 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -610,6 +610,69 @@ def test_completion_allowed_tools_tool_choice_is_a_400_where_the_provider_cannot assert not respx_mock.calls +_HOSTED_TOOL_CHOICE: Final = {"type": "web_search"} +_HOSTED_TOOLS: Final[list[dict[str, object]]] = [*_TWO_TOOLS, {"type": "web_search"}] + + +def test_completion_forwards_a_hosted_tool_choice_to_the_responses_api_unchanged(respx_mock: respx.MockRouter): + route: Final = respx_mock.post("https://api.openai.com/v1/responses").respond( + json={ + "id": "resp_hosted", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5.6", + "output": [ + { + "type": "message", + "id": "msg_hosted", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "18C and foggy", "annotations": []}], + } + ], + "usage": {"input_tokens": 40, "output_tokens": 5, "total_tokens": 45}, + } + ) + + response: Final = litellm.completion( + model="openai/responses/gpt-5.6", + api_key="sk-test", + api_base="https://api.openai.com/v1", + messages=[{"role": "user", "content": "What is the weather in Paris right now?"}], + tools=_HOSTED_TOOLS, + tool_choice=_HOSTED_TOOL_CHOICE, + ) + + sent: Final = json.loads(route.calls.last.request.content) + assert sent["tool_choice"] == _HOSTED_TOOL_CHOICE + assert {"type": "web_search"} in sent["tools"] + assert response.choices[0].message.content == "18C and foggy" + + +@pytest.mark.parametrize("drop_params", [False, True]) +@pytest.mark.parametrize("model", ["anthropic/claude-sonnet-4-5", "gemini/gemini-2.5-flash", "bedrock/converse/x"]) +def test_completion_hosted_tool_choice_is_a_400_where_the_provider_cannot_run_it( + respx_mock: respx.MockRouter, model: str, drop_params: bool +): + with pytest.raises( + litellm.BadRequestError, match=r"tool_choice type 'web_search'.*only supported for azure, litellm_proxy, openai" + ) as exc: + litellm.completion( + model=model, + messages=[{"role": "user", "content": "What is the weather in Paris right now?"}], + tools=_TWO_TOOLS, + tool_choice=_HOSTED_TOOL_CHOICE, + drop_params=drop_params, + api_key="fake_api_key", + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + assert exc.value.status_code == 400 + assert not respx_mock.calls + + _CLAUDE_TOOL_SEARCH_CALL: Final = { "id": "srvtoolu_01search", "type": "function", @@ -6390,7 +6453,9 @@ def test_aiohttp_openai_warns_only_when_http2_enabled( assert warned is http2_on -@pytest.mark.parametrize("tool_choice", [{"type": "bogus"}, {"name": "lookup_fruit"}, {"type": "file_search"}]) +@pytest.mark.parametrize( + "tool_choice", [{"type": "bogus"}, {"name": "lookup_fruit"}, {"type": "function", "name": "lookup_fruit"}] +) def test_completion_rejects_untranslatable_tool_choice_with_a_400(tool_choice): with pytest.raises(litellm.BadRequestError) as exc_info: litellm.completion( diff --git a/tests/unit/test_router/core/test_callback_filters_step.py b/tests/unit/test_router/core/test_callback_filters_step.py new file mode 100644 index 00000000000..ee1253f900a --- /dev/null +++ b/tests/unit/test_router/core/test_callback_filters_step.py @@ -0,0 +1,84 @@ +""" +The ``callback_filters`` host step (``RustCallbackChecks._answer_callback_filters``, +litellm/router_rust/steps/callbacks.py): a pick reads the candidates' deployments back only when a +routing callback filters them, and the kept ids come back in order, once each. +""" + +from collections.abc import Mapping +from typing import Final + +import pytest + +from litellm.integrations.custom_logger import CustomLogger +from litellm.router_rust.steps.callbacks import filtered_ids, filters_deployments +from tests.unit.test_router.core.registry_helpers import MakeCore +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment + +GROUP: Final = [deployment("g", "a"), deployment("g", "b"), deployment("g", "c")] + + +class Narrow(CustomLogger): + """Keeps the deployments other than ``dropped``.""" + + def __init__(self, dropped: str) -> None: + super().__init__() + self.dropped = dropped + + async def async_filter_deployments( + self, + model: str, + healthy_deployments: list, + messages: object, + request_kwargs: object = None, + parent_otel_span: object = None, + ) -> list: + return [d for d in healthy_deployments if d["model_info"]["id"] != self.dropped] + + +class Observer(CustomLogger): + """A callback that does not override ``async_filter_deployments``.""" + + +def _entry(deployment_id: object) -> Mapping[str, object]: + return {"model_info": {"id": deployment_id}} + + +def test_filtered_ids_keep_the_filter_order_once_each_among_the_candidates() -> None: + remaining: Final = [_entry("c"), _entry("x"), _entry("a"), _entry("c"), _entry(7), {"model_info": None}, "junk"] + + assert filtered_ids(remaining, ["a", "b", "c"]) == ["c", "a"] + + +def test_only_an_overriding_filter_counts_as_a_filter() -> None: + assert filters_deployments(Narrow(dropped="a")) + assert not filters_deployments(Observer()) + assert not filters_deployments(object()) + + +@pytest.mark.asyncio +async def test_a_pick_without_filters_reads_no_deployment_back( + make_core: MakeCore, host: ScriptedHost, monkeypatch: pytest.MonkeyPatch +) -> None: + host.lists["callbacks"].append(Observer()) + core: Final = make_core(model_list=list(GROUP)) + + def refuse(_self: object, _ids: object) -> list[dict[str, object]]: # mutable-ok: the patched signature + raise AssertionError("no filter runs, so no deployment is read back") + + monkeypatch.setattr(type(core), "_records_of", refuse) + + picked: Final = {(await core.select("g"))["model_info"]["id"] for _ in range(30)} # pyright: ignore[reportIndexIssue] + + assert picked == {"a", "b", "c"} + + +@pytest.mark.asyncio +async def test_a_pick_with_a_filter_reads_the_candidates_and_honours_it( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.lists["callbacks"].append(Narrow(dropped="b")) + core: Final = make_core(model_list=list(GROUP)) + + picked: Final = {(await core.select("g"))["model_info"]["id"] for _ in range(30)} # pyright: ignore[reportIndexIssue] + + assert picked == {"a", "c"} diff --git a/tests/unit/test_router/core/test_retries_cooldowns.py b/tests/unit/test_router/core/test_retries_cooldowns.py index 004b1762240..d5ffa8491b1 100644 --- a/tests/unit/test_router/core/test_retries_cooldowns.py +++ b/tests/unit/test_router/core/test_retries_cooldowns.py @@ -13,6 +13,7 @@ import pytest import litellm from litellm.constants import DEFAULT_COOLDOWN_TIME_SECONDS, INITIAL_RETRY_DELAY, JITTER +from litellm.litellm_core_utils.total_deadline import TotalTimeoutNotApplicable, total_deadline from litellm.router_api import ChangeSet from litellm.router_core import RouterCore from litellm.types.router import AllowedFailsPolicy, RetryPolicy @@ -430,6 +431,60 @@ async def test_lone_deployment_waits_the_retry_after_header(make_core: MakeCore, assert all(7 <= slept < 7 + JITTER for slept in host.sleeps[:2]) +@pytest.mark.parametrize( + ("budget", "expected_calls"), + [(10, 2), (100, 4)], + ids=["second_backoff_outlasts_the_budget", "every_backoff_fits_the_budget"], +) +@pytest.mark.asyncio +async def test_a_retry_runs_only_while_its_backoff_fits_the_total_timeout( + make_core: MakeCore, host: ScriptedHost, budget: float, expected_calls: int +) -> None: + core: Final = make_core(model_list=[deployment("g", "a")], num_retries=3) + host.answer("openai/a", lambda **_: rate_limited(retry_after="6")) + + async with total_deadline(budget, from_body=True, clock=lambda: sum(host.sleeps)): + failure: Final = await attempt(core) + + assert isinstance(failure, litellm.RateLimitError) + assert len(host.calls_to("acompletion")) == expected_calls + assert sum(host.sleeps) < budget + + +@pytest.mark.asyncio +async def test_a_streaming_attempt_under_a_body_total_timeout_is_refused_before_the_provider_is_called( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=[deployment("g", "a")], + num_retries=3, + ) + host.answer("openai/a", lambda **_: litellm.ModelResponse()) + + async with total_deadline(10, from_body=True): + refused: Final = await attempt(core, stream=True) + + assert isinstance(refused, TotalTimeoutNotApplicable) + assert host.calls_to("acompletion") == [] + + +@pytest.mark.asyncio +async def test_a_streaming_attempt_under_a_header_total_timeout_runs_unbounded( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=[deployment("g", "a")], + num_retries=3, + ) + host.answer("openai/a", lambda **_: rate_limited(retry_after="6")) + + async with total_deadline(10, from_body=False, clock=lambda: sum(host.sleeps)): + failure: Final = await attempt(core, stream=True) + + assert isinstance(failure, litellm.RateLimitError) + assert len(host.calls_to("acompletion")) == 4 + + @pytest.mark.asyncio async def test_lone_deployment_ignores_a_retry_after_header_over_a_minute( make_core: MakeCore, host: ScriptedHost diff --git a/tests/unit/test_router/test_io_token_rate_limits.py b/tests/unit/test_router/test_io_token_rate_limits.py index f09ad5973e2..bb4813a273b 100644 --- a/tests/unit/test_router/test_io_token_rate_limits.py +++ b/tests/unit/test_router/test_io_token_rate_limits.py @@ -3,12 +3,15 @@ Tests for separate ITPM/OTPM deployment rate limits (enforce_model_rate_limits). """ import asyncio -from collections.abc import Awaitable, Callable, Mapping +import contextvars +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping from datetime import datetime, timezone from typing import Final from unittest.mock import AsyncMock, MagicMock +import httpx import pytest +import respx from fakeredis import FakeRedis, FakeServer from fakeredis.aioredis import FakeRedis as FakeAsyncRedis @@ -18,6 +21,8 @@ from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache, RedisCircuitBreakerOpenError from litellm.constants import NO_MAX_TOKENS_TPM_RESERVATION_FRACTION +from litellm.litellm_core_utils.attempt_settlement import routed_attempt +from litellm.litellm_core_utils.internal_call_metadata import IO_TOKEN_RESERVATION_ID_KEY from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( @@ -950,11 +955,8 @@ class TestModelRateLimitingCheckIOTokens: usage=Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3), ) - with pytest.raises(RuntimeError): - await async_io_token_reconcile_success(dual_cache, kwargs, response) + await async_io_token_reconcile_success(dual_cache, kwargs, response) - # The stash is cleared even though reconciliation raised, so a duplicate - # success event can't re-process it. assert ITPM_RESERVED_KEY not in metadata assert ITPM_CACHE_KEY not in metadata @@ -1853,3 +1855,332 @@ async def test_a_later_turn_is_admitted_on_this_instance_s_count_when_the_shared admitted: Final = await async_admit_io_token_turn(dual_cache, _turn_deployment(kind), _TURN_AT) assert admitted[ITPM_CACHE_KEY if kind == "itpm" else OTPM_CACHE_KEY] == _turn_key(kind) + + +_HELD_DEPLOYMENT: Final = { + "litellm_params": {"model": ROUTER_IO_MODEL, "itpm": 50000, "otpm": 50000}, + "model_info": {"id": "held-id"}, + "model_name": "limited", +} + + +def _held_keys() -> tuple[str, str]: + minute: Final = _current_minute() + return ( + f"global_router:held-id:{ROUTER_IO_MODEL}:itpm:{minute}", + f"global_router:held-id:{ROUTER_IO_MODEL}:otpm:{minute}", + ) + + +class _HeldIncrement(DualCache): + """Holds the first increment matching ``holds`` until ``release`` is set.""" + + def __init__(self, holds: Callable[[str, int], bool]) -> None: + super().__init__() + self._holds: Final = holds + self.holding: Final = asyncio.Event() + self.release: Final = asyncio.Event() + + async def async_increment_cache(self, key: str, value: float, **kwargs: object) -> float | None: + if not self.holding.is_set() and self._holds(key, int(value)): + self.holding.set() + await self.release.wait() + return await super().async_increment_cache(key=key, value=value, **kwargs) + + +async def _counts(cache: DualCache) -> tuple[int, int]: + for _ in range(10): + await asyncio.sleep(0) + itpm_key, otpm_key = _held_keys() + itpm: Final[object] = await cache.async_get_cache(key=itpm_key) + otpm: Final[object] = await cache.async_get_cache(key=otpm_key) + return (itpm if isinstance(itpm, int) else 0), (otpm if isinstance(otpm, int) else 0) + + +async def _cancel_twice_while_held(cache: _HeldIncrement, held: "asyncio.Task[object]") -> None: + await cache.holding.wait() + held.cancel() + held.cancel() + with pytest.raises(asyncio.CancelledError): + await held + cache.release.set() + + +@pytest.mark.asyncio +async def test_a_request_cancelled_mid_reservation_hands_back_what_it_reserved() -> None: + cache: Final = _HeldIncrement(lambda key, value: ":otpm:" in key and value > 0) + request: Final[dict[str, object]] = {"messages": PROXY_MESSAGES, "max_tokens": 100, "metadata": {}} + + async def reserve() -> object: + set_io_token_rate_limit_request_kwargs(request, metadata_channel="metadata") + return await async_io_token_pre_call_check(cache, dict(_HELD_DEPLOYMENT)) + + await _cancel_twice_while_held(cache, asyncio.create_task(reserve())) + + assert await _counts(cache) == (0, 0) + assert ITPM_RESERVED_KEY not in request["metadata"] + + +def _stashed(itpm_reserved: int, otpm_reserved: int) -> dict[str, object]: + itpm_key, otpm_key = _held_keys() + return { + "metadata": { + ITPM_RESERVED_KEY: itpm_reserved, + OTPM_RESERVED_KEY: otpm_reserved, + ITPM_CACHE_KEY: itpm_key, + OTPM_CACHE_KEY: otpm_key, + } + } + + +async def _reserved(cache: DualCache, itpm_reserved: int, otpm_reserved: int) -> dict[str, object]: + itpm_key, otpm_key = _held_keys() + await cache.async_increment_cache(key=itpm_key, value=itpm_reserved) + await cache.async_increment_cache(key=otpm_key, value=otpm_reserved) + return _stashed(itpm_reserved, otpm_reserved) + + +@pytest.mark.asyncio +async def test_a_refund_cancelled_between_its_decrements_still_hands_back_both() -> None: + cache: Final = _HeldIncrement(lambda key, value: ":itpm:" in key and value < 0) + kwargs: Final = await _reserved(cache, 9, 100) + + await _cancel_twice_while_held( + cache, asyncio.create_task(ModelRateLimitingCheck(cache).async_log_failure_event(kwargs, None, None, None)) + ) + + assert await _counts(cache) == (0, 0) + + +class _YieldingIncrement(DualCache): + async def async_increment_cache(self, key: str, value: float, **kwargs: object) -> float | None: + await asyncio.sleep(0) + return await super().async_increment_cache(key=key, value=value, **kwargs) + + +@pytest.mark.asyncio +async def test_competing_failure_callbacks_hand_a_reservation_back_once() -> None: + cache: Final = _YieldingIncrement() + kwargs: Final = await _reserved(cache, 9, 100) + check: Final = ModelRateLimitingCheck(cache) + + await asyncio.gather(*(check.async_log_failure_event(kwargs, None, None, None) for _ in range(2))) + + assert await _counts(cache) == (0, 0) + + +@pytest.mark.asyncio +async def test_a_late_failure_callback_of_one_attempt_leaves_the_next_attempts_reservation_held() -> None: + cache: Final = DualCache() + check: Final = ModelRateLimitingCheck(cache) + request: Final[dict[str, object]] = {"messages": PROXY_MESSAGES, "max_tokens": 100, "metadata": {}} + + async def attempt() -> contextvars.Context: + with routed_attempt(): + set_io_token_rate_limit_request_kwargs(request, metadata_channel="metadata") + await async_io_token_pre_call_check(cache, dict(_HELD_DEPLOYMENT)) + return contextvars.copy_context() + + def failure_logged_in(context: contextvars.Context) -> "asyncio.Task[None]": + return context.run(asyncio.create_task, check.async_log_failure_event(request, None, None, None)) + + first: Final = await attempt() + await failure_logged_in(first) + _ = await attempt() + held_by_second: Final = await _counts(cache) + + await failure_logged_in(first) + + assert held_by_second[1] == 100 + assert await _counts(cache) == held_by_second + + +def _anthropic_reply() -> dict[str, object]: + return { + "id": "msg_held", + "type": "message", + "role": "assistant", + "model": "claude-haiku-5-5", + "content": [{"type": "text", "text": "done"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 3, "output_tokens": 1}, + } + + +@pytest.mark.parametrize( + "attempts", + [("hang",), ("503", "hang")], + ids=["cancelled_while_waiting_on_the_provider", "cancelled_on_its_retry"], +) +@pytest.mark.asyncio +async def test_a_cancelled_request_sharing_one_logging_object_hands_back_every_reservation( + attempts: tuple[str, ...], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + waiting_on_provider: Final = asyncio.Event() + script: Final = iter((*attempts, "ok")) + + async def upstream(_: httpx.Request) -> httpx.Response: + match next(script): + case "503": + return httpx.Response( + 503, json={"type": "error", "error": {"type": "overloaded_error", "message": "x"}} + ) + case "hang": + waiting_on_provider.set() + await asyncio.Event().wait() + return httpx.Response(200, json=_anthropic_reply()) + + respx_mock.post("https://api.anthropic.com/v1/messages").mock(side_effect=upstream) + router: Final = _router_with_io_limits() + first_minute: Final = _current_minute() + call: Final["asyncio.Task[object]"] = asyncio.create_task( + router.acompletion( + model="limited", + messages=PROXY_MESSAGES, + max_tokens=100, + metadata={"user_api_key_hash": "hashed-key"}, + litellm_call_id="proxy-call", + litellm_logging_obj=_logging_obj_shared_across_attempts("acompletion"), + ) + ) + await waiting_on_provider.wait() + call.cancel() + + with pytest.raises(asyncio.CancelledError): + await call + + assert await _io_counters_since(router, first_minute) == (0, 0) + assert await _proxy_chat(router, "hi") == await _io_counters_since(router, first_minute) + + +@pytest.mark.asyncio +async def test_a_reservation_refunded_through_a_logged_copy_of_its_metadata_is_refunded_once() -> None: + cache: Final = DualCache() + check: Final = ModelRateLimitingCheck(cache) + request: Final[dict[str, object]] = {"messages": PROXY_MESSAGES, "max_tokens": 100, "litellm_metadata": {}} + + async def cancelled_attempt() -> None: + with routed_attempt(): + set_io_token_rate_limit_request_kwargs(request, metadata_channel="litellm_metadata") + await async_io_token_pre_call_check(cache, dict(_HELD_DEPLOYMENT)) + logged.update({"litellm_params": {"litellm_metadata": dict(request["litellm_metadata"])}}) # pyright: ignore[reportArgumentType] # the stashed metadata dict + await check.async_log_failure_event(request, None, None, None) + + logged: Final[dict[str, object]] = {} + await asyncio.create_task(cancelled_attempt()) + + await check.async_log_failure_event(logged, None, None, None) + + assert await _counts(cache) == (0, 0) + + +@pytest.mark.asyncio +async def test_a_retry_refunds_the_reservation_its_previous_attempt_left_behind() -> None: + cache: Final = DualCache() + request: Final[dict[str, object]] = {"messages": PROXY_MESSAGES, "max_tokens": 100, "metadata": {}} + + async def first_attempt() -> None: + with routed_attempt(): + set_io_token_rate_limit_request_kwargs(request, metadata_channel="metadata") + await async_io_token_pre_call_check(cache, dict(_HELD_DEPLOYMENT)) + + def retry() -> None: + with routed_attempt(): + refund_stale_reservation_before_retry(cache, request) + + await asyncio.create_task(first_attempt()) + contextvars.copy_context().run(retry) + + assert await _counts(cache) == (0, 0) + + +@pytest.mark.asyncio +async def test_a_request_cancelled_after_its_provider_answered_keeps_the_tokens_it_spent_counted( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.integrations.custom_logger import CustomLogger + + answered: Final = asyncio.Event() + + class SlowSuccessHook(CustomLogger): + async def async_post_call_success_deployment_hook(self, request_data, response, call_type): + answered.set() + await asyncio.Event().wait() + + monkeypatch.setattr(litellm, "callbacks", [*litellm.callbacks, SlowSuccessHook()]) + router: Final = _router_with_io_limits() + first_minute: Final = _current_minute() + call: Final["asyncio.Task[tuple[int, int]]"] = asyncio.create_task(_proxy_chat(router, "hi")) + await answered.wait() + call.cancel() + + with pytest.raises(asyncio.CancelledError): + await call + + _, otpm = await _io_counters_since(router, first_minute) + assert otpm == 100 + + +def _stashed_as(reservation_id: str, itpm_reserved: int, otpm_reserved: int) -> dict[str, object]: + stashed: Final = _stashed(itpm_reserved, otpm_reserved) + return {"metadata": {**stashed["metadata"], IO_TOKEN_RESERVATION_ID_KEY: reservation_id}} # pyright: ignore[reportGeneralTypeIssues] # the stash is a dict + + +@pytest.mark.asyncio +async def test_a_settled_reservation_stays_settled_however_many_others_settle_after_it() -> None: + cache: Final = DualCache() + check: Final = ModelRateLimitingCheck(cache) + itpm_key, otpm_key = _held_keys() + await cache.async_increment_cache(key=itpm_key, value=9) + await cache.async_increment_cache(key=otpm_key, value=100) + first: Final = _stashed_as("first-reservation", 9, 100) + logged_copy: Final = {"litellm_params": {"metadata": dict(first["metadata"])}} # pyright: ignore[reportArgumentType] # the stash is a dict + await check.async_log_failure_event(first, None, None, None) + for n in range(10_001): + io_token_refund_failure(cache, _stashed_as(f"later-reservation-{n}", 0, 0)) + + await check.async_log_failure_event(logged_copy, None, None, None) + + assert await _counts(cache) == (0, 0) + + +class _FirstChunkNeverArrives(httpx.AsyncByteStream): + def __init__(self, reading: asyncio.Event) -> None: + self._reading: Final = reading + + async def __aiter__(self) -> AsyncIterator[bytes]: + self._reading.set() + await asyncio.Event().wait() + yield b"" + + +@pytest.mark.asyncio +async def test_a_stream_cancelled_before_its_first_chunk_hands_back_its_reservation( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + reading_first_chunk: Final = asyncio.Event() + respx_mock.post("https://api.anthropic.com/v1/messages").mock( + return_value=httpx.Response( + 200, headers={"content-type": "text/event-stream"}, stream=_FirstChunkNeverArrives(reading_first_chunk) + ) + ) + router: Final = _router_with_io_limits() + first_minute: Final = _current_minute() + call: Final["asyncio.Task[object]"] = asyncio.create_task( + router.acompletion( + model="limited", + messages=PROXY_MESSAGES, + max_tokens=100, + stream=True, + metadata={"user_api_key_hash": "hashed-key"}, + ) + ) + await reading_first_chunk.wait() + call.cancel() + + with pytest.raises(asyncio.CancelledError): + await call + + assert await _io_counters_since(router, first_minute) == (0, 0) diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index ca61fcd0905..24ade9b48fb 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -20331,3 +20331,22 @@ def test_router_warns_when_a_deployment_enforces_endpoints_without_a_list(caplog assert "Deployment 'embed-from-config' sets enforce_supported_endpoints without" in caplog.text assert "Deployment 'embed-from-api' sets enforce_supported_endpoints without" in caplog.text + + +def test_registry_finds_a_team_deployment_by_public_and_internal_name_for_its_team_only() -> None: + internal_name: Final = "model_name_team-a_row" + router: Final = Router( + model_list=[ + { + "model_name": internal_name, + "litellm_params": {"model": "openai/gpt-5.5", "api_key": "x"}, + "model_info": {"id": "team-row", "team_id": "team-a", "team_public_model_name": "public-chat"}, + } + ] + ) + registry: Final = router.registry() + + assert [row["model_info"]["id"] for row in registry.deployments("public-chat", team_id="team-a")] == ["team-row"] + assert [row["model_info"]["id"] for row in registry.deployments(internal_name, team_id="team-a")] == ["team-row"] + assert list(registry.deployments("public-chat", team_id="team-b")) == [] + assert list(registry.deployments("public-chat")) == [] diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 44bdd187da4..a8421d8afbd 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -6511,6 +6511,34 @@ async def test_wrapper_async_raises_original_exception_even_if_hook_callback_err ) +@pytest.mark.asyncio +async def test_wrapper_async_lets_a_cancellation_during_the_failure_hook_cancel_the_call( + monkeypatch: pytest.MonkeyPatch, +) -> None: + hook_entered: Final = asyncio.Event() + + class HangingLogger(CustomLogger): + async def async_post_call_failure_deployment_hook( + self, request_data, exception, call_type, fallback_depth=None + ): + hook_entered.set() + await asyncio.Event().wait() + + monkeypatch.setattr(litellm, "callbacks", [HangingLogger()]) + call: Final = asyncio.create_task( + litellm.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + mock_response=litellm.AuthenticationError(message="bad key", llm_provider="openai", model="gpt-4o-mini"), + ) + ) + await hook_entered.wait() + call.cancel() + + with pytest.raises(asyncio.CancelledError): + await call + + @pytest.mark.asyncio async def test_router_fallback_chain_reports_increasing_fallback_depth(monkeypatch: pytest.MonkeyPatch) -> None: """Regression: a real Router fallback chain must report fallback_depth=None on the @@ -6898,41 +6926,6 @@ async def test_router_fallback_not_skipped_when_failure_hook_callback_touches_at assert response.choices[0].message.content == "fallback worked" -@pytest.mark.asyncio -async def test_wrapper_async_preserves_original_exception_when_hook_await_is_cancelled( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """Regression: if the caller's own timeout (e.g. asyncio.wait_for) fires while the - failure hook is still being awaited, the real deployment exception must still reach - the caller - not get replaced by CancelledError/TimeoutError from the hook's own - await getting cancelled.""" - - hook_entered: Final = asyncio.Event() - hook_released: Final = asyncio.Event() - - class SlowLogger(CustomLogger): - async def async_post_call_failure_deployment_hook( - self, request_data, exception, call_type, fallback_depth=None - ): - hook_entered.set() - await hook_released.wait() - - monkeypatch.setattr(litellm, "callbacks", [SlowLogger()]) - - call: Final = asyncio.create_task( - litellm.acompletion( - model="gpt-4o-mini", - messages=[{"role": "user", "content": "hi"}], - mock_response=litellm.AuthenticationError(message="bad key", llm_provider="openai", model="gpt-4o-mini"), - ) - ) - assert await asyncio.wait_for(hook_entered.wait(), timeout=5) - call.cancel() - - with pytest.raises(litellm.AuthenticationError): - await call - - @pytest.mark.asyncio async def test_wrapper_async_failure_hook_latency_does_not_inflate_reported_duration( monkeypatch: pytest.MonkeyPatch, @@ -8809,7 +8802,8 @@ def test_validate_tool_choice_cursor_format(): {"type": "invalid"}, {"type": "function"}, {"name": "lookup_fruit"}, - {"type": "file_search"}, + {"type": "web_search_tool"}, + {"type": ["web_search"]}, ], ) def test_validate_tool_choice_invalid_dict_is_a_400(tool_choice): @@ -8823,6 +8817,24 @@ def test_validate_tool_choice_invalid_dict_is_a_400(tool_choice): assert exc_info.value.model == MODEL +@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") +@pytest.mark.parametrize( + "tool_choice", + [ + {"type": "web_search"}, + {"type": "web_search_preview"}, + {"type": "file_search"}, + {"type": "tool_search"}, + {"type": "code_interpreter"}, + {"type": "mcp", "server_label": "docs", "name": "search"}, + ], + ids=lambda choice: str(choice["type"]), +) +def test_validate_tool_choice_forcing_a_hosted_tool_passes_through(tool_choice: dict[str, object]): + """Forcing an OpenAI hosted tool is a Responses API choice the bridge sends as is, so validation leaves it alone.""" + assert validate_chat_completion_tool_choice(tool_choice, model=MODEL) is tool_choice + + @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown") @pytest.mark.parametrize("mode", ["auto", "required"]) def test_validate_tool_choice_allowed_tools_chat_layout_passes_through(mode: str): diff --git a/tests/unit/test_video_generation.py b/tests/unit/test_video_generation.py index 9d9ca3bb734..93243607b6f 100644 --- a/tests/unit/test_video_generation.py +++ b/tests/unit/test_video_generation.py @@ -17,6 +17,7 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.llms.gemini.videos.transformation import GeminiVideoConfig from litellm.llms.openai.videos.transformation import OpenAIVideoConfig +from litellm.proxy._types import UserAPIKeyAuth from litellm.types.videos.main import VideoObject, VideoResponse from litellm.videos import main as videos_main from litellm.videos.main import ( @@ -2232,7 +2233,7 @@ def video_proxy_test_client(): app = FastAPI() app.include_router(video_router) - app.dependency_overrides[user_api_key_auth] = lambda: MagicMock() + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(api_key="sk-9876") return TestClient(app) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentIdentityFields.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentIdentityFields.tsx index d4af3ff3833..00c8fce98d1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentIdentityFields.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentIdentityFields.tsx @@ -8,6 +8,7 @@ import { Input } from "@/components/ui/input"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { AgentFormField, type AgentFormValues } from "./AgentFormKit"; import { entraTenantFromIssuer, IDENTITY_UUID_PATTERN } from "./agent_identity"; +import { codeSamples } from "@/i18n/untranslated/codeSamples"; const PROVIDER_OPTIONS: { value: string; labelKey: ParseKeys }[] = [ { value: "none", labelKey: "agents.identityFields.providers.none" }, @@ -206,7 +207,7 @@ export const AgentIdentityFields = ({ accessToken }: { accessToken: string | nul ref={ref} value={typeof value === "string" ? value : ""} onChange={onChange} - placeholder="Agent.Invoke" + placeholder={codeSamples.agents.identityScope} /> )} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/KillSwitchFormFields.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/KillSwitchFormFields.tsx index 2b9a0edb70d..6abeb5d18ea 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/KillSwitchFormFields.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/KillSwitchFormFields.tsx @@ -11,6 +11,7 @@ import { PasswordInput } from "@/components/shared/PasswordInput"; import i18n from "@/i18n/i18n"; import { AgentFormField, AgentFormValues, labelWithHint } from "./AgentFormKit"; import { KILL_SWITCH_AUTH_TYPES, KILL_SWITCH_METHODS, validateKillSwitchBody } from "./kill_switch_config"; +import { codeSamples } from "@/i18n/untranslated/codeSamples"; const KeyValueFieldArray = ({ name, @@ -135,7 +136,7 @@ const KillSwitchAuthFields = () => { ( {...control} ref={ref} rows={4} - placeholder='{"reason": "manual kill switch"}' + placeholder={codeSamples.agents.killSwitchPayloadJson} value={typeof value === "string" ? value : ""} onChange={onChange} /> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/api-reference/_components/APIReferenceView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/api-reference/_components/APIReferenceView.tsx index 7f6f7ea1d17..f55dc2e4c82 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/api-reference/_components/APIReferenceView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/api-reference/_components/APIReferenceView.tsx @@ -4,6 +4,7 @@ import { useTranslation } from "react-i18next"; import CodeBlock from "@/components/CodeBlock"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import DocLink from "./DocLink"; +import { productNames } from "@/i18n/untranslated/productNames"; interface ApiRefProps { proxySettings: { @@ -36,13 +37,13 @@ const APIReferenceView: React.FC = ({ proxySettings }) => { - OpenAI Python SDK + {productNames.openAIPythonSdk} - LlamaIndex + {productNames.llamaIndex} - Langchain Py + {productNames.langchainPython} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_health.tsx b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_health.tsx index 382438d9efe..db532fafefe 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_health.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_health.tsx @@ -5,6 +5,8 @@ import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { CheckCircle2, XCircle, ClipboardCopy } from "lucide-react"; import { ResponseTimeIndicator } from "./response_time_indicator"; +const UNHEALTHY_STATUS = "unhealthy"; + // Helper function to deep-parse a JSON string if possible const deepParse = (input: any) => { let parsed = input; @@ -179,7 +181,7 @@ const HealthCheckDetails: React.FC<{ response: any }> = ({ response }) => { className={`text-sm font-medium ${response?.status === "healthy" ? "text-success" : "text-destructive"}`} > {t("caching.health.cacheStatus")} - {response?.status || "unhealthy"} + {response?.status || UNHEALTHY_STATUS}

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/EvaluationSettingsModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/EvaluationSettingsModal.tsx index 12ff10bc0c5..0c3cb62a57f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/EvaluationSettingsModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/EvaluationSettingsModal.tsx @@ -14,6 +14,7 @@ import { DialogTitle, } from "@/components/ui/dialog"; import { Textarea } from "@/components/ui/textarea"; +import { codeSamples } from "@/i18n/untranslated/codeSamples"; const DEFAULT_PROMPT = `Evaluate whether this guardrail's decision was correct. Analyze the user input, the guardrail action taken, and determine if it was appropriate. @@ -130,7 +131,7 @@ export function EvaluationSettingsModal({ -

response_format: json_schema

+

{codeSamples.guardrails.evaluationResponseFormat}