mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
Merge remote-tracking branch 'origin/main' into HEAD
# Conflicts: # litellm/proxy/management_endpoints/internal_user_endpoints.py
This commit is contained in:
commit
0c3e369e3e
390 changed files with 18156 additions and 3809 deletions
|
|
@ -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]:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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())]
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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>(
|
||||
|
|
|
|||
|
|
@ -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`
|
||||
|
|
|
|||
|
|
@ -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())),
|
||||
|
|
|
|||
|
|
@ -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?)
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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])
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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")),
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 ---
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
90
litellm/litellm_core_utils/attempt_settlement.py
Normal file
90
litellm/litellm_core_utils/attempt_settlement.py
Normal 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,
|
||||
)
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
197
litellm/litellm_core_utils/total_deadline.py
Normal file
197
litellm/litellm_core_utils/total_deadline.py
Normal 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
|
||||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
(
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
(
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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__()
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
362
litellm/llms/bedrock/web_identity_session_policy.py
Normal file
362
litellm/llms/bedrock/web_identity_session_policy.py
Normal 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),
|
||||
)
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
379
litellm/llms/opensearch/vector_stores/index_passthrough.py
Normal file
379
litellm/llms/opensearch/vector_stores/index_passthrough.py
Normal 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
|
||||
)
|
||||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
|
@ -243,6 +243,7 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = (
|
|||
"/mistral/",
|
||||
"/nvidia_nim/",
|
||||
"/openai/",
|
||||
"/opensearch/",
|
||||
"/openai_passthrough/",
|
||||
"/tinyfish/",
|
||||
"/transcribe",
|
||||
|
|
|
|||
|
|
@ -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)",
|
||||
|
|
|
|||
|
|
@ -427,6 +427,7 @@ class LiteLLMRoutes(enum.Enum):
|
|||
"/bespoke",
|
||||
"/openrouter",
|
||||
"/milvus",
|
||||
"/opensearch",
|
||||
"/gigachat",
|
||||
"/watsonx",
|
||||
"/nvidia_nim",
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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`.
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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]:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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 _:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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]] = {
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)}"
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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, ...]:
|
||||
|
|
|
|||
|
|
@ -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 = (
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
Loading…
Add table
Reference in a new issue