Merge remote-tracking branch 'origin/main' into HEAD

# Conflicts:
#	litellm/proxy/management_endpoints/internal_user_endpoints.py
This commit is contained in:
Jesus Velarde 2026-11-09 12:41:57 -08:00
commit 0c3e369e3e
390 changed files with 18156 additions and 3809 deletions

View file

@ -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]:

View file

@ -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,

View file

@ -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,

View file

@ -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<Value>,
) {
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())]

View file

@ -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,

View file

@ -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",

View file

@ -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<SpendCounter>, 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<SpendCounter>, PyNumber)>,
) -> bool {
) -> Option<PyPostCallSettlement> {
self.0
.increment_many_post_call(&request.0, &self::increments(increments))
.map(PyPostCallSettlement)
}
fn set_max<'py>(

View file

@ -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<T>`, 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<T>`, 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<RemoteSlot>`. 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`

View file

@ -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<SpendCounter>),
Key {
token: String,
window: Option<String>,
@ -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<Self, NotASpendCounter> {
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())),

View file

@ -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<Number>,
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::<Vec<_>>()).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<Self>,
request: &RequestRedis,
increments: &[(SpendCounter, Number)],
) -> bool {
let Some(remote) = self.cache.redis().get() else {
return false;
};
) -> Option<PostCallSettlement> {
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::<Vec<_>>();
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<f64> = 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?)
}

View file

@ -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<AtomicBool>);
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 {

View file

@ -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<RedisClient>) -> Arc<SpendCounters> {
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::<String>(&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::<String>(&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])

View file

@ -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

View file

@ -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)

View file

@ -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())

View file

@ -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)

View file

@ -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,

View file

@ -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")),

View file

@ -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

View file

@ -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,
)

View file

@ -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

View file

@ -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 ---

View file

@ -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": [<at least one tool>]}}'
)
@ -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)

View file

@ -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,
)

View file

@ -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)

View file

@ -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

View file

@ -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(

View file

@ -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.

View file

@ -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

View file

@ -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

View file

@ -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(

View file

@ -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(

View file

@ -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:

View file

@ -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)

View file

@ -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

View file

@ -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,
(

View file

@ -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,
(

View file

@ -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")

View file

@ -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:

View file

@ -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(

View file

@ -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:

View file

@ -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:
"""

View file

@ -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,

View file

@ -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(

View file

@ -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__()

View file

@ -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)

View file

@ -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,

View file

@ -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,

View file

@ -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),
)

View file

@ -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)

View file

@ -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)

View file

@ -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/<index>/<api>, for example /opensearch/<index>/_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
)

View file

@ -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"

View file

@ -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)

View file

@ -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,

View file

@ -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)

View file

@ -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'] = <your-region_name>
class SagemakerLLM(BaseAWSLLM):
_web_identity_session_policy: ClassVar[WebIdentitySessionPolicy] = SAGEMAKER_SESSION_POLICY
def _load_credentials(
self,
optional_params: dict,

View file

@ -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(

View file

@ -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"}
}

File diff suppressed because it is too large Load diff

View file

@ -243,6 +243,7 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = (
"/mistral/",
"/nvidia_nim/",
"/openai/",
"/opensearch/",
"/openai_passthrough/",
"/tinyfish/",
"/transcribe",

View file

@ -35987,6 +35987,386 @@
]
}
},
"/opensearch/{endpoint}": {
"delete": {
"description": "Call an OpenSearch index through the proxy, e.g. `POST /opensearch/<index>/_search`.\n\n`<index>` 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/<index>/_search`.\n\n`<index>` 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/<index>/_search`.\n\n`<index>` 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/<index>/_search`.\n\n`<index>` 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/<index>/_search`.\n\n`<index>` 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/<index>/_search`.\n\n`<index>` 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)",

View file

@ -427,6 +427,7 @@ class LiteLLMRoutes(enum.Enum):
"/bespoke",
"/openrouter",
"/milvus",
"/opensearch",
"/gigachat",
"/watsonx",
"/nvidia_nim",

View file

@ -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,
)

View file

@ -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)

View file

@ -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

View file

@ -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,

View file

@ -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,

View file

@ -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`.

View file

@ -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,
)

View file

@ -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

View file

@ -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:

View file

@ -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:
"""

View file

@ -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"

View file

@ -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,

View file

@ -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]:

View file

@ -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)

View file

@ -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 _:

View file

@ -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

View file

@ -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()

View file

@ -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:

View file

@ -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):

View file

@ -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": "(?:(?<![A-Za-z0-9])|(?<=\\\\[nrt])|(?<=\\\\u[0-9A-F]{4})|(?<=%[0-9A-F]{2})|(?<=%25[0-9A-F]{2}))sk-ant-[A-Za-z0-9_\\-]{20,}",
"category": "Credential Patterns",
"description": "Detects Anthropic API keys (sk-ant-... format)"
},
{
"name": "gemini_api_key",
"display_name": "Gemini / Google API Key",
"pattern": "(?:(?<![A-Za-z0-9])|(?<=\\\\[nrt])|(?<=\\\\u[0-9A-F]{4})|(?<=%[0-9A-F]{2})|(?<=%25[0-9A-F]{2}))(?:(?-i:AIza)[A-Za-z0-9_\\-]{33,36}|(?-i:AQ\\.)[A-Za-z0-9_\\-]{45,55})(?![A-Za-z0-9_])",
"category": "Credential Patterns",
"description": "Detects Google / Gemini API keys (legacy AIza... and new AQ.... authorization key formats)"
},
{
"name": "ipv4",
"display_name": "IP Address (IPv4)",
@ -626,4 +648,4 @@
"allow_word_numbers": false
}
]
}
}

View file

@ -11,10 +11,11 @@ import re
from collections.abc import Iterator, Sequence
from enum import Enum
from re import Pattern
from types import MappingProxyType
from typing import Final
from pydantic import ConfigDict, TypeAdapter, with_config
from typing_extensions import ReadOnly, TypedDict
from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm.proxy.guardrails.content_filter_data import DATA_ROOTS, category_dirs
@ -26,6 +27,7 @@ class _PatternEntry(TypedDict):
pattern: ReadOnly[str]
category: ReadOnly[str]
description: ReadOnly[str]
streaming_holdback_chars: ReadOnly[NotRequired[int]]
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
@ -84,6 +86,7 @@ KNOWN_PATTERN_KEYS: Final = {
"category",
"action",
"description",
"streaming_holdback_chars",
}
PATTERN_EXTRA_CONFIG: Final[dict[str, _PatternExtraConfig]] = {}
@ -132,6 +135,14 @@ for pattern_data in _PATTERNS_DATA["patterns"]:
PATTERN_CATEGORIES[category] = []
PATTERN_CATEGORIES[category].append(pattern_data["name"])
PATTERN_STREAMING_HOLDBACK_CHARS: Final = MappingProxyType(
{
pattern_data["name"]: pattern_data["streaming_holdback_chars"]
for pattern_data in _PATTERNS_DATA["patterns"]
if "streaming_holdback_chars" in pattern_data
}
)
# Build display names mapping from JSON
PATTERN_DISPLAY_NAMES: Final[dict[str, str]] = {

View file

@ -65,6 +65,7 @@ from litellm.litellm_core_utils.token_counter import (
offload_token_count,
split_embedding_input_tokens,
)
from litellm.litellm_core_utils.total_deadline import total_timeout_held
from litellm.llms.azure.ptu_capacity import normalized_tokens
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.auth_utils import (
@ -3801,6 +3802,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
raise cancellation
@with_service_target(RATE_LIMITS_TARGET)
@total_timeout_held()
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,

View file

@ -47,6 +47,7 @@ from litellm.proxy.db.db_spend_update_writer import (
from litellm.proxy.hooks.search_free_tier import apply_search_free_tier
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, metadata_variable_name_for_request
from litellm.proxy.route_llm_request import ProxyModelNotFoundError
from litellm.proxy.spend_tracking.committed_spend import committed_spend_tracked, warm_committed_counter_keys
from litellm.proxy.spend_tracking.spend_counter_batch import post_call_counter_keys, spend_counter_batch_scope
from litellm.proxy.spend_tracking.spend_event import (
ObjectMapping,
@ -75,6 +76,8 @@ from litellm.proxy.utils import (
BilledFailedAttempt,
ProxyUpdateSpend,
)
from litellm.types.proxy.auth.user_api_key_auth import committed_budget_spend_of_row
from litellm.types.proxy.carried_budget_state import EntitySpendCounterKeys
from litellm.types.router import Deployment
from litellm.types.utils import (
CallTypes,
@ -268,6 +271,11 @@ def _proxy_spend_writer() -> 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:

View file

@ -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)}"

View file

@ -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(

View file

@ -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,

View file

@ -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:

View file

@ -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,

View file

@ -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

View file

@ -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

View file

@ -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, ...]:

View file

@ -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 = (

View file

@ -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)

View file

@ -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

View file

@ -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,

View file

@ -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,

View file

@ -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

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