diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index dce8c07fb0b..51542654a02 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -42,6 +42,9 @@ mod _native { NativeAuthContinuation, request_auth_budget_alert, request_auth_budget_alert_plan, request_auth_budget_exemption, request_auth_budget_price, request_auth_budget_route, request_auth_budget_window_start, request_auth_budgets, request_auth_common_budget, + request_auth_context_end_user, request_auth_context_member, request_auth_context_models, + request_auth_context_percentage, request_auth_context_temporary, + request_auth_context_throttle, request_auth_context_token, request_auth_litellm_header_token, request_auth_member_limit, request_auth_model_budget_check, request_auth_model_budget_query, request_auth_model_budget_resolve, request_auth_scheme_credential, diff --git a/litellm-rust/crates/python-bridge/src/request_auth.rs b/litellm-rust/crates/python-bridge/src/request_auth.rs index ab4562e4adb..e8a5bf6f5cf 100644 --- a/litellm-rust/crates/python-bridge/src/request_auth.rs +++ b/litellm-rust/crates/python-bridge/src/request_auth.rs @@ -234,3 +234,84 @@ pub(crate) fn request_auth_budget_alert_plan( let facts: core::budget_alerts::Facts = from_py(facts)?; to_py(py, &py.detach(move || core::budget_alerts::plan(&facts))) } + +#[pyfunction] +pub(crate) fn request_auth_context_models( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_context::ModelCheckFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::budget_context::model_checks(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_context_end_user( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_context::EndUserFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::budget_context::end_user_source(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_context_token( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_context::TokenDefaultsFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::budget_context::token_assignments(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_context_member( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_context::MemberFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::budget_context::member_source(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_context_temporary( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_context::TemporaryFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::budget_context::temporary(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_context_throttle( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_context::ThrottleFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::budget_context::throttle(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_context_percentage( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_context::NumericFacts = from_py(facts)?; + Ok(py.detach(move || core::budget_context::percentage(&facts))) +} diff --git a/litellm-rust/crates/request-auth-core/src/budget_context.rs b/litellm-rust/crates/request-auth-core/src/budget_context.rs new file mode 100644 index 00000000000..9f9f9e1b107 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/budget_context.rs @@ -0,0 +1,481 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq, strum::VariantArray)] +pub enum ModelScope { + Team, + EndUser, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct ModelCheckFacts { + pub scope: ModelScope, + pub entity_id: Option, + pub budget_count: usize, + pub model_count: usize, +} + +pub fn model_checks(facts: &ModelCheckFacts) -> Vec { + if facts.budget_count == 0 + || facts.entity_id.is_none() + || (facts.scope == ModelScope::EndUser && facts.entity_id.as_deref() == Some("")) + { + return vec![]; + } + (0..facts.model_count).collect() +} + +#[derive(Clone, Debug, Default, Deserialize)] +pub struct BudgetFields { + pub max_budget: Option, + pub tpm_limit: Option, + pub rpm_limit: Option, + pub tpd_limit: Option, + pub model_budget_count: Option, +} + +#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq, strum::VariantArray)] +pub enum EndUserMode { + Models, + Token, + Default, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct EndUserFacts { + pub mode: EndUserMode, + pub object_present: bool, + pub end_user_id: Option, + pub prisma_present: bool, + pub key_budget_id: Option, + pub key_read: bool, + pub key_budget: Option, + pub global_config_read: bool, + pub global_budget_id: Option, + pub global_read: bool, + pub global_budget: Option, +} + +#[derive(Clone, Copy, Debug, Serialize, PartialEq, Eq, strum::VariantArray)] +pub enum EndUserSource { + Own, + ReadKey, + Key, + ReadGlobalConfig, + ReadGlobal, + Global, + None, +} + +pub fn end_user_source(facts: &EndUserFacts) -> EndUserSource { + use EndUserSource::*; + if facts.mode != EndUserMode::Default && facts.object_present { + return Own; + } + if facts.mode == EndUserMode::Models + && (!facts.prisma_present || facts.end_user_id.as_deref().is_none_or(str::is_empty)) + { + return None; + } + if facts.key_budget_id.is_some() { + if !facts.key_read { + return ReadKey; + } + if facts.key_budget.is_some() { + return Key; + } + } + if !facts.global_config_read { + return ReadGlobalConfig; + } + if facts.global_budget_id.is_none() { + return None; + } + if !facts.global_read { + return ReadGlobal; + } + if facts.global_budget.is_some() { + Global + } else { + None + } +} + +#[derive(Clone, Debug, Deserialize)] +pub struct TokenDefaultsFacts { + pub current: BudgetFields, + pub keep_token_limits: bool, + pub default_present: bool, +} + +#[derive(Clone, Copy, Debug, Serialize, PartialEq, Eq, strum::VariantArray)] +pub enum TokenField { + MaxBudget, + Tpm, + Rpm, + Tpd, + Models, +} + +pub fn token_assignments(facts: &TokenDefaultsFacts) -> Vec { + if !facts.default_present { + return vec![]; + } + [ + (TokenField::MaxBudget, facts.current.max_budget.is_none()), + (TokenField::Tpm, facts.current.tpm_limit.is_none()), + (TokenField::Rpm, facts.current.rpm_limit.is_none()), + (TokenField::Tpd, facts.current.tpd_limit.is_none()), + ( + TokenField::Models, + facts.current.model_budget_count.is_none(), + ), + ] + .into_iter() + .filter_map(|(field, missing)| (!facts.keep_token_limits || missing).then_some(field)) + .collect() +} + +#[derive(Clone, Debug, Deserialize)] +pub struct MembershipSnapshot { + pub budget_id: Option, + pub model_budget_count: Option, +} + +#[derive(Clone, Copy, Debug, Deserialize, strum::VariantArray)] +pub enum LookupFailure { + Membership, + Default, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct MemberFacts { + pub user_id: Option, + pub team_id: Option, + pub prisma_present: bool, + pub default_id: Option, + pub skip_budget_checks: bool, + pub membership_read: bool, + pub membership: Option, + pub default_read: bool, + pub default_budget: Option, + pub lookup_failure: Option, +} + +#[derive(Clone, Copy, Debug, Serialize, PartialEq, Eq, strum::VariantArray)] +pub enum MemberSource { + Done, + ReadMembership, + ReadDefault, + Member, + Default, + Clear, + RaiseLookup, + SkipLookup, +} + +pub fn member_source(facts: &MemberFacts) -> MemberSource { + use MemberSource::*; + if facts.user_id.is_none() || facts.team_id.is_none() || !facts.prisma_present { + return Done; + } + if facts.lookup_failure.is_some() { + return if facts.skip_budget_checks { + SkipLookup + } else { + RaiseLookup + }; + } + if !facts.membership_read { + return ReadMembership; + } + if facts.membership.as_ref().is_some_and(|member| { + member.budget_id.is_some() + && member.budget_id != facts.default_id + && member.model_budget_count.is_some_and(|count| count > 0) + }) { + return Member; + } + if facts.default_id.is_none() { + return Done; + } + if !facts.default_read { + return ReadDefault; + } + if facts + .default_budget + .as_ref() + .is_some_and(|budget| budget.model_budget_count.is_some()) + { + Default + } else { + Clear + } +} + +#[derive(Clone, Debug, Deserialize)] +pub struct IncreaseFacts { + pub number: Option, + pub truthy: bool, + pub type_name: String, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct TemporaryFacts { + pub max_budget: Option, + pub metadata_read: bool, + pub has_increase: bool, + pub has_expiry: bool, + pub expires_at: Option, + pub now: Option, + pub increase: Option, +} + +#[derive(Clone, Copy, Debug, Serialize, PartialEq, Eq, strum::VariantArray)] +pub enum TemporaryRead { + Metadata, + Expiry, + Now, + Increase, +} + +#[derive(Clone, Debug, Default, Serialize, PartialEq)] +pub struct TemporaryDecision { + pub read: Option, + pub copy: bool, + pub max_budget: Option, + pub invalid_increase_type: Option, +} + +fn temp_read(read: TemporaryRead) -> TemporaryDecision { + TemporaryDecision { + read: Some(read), + ..TemporaryDecision::default() + } +} + +pub fn temporary(facts: &TemporaryFacts) -> TemporaryDecision { + let Some(limit) = facts.max_budget else { + return TemporaryDecision::default(); + }; + if !facts.metadata_read { + return temp_read(TemporaryRead::Metadata); + } + if !facts.has_increase || !facts.has_expiry { + return TemporaryDecision::default(); + } + let Some(expiry) = facts.expires_at else { + return temp_read(TemporaryRead::Expiry); + }; + let Some(now) = facts.now else { + return temp_read(TemporaryRead::Now); + }; + if expiry <= now { + return TemporaryDecision::default(); + } + let Some(increase) = &facts.increase else { + return temp_read(TemporaryRead::Increase); + }; + if !increase.truthy { + return TemporaryDecision::default(); + } + match increase.number { + Some(amount) => TemporaryDecision { + copy: true, + max_budget: Some(limit + amount), + ..TemporaryDecision::default() + }, + None => TemporaryDecision { + invalid_increase_type: Some(increase.type_name.clone()), + ..TemporaryDecision::default() + }, + } +} + +#[derive(Clone, Debug, Deserialize)] +pub struct NumericFacts { + pub number: Option, + pub boolean: bool, +} + +pub fn percentage(facts: &NumericFacts) -> Option { + if facts.boolean { + return None; + } + facts.number.filter(|pct| *pct > 0.0 && *pct <= 1.0) +} + +#[derive(Clone, Debug, Deserialize)] +pub struct ThrottleFacts { + pub opt_in: bool, + pub tpm_limit: Option, + pub rpm_limit: Option, + pub percentage: Option, +} + +#[derive(Clone, Debug, Default, Serialize, PartialEq)] +pub struct ThrottleDecision { + pub read_percentage: bool, + pub percentage: Option, +} + +pub fn throttle(facts: &ThrottleFacts) -> ThrottleDecision { + if !facts.opt_in || (facts.tpm_limit.is_none() && facts.rpm_limit.is_none()) { + return ThrottleDecision::default(); + } + match &facts.percentage { + None => ThrottleDecision { + read_percentage: true, + ..ThrottleDecision::default() + }, + Some(raw) => ThrottleDecision { + percentage: percentage(raw), + ..ThrottleDecision::default() + }, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::team_empty(ModelScope::Team, Some(""), vec![0, 1])] + #[case::team_none(ModelScope::Team, None, vec![])] + #[case::enduser_empty(ModelScope::EndUser, Some(""), vec![])] + #[case::enduser(ModelScope::EndUser, Some("user"), vec![0, 1])] + fn model_identity_gates( + #[case] scope: ModelScope, + #[case] identity: Option<&str>, + #[case] expected: Vec, + ) { + assert_eq!( + model_checks(&ModelCheckFacts { + scope, + entity_id: identity.map(str::to_owned), + budget_count: 1, + model_count: 2 + }), + expected + ); + } + + #[rstest] + #[case::zero(0.0, false, None)] + #[case::negative(-0.5, false, None)] + #[case::nan(f64::NAN, false, None)] + #[case::infinite(f64::INFINITY, false, None)] + #[case::one(1.0, false, Some(1.0))] + #[case::boolean(1.0, true, None)] + fn numeric_percentages( + #[case] number: f64, + #[case] boolean: bool, + #[case] expected: Option, + ) { + assert_eq!( + percentage(&NumericFacts { + number: Some(number), + boolean + }), + expected + ); + } + + #[rstest] + #[case::expired(10, 10, 3.0, false, None)] + #[case::before(9, 10, 3.0, false, None)] + #[case::active(11, 10, 3.0, true, Some(13.0))] + #[case::negative(11, 10, -3.0, true, Some(7.0))] + #[case::zero(11, 10, 0.0, false, None)] + fn temporary_copy_boundary( + #[case] expiry: i64, + #[case] now: i64, + #[case] amount: f64, + #[case] copy: bool, + #[case] expected: Option, + ) { + let facts = TemporaryFacts { + max_budget: Some(10.0), + metadata_read: true, + has_increase: true, + has_expiry: true, + expires_at: Some(expiry), + now: Some(now), + increase: Some(IncreaseFacts { + number: Some(amount), + truthy: amount != 0.0, + type_name: "float".to_owned(), + }), + }; + let result = temporary(&facts); + assert_eq!( + (result.copy, result.max_budget, result.read), + (copy, expected, None) + ); + } + + #[rstest] + #[case::member_override(Some("member"), Some(1), Some("default"), false, MemberSource::Member)] + #[case::no_override(None, Some(1), None, false, MemberSource::Done)] + #[case::empty_override( + Some("member"), + Some(0), + Some("default"), + false, + MemberSource::ReadDefault + )] + #[case::current_default( + Some("default"), + Some(1), + Some("default"), + false, + MemberSource::ReadDefault + )] + #[case::missing_default(None, None, Some("default"), true, MemberSource::Clear)] + fn member_default_staging( + #[case] member_id: Option<&str>, + #[case] count: Option, + #[case] default_id: Option<&str>, + #[case] default_read: bool, + #[case] expected: MemberSource, + ) { + let facts = MemberFacts { + user_id: Some(String::new()), + team_id: Some(String::new()), + prisma_present: true, + default_id: default_id.map(str::to_owned), + skip_budget_checks: false, + membership_read: true, + membership: Some(MembershipSnapshot { + budget_id: member_id.map(str::to_owned), + model_budget_count: count, + }), + default_read, + default_budget: None, + lookup_failure: None, + }; + assert_eq!(member_source(&facts), expected); + } + + #[rstest] + #[case::replace(false, true, vec![TokenField::MaxBudget, TokenField::Tpm, TokenField::Rpm, TokenField::Tpd, TokenField::Models])] + #[case::keep_empty(true, true, vec![TokenField::Rpm, TokenField::Tpd])] + #[case::no_default(false, false, vec![])] + fn token_assignment_order( + #[case] keep: bool, + #[case] present: bool, + #[case] expected: Vec, + ) { + let facts = TokenDefaultsFacts { + current: BudgetFields { + max_budget: Some(0.0), + tpm_limit: Some(0), + model_budget_count: Some(0), + ..BudgetFields::default() + }, + keep_token_limits: keep, + default_present: present, + }; + assert_eq!(token_assignments(&facts), expected); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/lib.rs b/litellm-rust/crates/request-auth-core/src/lib.rs index e74b05be20a..eaf650b66c2 100644 --- a/litellm-rust/crates/request-auth-core/src/lib.rs +++ b/litellm-rust/crates/request-auth-core/src/lib.rs @@ -1,6 +1,7 @@ #![forbid(unsafe_code)] pub mod budget_alerts; +pub mod budget_context; pub mod budget_pricing; pub mod common_budgets; pub mod credential; diff --git a/litellm/_auth/budget_context.py b/litellm/_auth/budget_context.py new file mode 100644 index 00000000000..ef67468621b --- /dev/null +++ b/litellm/_auth/budget_context.py @@ -0,0 +1,174 @@ +from dataclasses import dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class ModelCheckFacts: + scope: Literal["Team", "EndUser"] + entity_id: str | None + budget_count: int + model_count: int + + +@dataclass(frozen=True, slots=True) +class BudgetFields: + max_budget: float | None = None + tpm_limit: int | None = None + rpm_limit: int | None = None + tpd_limit: int | None = None + model_budget_count: int | None = None + + +@dataclass(frozen=True, slots=True) +class EndUserFacts: + mode: Literal["Models", "Token", "Default"] + object_present: bool + end_user_id: str | None + prisma_present: bool + key_budget_id: str | None + key_read: bool = False + key_budget: BudgetFields | None = None + global_config_read: bool = False + global_budget_id: str | None = None + global_read: bool = False + global_budget: BudgetFields | None = None + + +EndUserSource: TypeAlias = Literal["Own", "ReadKey", "Key", "ReadGlobalConfig", "ReadGlobal", "Global", "None"] + + +@dataclass(frozen=True, slots=True) +class TokenDefaultsFacts: + current: BudgetFields + keep_token_limits: bool + default_present: bool + + +TokenField: TypeAlias = Literal["MaxBudget", "Tpm", "Rpm", "Tpd", "Models"] + + +@dataclass(frozen=True, slots=True) +class MembershipSnapshot: + budget_id: str | None + model_budget_count: int | None + + +@dataclass(frozen=True, slots=True) +class MemberFacts: + user_id: str | None + team_id: str | None + prisma_present: bool + default_id: str | None + skip_budget_checks: bool + membership_read: bool = False + membership: MembershipSnapshot | None = None + default_read: bool = False + default_budget: BudgetFields | None = None + lookup_failure: Literal["Membership", "Default"] | None = None + + +MemberSource: TypeAlias = Literal[ + "Done", "ReadMembership", "ReadDefault", "Member", "Default", "Clear", "RaiseLookup", "SkipLookup" +] + + +@dataclass(frozen=True, slots=True) +class IncreaseFacts: + number: float | None + truthy: bool + type_name: str + + +@dataclass(frozen=True, slots=True) +class TemporaryFacts: + max_budget: float | None + metadata_read: bool = False + has_increase: bool = False + has_expiry: bool = False + expires_at: int | None = None + now: int | None = None + increase: IncreaseFacts | None = None + + +@dataclass(frozen=True, slots=True) +class TemporaryDecision: + read: Literal["Metadata", "Expiry", "Now", "Increase"] | None = None + copy: bool = False + max_budget: float | None = None + invalid_increase_type: str | None = None + + +@dataclass(frozen=True, slots=True) +class NumericFacts: + number: float | None + boolean: bool + + +@dataclass(frozen=True, slots=True) +class ThrottleFacts: + opt_in: bool + tpm_limit: int | None + rpm_limit: int | None + percentage: NumericFacts | None = None + + +@dataclass(frozen=True, slots=True) +class ThrottleDecision: + read_percentage: bool = False + percentage: float | None = None + + +_MODELS: Final = TypeAdapter(ModelCheckFacts) +_INDICES: Final = TypeAdapter(tuple[int, ...]) +_END: Final = TypeAdapter(EndUserFacts) +_END_SOURCE: Final = TypeAdapter[EndUserSource](EndUserSource) +_TOKEN: Final = TypeAdapter(TokenDefaultsFacts) +_FIELDS: Final = TypeAdapter(tuple[TokenField, ...]) +_MEMBER: Final = TypeAdapter(MemberFacts) +_MEMBER_SOURCE: Final = TypeAdapter[MemberSource](MemberSource) +_TEMP: Final = TypeAdapter(TemporaryFacts) +_TEMP_DECISION: Final = TypeAdapter(TemporaryDecision) +_THROTTLE: Final = TypeAdapter(ThrottleFacts) +_THROTTLE_DECISION: Final = TypeAdapter(ThrottleDecision) +_NUMBER: Final = TypeAdapter(NumericFacts) +_PERCENTAGE: Final = TypeAdapter[float | None](float | None) + + +@native_boundary +def model_checks(facts: ModelCheckFacts) -> tuple[int, ...]: + return _INDICES.validate_python(_native.request_auth_context_models(_MODELS.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def end_user_source(facts: EndUserFacts) -> EndUserSource: + return _END_SOURCE.validate_python(_native.request_auth_context_end_user(_END.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def token_assignments(facts: TokenDefaultsFacts) -> tuple[TokenField, ...]: + return _FIELDS.validate_python(_native.request_auth_context_token(_TOKEN.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def member_source(facts: MemberFacts) -> MemberSource: + return _MEMBER_SOURCE.validate_python(_native.request_auth_context_member(_MEMBER.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def temporary(facts: TemporaryFacts) -> TemporaryDecision: + return _TEMP_DECISION.validate_python(_native.request_auth_context_temporary(_TEMP.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def percentage(facts: NumericFacts) -> float | None: + return _PERCENTAGE.validate_python(_native.request_auth_context_percentage(_NUMBER.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts + + +@native_boundary +def throttle(facts: ThrottleFacts) -> ThrottleDecision: + return _THROTTLE_DECISION.validate_python(_native.request_auth_context_throttle(_THROTTLE.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts diff --git a/litellm/proxy/auth/budget_context_adapter.py b/litellm/proxy/auth/budget_context_adapter.py new file mode 100644 index 00000000000..aab2286515b --- /dev/null +++ b/litellm/proxy/auth/budget_context_adapter.py @@ -0,0 +1,377 @@ +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, replace +from typing import Final, Protocol, assert_never, cast # noqa: TID251 # preserve identity of existing JSON mappings + +from opentelemetry.trace import Span + +import litellm +from litellm._auth import budget_context as core +from litellm._auth.failure import NativeAuthFailure +from litellm._logging import verbose_proxy_logger +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.proxy._types import LiteLLM_EndUserTable, LiteLLM_TeamMembership, LiteLLM_TeamTable, UserAPIKeyAuth +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging + + +class TeamModelLimiter(Protocol): + async def is_team_within_model_budget( + self, + team_id: str, + team_model_max_budget: Mapping[str, object], + key_model_max_budget: Mapping[str, object] | None, + model: str, + ) -> bool: ... + + +class EndUserModelLimiter(Protocol): + async def is_end_user_within_model_budget( + self, + end_user_id: str, + end_user_model_max_budget: Mapping[str, object], + model: str, + team_id: str | None = None, + ) -> bool: ... + + +def mapping_value(value: object) -> Mapping[str, object] | None: + if isinstance(value, Mapping): + return cast(Mapping[str, object], value) # cast-ok: preserve the original JSON map's identity + return None + + +def mapping_count(value: object) -> int | None: + mapping: Final = mapping_value(value) + return len(mapping) if mapping is not None else None + + +def budget_fields(row: LiteLLM_BudgetTable | None) -> core.BudgetFields | None: + if row is None: + return None + return core.BudgetFields( + row.max_budget, + row.tpm_limit, + row.rpm_limit, + row.tpd_limit, + mapping_count(getattr(row, "model_max_budget", None)), + ) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # legacy budget table declares an unparameterized dict + + +async def check_team_model_budget( + valid_token: UserAPIKeyAuth, + model_max_budget_limiter: TeamModelLimiter, + models: Sequence[str], +) -> None: + budgets: Final = valid_token.team_model_max_budget + selected: Final = core.model_checks( + core.ModelCheckFacts("Team", valid_token.team_id, len(budgets) if budgets is not None else 0, len(models)) + ) + for index in selected: + assert valid_token.team_id is not None and budgets is not None + await model_max_budget_limiter.is_team_within_model_budget( + team_id=valid_token.team_id, + team_model_max_budget=budgets, + key_model_max_budget=mapping_value(getattr(valid_token, "model_max_budget", None)), + model=models[index], + ) + + +async def check_end_user_model_budget( + end_user_id: str | None, + end_user_model_max_budget: Mapping[str, object] | None, + model_max_budget_limiter: EndUserModelLimiter, + models: Sequence[str], + team_id: str | None, +) -> None: + selected: Final = core.model_checks( + core.ModelCheckFacts( + "EndUser", + end_user_id, + len(end_user_model_max_budget) if end_user_model_max_budget is not None else 0, + len(models), + ) + ) + for index in selected: + assert end_user_id is not None and end_user_model_max_budget is not None + await model_max_budget_limiter.is_end_user_within_model_budget( + end_user_id=end_user_id, + end_user_model_max_budget=end_user_model_max_budget, + model=models[index], + team_id=team_id, + ) + + +@dataclass(frozen=True, slots=True) +class EndUserIO: + prisma_client: PrismaClient | None + cache: UserApiKeyCache + span: Span | None + + +async def _end_user_row( + facts: core.EndUserFacts, + io: EndUserIO, + own: LiteLLM_EndUserTable | None, + key: LiteLLM_BudgetTable | None = None, + proxy: LiteLLM_BudgetTable | None = None, +) -> LiteLLM_BudgetTable | None: + from litellm.proxy.auth.auth_checks import get_default_end_user_budget + + decision: Final = core.end_user_source(facts) + match decision: + case "Own": + assert own is not None + return own.litellm_budget_table + case "Key": + return key + case "Global": + return proxy + case "None": + return None + case "ReadKey": + row: Final = await get_default_end_user_budget( + prisma_client=io.prisma_client, + user_api_key_cache=io.cache, + parent_otel_span=io.span, + budget_id=facts.key_budget_id, + ) + return await _end_user_row( + replace(facts, key_read=True, key_budget=budget_fields(row)), io, own, row, proxy + ) + case "ReadGlobalConfig": + return await _end_user_row( + replace(facts, global_config_read=True, global_budget_id=litellm.max_end_user_budget_id), + io, + own, + key, + proxy, + ) + case "ReadGlobal": + row_global: Final = await get_default_end_user_budget( + prisma_client=io.prisma_client, + user_api_key_cache=io.cache, + parent_otel_span=io.span, + budget_id=facts.global_budget_id, + ) + return await _end_user_row( + replace(facts, global_read=True, global_budget=budget_fields(row_global)), io, own, key, row_global + ) + case _: + assert_never(decision) + + +async def resolve_default_end_user_budget( + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + key_end_user_budget_id: str | None, + parent_otel_span: Span | None = None, +) -> LiteLLM_BudgetTable | None: + return await _end_user_row( + core.EndUserFacts("Default", False, None, prisma_client is not None, key_end_user_budget_id), # pyright: ignore[reportUnnecessaryComparison] # legacy callers can pass absent database clients at runtime + EndUserIO(prisma_client, user_api_key_cache, parent_otel_span), + None, + ) + + +async def end_user_model_max_budget( + end_user_id: str | None, + end_user_object: LiteLLM_EndUserTable | None, + key_end_user_budget_id: str | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, +) -> Mapping[str, object] | None: + row: Final = await _end_user_row( + core.EndUserFacts( + "Models", end_user_object is not None, end_user_id, prisma_client is not None, key_end_user_budget_id + ), + EndUserIO(prisma_client, user_api_key_cache, parent_otel_span), + end_user_object, + ) + return getattr(row, "model_max_budget", None) if row is not None else None # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # legacy row field + + +async def apply_key_end_user_default_budget_to_token( + valid_token: UserAPIKeyAuth, + end_user_object: LiteLLM_EndUserTable | None, + key_end_user_budget_id: str, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + parent_otel_span: Span | None, + keep_token_limits: bool, +) -> None: + row: Final = await _end_user_row( + core.EndUserFacts("Token", end_user_object is not None, valid_token.end_user_id, True, key_end_user_budget_id), + EndUserIO(prisma_client, user_api_key_cache, parent_otel_span), + end_user_object, + ) + current: Final = core.BudgetFields( + valid_token.end_user_max_budget, + valid_token.end_user_tpm_limit, + valid_token.end_user_rpm_limit, + valid_token.end_user_tpd_limit, + mapping_count(valid_token.end_user_model_max_budget), + ) + for field in core.token_assignments(core.TokenDefaultsFacts(current, keep_token_limits, row is not None)): + assert row is not None + match field: + case "MaxBudget": + valid_token.end_user_max_budget = ( # rebind-ok: request-local native projection + row.max_budget + ) + case "Tpm": + valid_token.end_user_tpm_limit = ( # rebind-ok: request-local native projection + row.tpm_limit + ) + case "Rpm": + valid_token.end_user_rpm_limit = ( # rebind-ok: request-local native projection + row.rpm_limit + ) + case "Tpd": + valid_token.end_user_tpd_limit = ( # rebind-ok: request-local native projection + row.tpd_limit + ) + case "Models": + valid_token.end_user_model_max_budget = row.model_max_budget # pyright: ignore[reportUnknownMemberType] # preserve the validated row dict and its identity # rebind-ok: native-selected request-local budget assignment + case _: + assert_never(field) + + +@dataclass(frozen=True, slots=True) +class MemberIO: + prisma_client: PrismaClient | None + cache: UserApiKeyCache + logging: ProxyLogging | None + + +@dataclass(frozen=True, slots=True) +class MemberRead: + member: LiteLLM_TeamMembership | None = None + default: LiteLLM_BudgetTable | None = None + error: Exception | None = None + + +async def _read_member(facts: core.MemberFacts, io: MemberIO) -> MemberRead: + from litellm.proxy.auth.auth_checks import get_team_membership + + assert facts.user_id is not None and facts.team_id is not None + try: + row: Final = await get_team_membership( + user_id=facts.user_id, + team_id=facts.team_id, + prisma_client=io.prisma_client, + user_api_key_cache=io.cache, + proxy_logging_obj=io.logging, + ) + except NativeAuthFailure: + raise + except Exception as error: # noqa: BLE001 # only the database/cache read is caught; native failures rethrow + return MemberRead(error=error) + return MemberRead(member=row) + + +async def _read_member_default(facts: core.MemberFacts, io: MemberIO) -> MemberRead: + from litellm.proxy.auth.auth_checks import get_team_member_default_budget + + assert facts.default_id is not None + try: + row: Final = await get_team_member_default_budget( + budget_id=facts.default_id, + prisma_client=io.prisma_client, + user_api_key_cache=io.cache, + raise_on_lookup_error=True, + ) + except NativeAuthFailure: + raise + except Exception as error: # noqa: BLE001 # only the database/cache read is caught; native failures rethrow + return MemberRead(error=error) + return MemberRead(default=row) + + +async def _member_model_budget( + token: UserAPIKeyAuth, + facts: core.MemberFacts, + io: MemberIO, + reading: MemberRead, +) -> None: + decision: Final = core.member_source(facts) + match decision: + case "Done": + return + case "RaiseLookup": + assert reading.error is not None + raise reading.error + case "SkipLookup": + verbose_proxy_logger.warning( + "Skipping team-member model budget lookup after a lookup error", exc_info=reading.error + ) + return + case "ReadMembership": + member_read: Final = await _read_member(facts, io) + member: Final = member_read.member + member_row: Final = member.litellm_budget_table if member is not None else None + snapshot: Final = ( + core.MembershipSnapshot(member.budget_id, mapping_count(getattr(member_row, "model_max_budget", None))) + if member is not None + else None + ) + await _member_model_budget( + token, + replace( + facts, + membership_read=True, + membership=snapshot, + lookup_failure="Membership" if member_read.error is not None else None, + ), + io, + member_read, + ) + case "ReadDefault": + default_read: Final = await _read_member_default(facts, io) + await _member_model_budget( + token, + replace( + facts, + default_read=True, + default_budget=budget_fields(default_read.default), + lookup_failure="Default" if default_read.error is not None else None, + ), + io, + default_read, + ) + case "Member": + assert reading.member is not None + token.team_member_model_max_budget = mapping_value( # rebind-ok: request-local native projection + getattr(reading.member.litellm_budget_table, "model_max_budget", None) + ) + case "Default": + assert reading.default is not None + token.team_member_model_max_budget = mapping_value( # rebind-ok: request-local native projection + getattr(reading.default, "model_max_budget", None) + ) + case "Clear": + token.team_member_model_max_budget = None # rebind-ok: request-local native projection + case _: + assert_never(decision) + + +async def load_team_member_default_model_budget( + user_api_key_auth_obj: UserAPIKeyAuth, + team_object: LiteLLM_TeamTable | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + skip_budget_checks: bool, + proxy_logging_obj: ProxyLogging | None, +) -> None: + metadata: Final = mapping_value(getattr(team_object, "metadata", None)) + default_id: Final = metadata.get("team_member_budget_id") if metadata is not None else None + facts: Final = core.MemberFacts( + user_id=user_api_key_auth_obj.user_id if isinstance(user_api_key_auth_obj.user_id, str) else None, + team_id=user_api_key_auth_obj.team_id if isinstance(user_api_key_auth_obj.team_id, str) else None, + prisma_present=prisma_client is not None, + default_id=default_id if isinstance(default_id, str) else None, + skip_budget_checks=skip_budget_checks, + ) + await _member_model_budget( + user_api_key_auth_obj, facts, MemberIO(prisma_client, user_api_key_cache, proxy_logging_obj), MemberRead() + ) diff --git a/litellm/proxy/auth/key_budget_context_adapter.py b/litellm/proxy/auth/key_budget_context_adapter.py new file mode 100644 index 00000000000..3f7d245704b --- /dev/null +++ b/litellm/proxy/auth/key_budget_context_adapter.py @@ -0,0 +1,95 @@ +from collections.abc import Mapping +from dataclasses import replace +from datetime import datetime, timezone +from typing import Final, assert_never + +import litellm +from litellm._auth import budget_context as core +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.budget_context_adapter import mapping_value + + +def numeric_facts(value: object) -> core.NumericFacts: + return core.NumericFacts(float(value) if isinstance(value, (int, float)) else None, isinstance(value, bool)) + + +def budget_throttle_percentage() -> float | None: + return core.percentage(numeric_facts(litellm.budget_exceeded_throttle_percentage)) + + +def read_key_throttle_facts(valid_token: UserAPIKeyAuth) -> core.ThrottleFacts: + metadata: Final = mapping_value(getattr(valid_token, "metadata", None)) + facts: Final = core.ThrottleFacts( + (metadata or {}).get("throttle_on_budget_exceeded") is True, + valid_token.tpm_limit, + valid_token.rpm_limit, + ) + decision: Final = core.throttle(facts) + return ( + replace(facts, percentage=numeric_facts(litellm.budget_exceeded_throttle_percentage)) + if decision.read_percentage + else facts + ) + + +def should_throttle_budget_exceeded(valid_token: UserAPIKeyAuth) -> bool: + return core.throttle(read_key_throttle_facts(valid_token)).percentage is not None + + +def inspect_key_budget_throttle(valid_token: UserAPIKeyAuth) -> float | None: + decision: Final = core.throttle(read_key_throttle_facts(valid_token)) + return budget_throttle_percentage() if decision.percentage is not None else None + + +def _microseconds(value: datetime) -> int: + delta: Final = value - datetime(1970, 1, 1, tzinfo=timezone.utc) + return (delta.days * 86400 + delta.seconds) * 1_000_000 + delta.microseconds + + +def _temporary( + token: UserAPIKeyAuth, + facts: core.TemporaryFacts, + metadata: Mapping[str, object] | None, +) -> UserAPIKeyAuth: + decision: Final = core.temporary(facts) + if decision.invalid_increase_type is not None: + raise TypeError( + f"unsupported operand type(s) for +: '{type(token.max_budget).__name__}' and '{decision.invalid_increase_type}'" + ) + if decision.copy: + return token.model_copy(update={"max_budget": decision.max_budget}) + match decision.read: + case None: + return token + case "Metadata": + raw: Final = mapping_value(getattr(token, "metadata", None)) + return _temporary( + token, + replace( + facts, + metadata_read=True, + has_increase=raw is not None and "temp_budget_increase" in raw, + has_expiry=raw is not None and "temp_budget_expiry" in raw, + ), + raw, + ) + case "Expiry": + assert metadata is not None + expiry: Final = datetime.fromisoformat(metadata["temp_budget_expiry"]) # pyright: ignore[reportArgumentType] # preserve datetime's public TypeError for malformed metadata + aware: Final = expiry.replace(tzinfo=timezone.utc) if expiry.tzinfo is None else expiry + return _temporary(token, replace(facts, expires_at=_microseconds(aware)), metadata) + case "Now": + return _temporary(token, replace(facts, now=_microseconds(datetime.now(timezone.utc))), metadata) + case "Increase": + assert metadata is not None + value: Final = metadata["temp_budget_increase"] + raw_increase: Final = core.IncreaseFacts( + float(value) if isinstance(value, (int, float)) else None, bool(value), type(value).__name__ + ) + return _temporary(token, replace(facts, increase=raw_increase), metadata) + case _: + assert_never(decision.read) + + +def update_key_budget_with_temp_budget_increase(valid_token: UserAPIKeyAuth) -> UserAPIKeyAuth: + return _temporary(valid_token, core.TemporaryFacts(valid_token.max_budget), None) diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 2c889910359..c8e429ab922 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -812,3 +812,17 @@ class RouterCall(Generic[_S]): def request_auth_budget_exemption(facts: object) -> dict[str, object]: ... def request_auth_budget_alert_plan(facts: object) -> dict[str, object]: ... + +def request_auth_context_models(facts: object) -> object: ... + +def request_auth_context_end_user(facts: object) -> object: ... + +def request_auth_context_token(facts: object) -> object: ... + +def request_auth_context_member(facts: object) -> object: ... + +def request_auth_context_temporary(facts: object) -> object: ... + +def request_auth_context_throttle(facts: object) -> object: ... + +def request_auth_context_percentage(facts: object) -> float | None: ... diff --git a/tests/unit/proxy/auth/test_budget_context_adapter.py b/tests/unit/proxy/auth/test_budget_context_adapter.py new file mode 100644 index 00000000000..707b918352c --- /dev/null +++ b/tests/unit/proxy/auth/test_budget_context_adapter.py @@ -0,0 +1,327 @@ +from collections.abc import Mapping +from typing import Final, Literal +from unittest.mock import AsyncMock, Mock + +import pytest + +import litellm +from litellm._auth import budget_context as core +from litellm._auth.failure import NativeAuthRuntimeError +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.proxy._types import LiteLLM_EndUserTable, LiteLLM_TeamMembership, LiteLLM_TeamTable, UserAPIKeyAuth +from litellm.proxy.auth import auth_checks +from litellm.proxy.auth import budget_context_adapter as adapter +from litellm.proxy.auth import user_api_key_auth as legacy +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + +@pytest.fixture(params=(False, True), ids=("reference", "native")) +def native(request: pytest.FixtureRequest) -> bool: + return request.param is True # pyright: ignore[reportAny] # fixture parameters are the two literal booleans above + + +@pytest.mark.parametrize("scope", ("Team", "EndUser")) +@pytest.mark.parametrize("identity", (None, "", "entity")) +@pytest.mark.parametrize("budget_count", (0, 1, 2)) +@pytest.mark.parametrize("model_count", (0, 1, 3)) +def test_native_model_read_selection( + scope: Literal["Team", "EndUser"], identity: str | None, budget_count: int, model_count: int +) -> None: + facts: Final = core.ModelCheckFacts(scope, identity, budget_count, model_count) + expected: Final = ( + tuple(range(model_count)) if budget_count and identity is not None and (scope == "Team" or identity) else () + ) + assert core.model_checks(facts) == expected + + +@pytest.mark.parametrize("keep", (False, True)) +@pytest.mark.parametrize("present", (False, True)) +@pytest.mark.parametrize("limit", (None, 0, -1, 1)) +@pytest.mark.parametrize("models", (None, 0, 2)) +def test_native_token_assignment_table(keep: bool, present: bool, limit: int | None, models: int | None) -> None: + facts: Final = core.TokenDefaultsFacts(core.BudgetFields(limit, limit, limit, limit, models), keep, present) + expected: Final = tuple( + name + for name, value in (("MaxBudget", limit), ("Tpm", limit), ("Rpm", limit), ("Tpd", limit), ("Models", models)) + if present and (not keep or value is None) + ) + assert core.token_assignments(facts) == expected + + +@pytest.mark.parametrize("mode", ("Models", "Token", "Default")) +@pytest.mark.parametrize("object_present", (False, True)) +@pytest.mark.parametrize("identity", (None, "", "eu")) +@pytest.mark.parametrize("db", (False, True)) +@pytest.mark.parametrize("key_id", (None, "", "key")) +def test_end_user_initial_table( + mode: Literal["Models", "Token", "Default"], + object_present: bool, + identity: str | None, + db: bool, + key_id: str | None, +) -> None: + facts: Final = core.EndUserFacts(mode, object_present, identity, db, key_id) + expected: Final = ( + "Own" + if mode != "Default" and object_present + else "None" + if mode == "Models" and (not identity or not db) + else "ReadKey" + if key_id is not None + else "ReadGlobalConfig" + ) + assert core.end_user_source(facts) == expected + + +@pytest.mark.parametrize("key_row", (False, True)) +@pytest.mark.parametrize("global_id", (None, "", "proxy")) +@pytest.mark.parametrize("global_read", (False, True)) +@pytest.mark.parametrize("global_row", (False, True)) +def test_end_user_fallback_stages(key_row: bool, global_id: str | None, global_read: bool, global_row: bool) -> None: + facts: Final = core.EndUserFacts( + "Default", + False, + None, + True, + "key", + key_read=True, + key_budget=core.BudgetFields() if key_row else None, + global_config_read=True, + global_budget_id=global_id, + global_read=global_read, + global_budget=core.BudgetFields() if global_row else None, + ) + expected: Final = ( + "Key" + if key_row + else "None" + if global_id is None + else "ReadGlobal" + if not global_read + else "Global" + if global_row + else "None" + ) + assert core.end_user_source(facts) == expected + + +@pytest.mark.parametrize("default_id", (None, "", "default")) +@pytest.mark.parametrize("member_id", (None, "", "member", "default")) +@pytest.mark.parametrize("member_models", (None, 0, 2)) +@pytest.mark.parametrize("default_models", (None, 0, 2)) +@pytest.mark.parametrize("skip", (False, True)) +def test_member_snapshot_table( + default_id: str | None, member_id: str | None, member_models: int | None, default_models: int | None, skip: bool +) -> None: + facts: Final = core.MemberFacts( + "", + "", + True, + default_id, + skip, + membership_read=True, + membership=core.MembershipSnapshot(member_id, member_models), + default_read=True, + default_budget=core.BudgetFields(model_budget_count=default_models), + ) + expected: Final = ( + "Member" + if member_id is not None and member_id != default_id and member_models + else "Done" + if default_id is None + else "Default" + if default_models is not None + else "Clear" + ) + assert core.member_source(facts) == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize("identity", (None, "", "entity")) +@pytest.mark.parametrize("scope", ("Team", "EndUser")) +@pytest.mark.parametrize("fail", (False, True)) +async def test_selected_model_reads_keep_order_and_stop_on_failure( + native: bool, identity: str | None, scope: Literal["Team", "EndUser"], fail: bool +) -> None: + failure: Final = RuntimeError("counter failed") + read: Final = AsyncMock(side_effect=(True, failure, True) if fail else None, return_value=True) + limiter: Final = Mock() + limiter.is_team_within_model_budget = read + limiter.is_end_user_within_model_budget = read + caps: Final[Mapping[str, object]] = {"a": {"max_budget": 1}} + token: Final = UserAPIKeyAuth(team_id=identity, team_model_max_budget=dict(caps)) + models: Final = ["a", "a", "b"] + coroutine: Final = ( + (adapter.check_team_model_budget if native else legacy._check_team_model_budget)(token, limiter, models) # pyright: ignore[reportPrivateUsage] # Phase 1 leaf is the unchanged reference implementation + if scope == "Team" + else (adapter.check_end_user_model_budget if native else legacy._check_end_user_model_budget)( # pyright: ignore[reportPrivateUsage] # Phase 1 leaf is the unchanged reference implementation + identity, caps, limiter, models, "team" + ) + ) + active: Final = identity is not None and (scope == "Team" or bool(identity)) + if fail and active: + with pytest.raises(RuntimeError, match="counter failed") as raised: + await coroutine + assert raised.value is failure + else: + await coroutine + assert read.await_count == (2 if fail else 3) if active else read.await_count == 0 + assert tuple(item.kwargs["model"] for item in read.await_args_list) == tuple(models[: read.await_count]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("own", (False, True)) +@pytest.mark.parametrize("own_row", (False, True)) +@pytest.mark.parametrize("identity", (None, "", "eu")) +@pytest.mark.parametrize("key_found", (False, True)) +async def test_end_user_source_read_order_and_identity( + native: bool, monkeypatch: pytest.MonkeyPatch, own: bool, own_row: bool, identity: str | None, key_found: bool +) -> None: + budget: Final = LiteLLM_BudgetTable(model_max_budget={"m": {"max_budget": 1}}) + end_user: Final = ( + LiteLLM_EndUserTable(user_id="eu", blocked=False, litellm_budget_table=budget if own_row else None) + if own + else None + ) + read: Final = AsyncMock(side_effect=(budget,) if key_found else (None, budget)) + monkeypatch.setattr(auth_checks, "get_default_end_user_budget", read) + monkeypatch.setattr(litellm, "max_end_user_budget_id", "proxy") + fn: Final = adapter.end_user_model_max_budget if native else legacy._end_user_model_max_budget # pyright: ignore[reportPrivateUsage] # Phase 1 leaf is the unchanged reference implementation + value: Final = await fn(identity, end_user, "key", Mock(), UserApiKeyCache(), None) + expected: Final = getattr(budget, "model_max_budget", None) if (own and own_row) or (not own and identity) else None + assert value is expected + assert read.await_count == (0 if own or not identity else 1 if key_found else 2) + if read.await_count: + assert read.await_args_list[0].kwargs["budget_id"] == "key" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("keep", (False, True)) +@pytest.mark.parametrize("present", (False, True)) +async def test_token_defaults_preserve_caps_and_projection_identity(native: bool, keep: bool, present: bool) -> None: + budget: Final = LiteLLM_BudgetTable( + max_budget=9, tpm_limit=8, rpm_limit=7, tpd_limit=6, model_max_budget={"m": {"max_budget": 5}} + ) + token: Final = UserAPIKeyAuth( + end_user_max_budget=0 if present else None, + end_user_tpm_limit=0 if present else None, + end_user_model_max_budget={} if present else None, + ) + original_models: Final = token.end_user_model_max_budget + own: Final = LiteLLM_EndUserTable(user_id="eu", blocked=False, litellm_budget_table=budget) + fn: Final = ( + adapter.apply_key_end_user_default_budget_to_token + if native + else legacy._apply_key_end_user_default_budget_to_token # pyright: ignore[reportPrivateUsage] # Phase 1 leaf is the unchanged reference implementation + ) + await fn(token, own, "key", Mock(), UserApiKeyCache(), None, keep) + assert token.end_user_max_budget == (0 if keep and present else 9) + assert token.end_user_tpm_limit == (0 if keep and present else 8) + assert (token.end_user_rpm_limit, token.end_user_tpd_limit) == (7, 6) + assert token.end_user_model_max_budget is ( + original_models if keep and present else getattr(budget, "model_max_budget", None) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("default_id", (None, "", "default")) +@pytest.mark.parametrize("member_id", (None, "default", "override")) +@pytest.mark.parametrize("member_caps", (None, {}, {"member": {"max_budget": 1}})) +@pytest.mark.parametrize("default_caps", (None, {}, {"default": {"max_budget": 2}})) +async def test_member_default_read_and_mutation_order( + native: bool, + monkeypatch: pytest.MonkeyPatch, + default_id: str | None, + member_id: str | None, + member_caps: Mapping[str, object] | None, + default_caps: Mapping[str, object] | None, +) -> None: + member_budget: Final = LiteLLM_BudgetTable(model_max_budget=dict(member_caps) if member_caps is not None else None) + default_budget: Final = LiteLLM_BudgetTable( + model_max_budget=dict(default_caps) if default_caps is not None else None + ) + membership: Final = LiteLLM_TeamMembership( + user_id="", team_id="", budget_id=member_id, litellm_budget_table=member_budget + ) + member_read: Final = AsyncMock(return_value=membership) + default_read: Final = AsyncMock(return_value=default_budget) + monkeypatch.setattr(auth_checks, "get_team_membership", member_read) + monkeypatch.setattr(legacy, "get_team_membership", member_read) + monkeypatch.setattr(auth_checks, "get_team_member_default_budget", default_read) + monkeypatch.setattr(legacy, "get_team_member_default_budget", default_read) + token: Final = UserAPIKeyAuth(user_id="", team_id="", team_member_model_max_budget={"old": {}}) + prior: Final = token.team_member_model_max_budget + team: Final = LiteLLM_TeamTable(team_id="", metadata={"team_member_budget_id": default_id}) + fn: Final = ( + adapter.load_team_member_default_model_budget if native else legacy._load_team_member_default_model_budget # pyright: ignore[reportPrivateUsage] # Phase 1 leaf is the unchanged reference implementation + ) + await fn(token, team, Mock(), UserApiKeyCache(), False, None) + override: Final = member_id is not None and member_id != default_id and bool(member_caps) + expected: Final = ( + getattr(member_budget, "model_max_budget", None) + if override + else getattr(default_budget, "model_max_budget", None) + if default_id is not None + else prior + ) + assert token.team_member_model_max_budget is expected + member_read.assert_awaited_once() + assert default_read.await_count == (0 if override or default_id is None else 1) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("skip", (False, True)) +@pytest.mark.parametrize("during_default", (False, True)) +async def test_member_lookup_error_policy( + native: bool, monkeypatch: pytest.MonkeyPatch, skip: bool, during_default: bool +) -> None: + error: Final = RuntimeError("budget read failed") + member_read: Final = AsyncMock(return_value=None, side_effect=None if during_default else error) + default_read: Final = AsyncMock(side_effect=error) + monkeypatch.setattr(auth_checks, "get_team_membership", member_read) + monkeypatch.setattr(legacy, "get_team_membership", member_read) + monkeypatch.setattr(auth_checks, "get_team_member_default_budget", default_read) + monkeypatch.setattr(legacy, "get_team_member_default_budget", default_read) + token: Final = UserAPIKeyAuth(user_id="u", team_id="t", team_member_model_max_budget={"old": {}}) + prior: Final = token.team_member_model_max_budget + team: Final = LiteLLM_TeamTable(team_id="t", metadata={"team_member_budget_id": "default"}) + fn: Final = ( + adapter.load_team_member_default_model_budget if native else legacy._load_team_member_default_model_budget # pyright: ignore[reportPrivateUsage] # Phase 1 leaf is the unchanged reference implementation + ) + pending: Final = fn(token, team, Mock(), UserApiKeyCache(), skip, None) + if skip: + await pending + else: + with pytest.raises(RuntimeError, match="budget read failed") as raised: + await pending + assert raised.value is error + assert token.team_member_model_max_budget is prior + + +@pytest.mark.asyncio +async def test_member_native_failure_never_becomes_skippable_lookup_error(monkeypatch: pytest.MonkeyPatch) -> None: + error: Final = NativeAuthRuntimeError("native planner failed") + read: Final = AsyncMock(side_effect=error) + monkeypatch.setattr(auth_checks, "get_team_membership", read) + with pytest.raises(NativeAuthRuntimeError, match="native planner failed") as raised: + await adapter.load_team_member_default_model_budget( + UserAPIKeyAuth(user_id="u", team_id="t"), None, Mock(), UserApiKeyCache(), True, None + ) + assert raised.value is error + read.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_member_planner_failure_stops_before_io(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.rust_bridge import _native + + native: Final = Mock(side_effect=RuntimeError("decode failed")) + read: Final = AsyncMock() + monkeypatch.setattr(_native, "request_auth_context_member", native) + monkeypatch.setattr(auth_checks, "get_team_membership", read) + with pytest.raises(NativeAuthRuntimeError, match="decode failed"): + await adapter.load_team_member_default_model_budget( + UserAPIKeyAuth(user_id="u", team_id="t"), None, Mock(), UserApiKeyCache(), True, None + ) + read.assert_not_awaited() + native.assert_called_once() diff --git a/tests/unit/proxy/auth/test_key_budget_context_adapter.py b/tests/unit/proxy/auth/test_key_budget_context_adapter.py new file mode 100644 index 00000000000..cb83565419e --- /dev/null +++ b/tests/unit/proxy/auth/test_key_budget_context_adapter.py @@ -0,0 +1,108 @@ +import math +from datetime import datetime, timezone +from typing import Final +from unittest.mock import Mock + +import pytest + +import litellm +from litellm._auth import budget_context as core +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth import budget_throttle +from litellm.proxy.auth import key_budget_context_adapter as adapter +from litellm.proxy.auth import user_api_key_auth as legacy + +_NOW: Final = datetime(2025, 1, 1, tzinfo=timezone.utc) + + +@pytest.mark.parametrize("native", (False, True), ids=("reference", "native")) +@pytest.mark.parametrize( + "increase", (None, False, True, 0, -3, 3, float("nan"), float("inf"), "", "3", [], [3], {}, {"n": 3}) +) +@pytest.mark.parametrize( + "expiry", ("2024-12-31", "2025-01-01T00:00:00+00:00", "2025-01-01T00:00:00.000001Z", "2026-01-01T09:00:00+09:00") +) +@pytest.mark.parametrize("limit", (None, 0, 10)) +def test_temporary_budget_reference_native( + native: bool, monkeypatch: pytest.MonkeyPatch, increase: object, expiry: str, limit: float | None +) -> None: + now: Final = Mock(return_value=_NOW) + clock: Final = Mock(wraps=datetime, now=now) + monkeypatch.setattr(adapter, "datetime", clock) + monkeypatch.setattr(legacy, "datetime", clock) + token: Final = UserAPIKeyAuth( + max_budget=limit, metadata={"temp_budget_increase": increase, "temp_budget_expiry": expiry} + ) + fn: Final = ( + adapter.update_key_budget_with_temp_budget_increase + if native + else legacy._update_key_budget_with_temp_budget_increase # pyright: ignore[reportPrivateUsage] # Phase 1 leaf is the unchanged reference implementation + ) + active: Final = expiry.startswith("2026") or expiry == "2025-01-01T00:00:00.000001Z" + changes: Final = limit is not None and active and bool(increase) + if changes and not isinstance(increase, (int, float)): + with pytest.raises(TypeError, match="unsupported operand type"): + fn(token) + return + result: Final = fn(token) + assert (result is token) is not changes + if changes: + assert limit is not None and isinstance(increase, (int, float)) + expected: Final = limit + increase + assert result.max_budget is not None + assert result.max_budget == expected or (math.isnan(result.max_budget) and math.isnan(expected)) + assert token.max_budget == limit + assert now.call_count == (0 if limit is None else 1) + + +@pytest.mark.parametrize("native", (False, True), ids=("reference", "native")) +@pytest.mark.parametrize("expiry", (None, 5, "", "invalid")) +@pytest.mark.parametrize("increase_present", (False, True)) +@pytest.mark.parametrize("limit", (None, 1)) +def test_temporary_metadata_error_timing( + native: bool, monkeypatch: pytest.MonkeyPatch, expiry: object, increase_present: bool, limit: float | None +) -> None: + token: Final = UserAPIKeyAuth( + max_budget=limit, + metadata={"temp_budget_expiry": expiry, **({"temp_budget_increase": 0} if increase_present else {})}, + ) + now: Final = Mock(return_value=_NOW) + clock: Final = Mock(wraps=datetime, now=now) + monkeypatch.setattr(adapter, "datetime", clock) + monkeypatch.setattr(legacy, "datetime", clock) + fn: Final = ( + adapter.update_key_budget_with_temp_budget_increase + if native + else legacy._update_key_budget_with_temp_budget_increase # pyright: ignore[reportPrivateUsage] # Phase 1 leaf is the unchanged reference implementation + ) + if limit is not None and increase_present: + with pytest.raises(ValueError if isinstance(expiry, str) else TypeError): + fn(token) + else: + assert fn(token) is token + now.assert_not_called() + + +@pytest.mark.parametrize("pct", (None, False, True, "", "0.5", 0, -0.5, 0.25, 1, 2, float("nan"), float("inf"))) +@pytest.mark.parametrize("opt_in", (None, False, True, 1, "true")) +@pytest.mark.parametrize("limit", (None, 0, 1)) +def test_throttle_reference_native( + monkeypatch: pytest.MonkeyPatch, pct: object, opt_in: object, limit: int | None +) -> None: + monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", pct) + token: Final = UserAPIKeyAuth(metadata={"throttle_on_budget_exceeded": opt_in}, tpm_limit=limit) + assert adapter.budget_throttle_percentage() == budget_throttle.budget_throttle_percentage() + assert adapter.should_throttle_budget_exceeded(token) == budget_throttle.should_throttle_budget_exceeded(token) + expected: Final = ( + budget_throttle.budget_throttle_percentage() if budget_throttle.should_throttle_budget_exceeded(token) else None + ) + assert adapter.inspect_key_budget_throttle(token) == expected + + +@pytest.mark.parametrize("tpm", (None, 0, 1)) +@pytest.mark.parametrize("rpm", (None, 0, 1)) +@pytest.mark.parametrize("opt_in", (False, True)) +def test_throttle_native_selects_config_read(tpm: int | None, rpm: int | None, opt_in: bool) -> None: + assert core.throttle(core.ThrottleFacts(opt_in, tpm, rpm)).read_percentage is ( + opt_in and (tpm is not None or rpm is not None) + ) diff --git a/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py b/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py index 55beaf5254d..e892fe43e65 100644 --- a/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py +++ b/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py @@ -131,26 +131,29 @@ def test_resolve_model_budget(): assert resolve_model_budget(model="non-existent", model_max_budget=model_max_budget) is None -# Test is_key_within_model_budget @pytest.mark.asyncio -async def test_is_key_within_model_budget(budget_limiter): - # Mock user API key dict - user_api_key = UserAPIKeyAuth( +@pytest.mark.parametrize("spend", (50.0, 100.0, 150.0)) +async def test_is_key_within_model_budget( + budget_limiter: PROXY_VirtualKeyModelMaxBudgetLimiter, + spend: float, +) -> None: + user_api_key: Final = UserAPIKeyAuth( token="test-key", key_alias="test-alias", model_max_budget={"gpt-4": {"budget_limit": 100.0, "time_period": "1d"}}, ) - - # Test when model is within budget - with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=50.0): - assert await budget_limiter.is_key_within_model_budget(user_api_key, "gpt-4") is True - - # Test when model exceeds budget - with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=150.0): + counter: Final = model_budget_spend_cache_key( + entity_type=Litellm_EntityType.KEY, + entity_id="test-key", + budget_model="gpt-4", + budget_duration="1d", + ) + await budget_limiter.dual_cache.async_set_cache(key=counter, value=spend) + if spend >= 100.0: with pytest.raises(litellm.BudgetExceededError): await budget_limiter.is_key_within_model_budget(user_api_key, "gpt-4") - - # Test model not in budget config + else: + assert await budget_limiter.is_key_within_model_budget(user_api_key, "gpt-4") is True assert await budget_limiter.is_key_within_model_budget(user_api_key, "non-existent") is True @@ -224,38 +227,26 @@ async def test_async_log_success_event_uses_per_model_budget_duration(budget_lim assert call_kwargs["response_cost"] == 0.05 -# Test is_end_user_within_model_budget @pytest.mark.asyncio -async def test_is_end_user_within_model_budget(budget_limiter): - # Test when model is within budget - with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=50.0): - assert ( - await budget_limiter.is_end_user_within_model_budget( - "test-user", - {"gpt-4": {"budget_limit": 100.0, "time_period": "1d"}}, - "gpt-4", - ) - is True - ) - - # Test when model exceeds budget - with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=150.0): - with pytest.raises(litellm.BudgetExceededError): - await budget_limiter.is_end_user_within_model_budget( - "test-user", - {"gpt-4": {"budget_limit": 100.0, "time_period": "1d"}}, - "gpt-4", - ) - - # Test model not in budget config - assert ( - await budget_limiter.is_end_user_within_model_budget( - "test-user", - {"gpt-4": {"budget_limit": 100.0, "time_period": "1d"}}, - "non-existent", - ) - is True +@pytest.mark.parametrize("spend", (50.0, 100.0, 150.0)) +async def test_is_end_user_within_model_budget( + budget_limiter: PROXY_VirtualKeyModelMaxBudgetLimiter, + spend: float, +) -> None: + counter: Final = model_budget_spend_cache_key( + entity_type=Litellm_EntityType.END_USER, + entity_id="test-user", + budget_model="gpt-4", + budget_duration="1d", ) + await budget_limiter.dual_cache.async_set_cache(key=counter, value=spend) + budgets: Final = {"gpt-4": {"budget_limit": 100.0, "time_period": "1d"}} + if spend >= 100.0: + with pytest.raises(litellm.BudgetExceededError): + await budget_limiter.is_end_user_within_model_budget("test-user", budgets, "gpt-4") + else: + assert await budget_limiter.is_end_user_within_model_budget("test-user", budgets, "gpt-4") is True + assert await budget_limiter.is_end_user_within_model_budget("test-user", budgets, "non-existent") is True # Test _get_spend_for_model_budget for the end-user scope