diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 1f4668b3a3f..dce8c07fb0b 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -39,13 +39,13 @@ mod _native { use crate::model_registry::bundled_model_cost_map; #[pymodule_export] use crate::request_auth::{ - NativeAuthContinuation, request_auth_budget_alert, 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_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, request_auth_select_api_key, - request_auth_select_websocket_api_key, request_auth_start, + 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_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, + request_auth_select_api_key, request_auth_select_websocket_api_key, request_auth_start, }; #[pymodule_export] use crate::router::{RouterCall, RouterCoreHandle}; diff --git a/litellm-rust/crates/python-bridge/src/request_auth.rs b/litellm-rust/crates/python-bridge/src/request_auth.rs index a13ea39eb42..ab4562e4adb 100644 --- a/litellm-rust/crates/python-bridge/src/request_auth.rs +++ b/litellm-rust/crates/python-bridge/src/request_auth.rs @@ -225,3 +225,12 @@ pub(crate) fn request_auth_budget_exemption( &py.detach(move || core::budget_pricing::exemption(&facts)), ) } + +#[pyfunction] +pub(crate) fn request_auth_budget_alert_plan( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_alerts::Facts = from_py(facts)?; + to_py(py, &py.detach(move || core::budget_alerts::plan(&facts))) +} diff --git a/litellm-rust/crates/request-auth-core/src/budget_alerts.rs b/litellm-rust/crates/request-auth-core/src/budget_alerts.rs new file mode 100644 index 00000000000..e6b6ad2a274 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/budget_alerts.rs @@ -0,0 +1,233 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)] +pub struct Failure { + pub kind: String, + pub message: String, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct ThresholdFacts { + pub key: String, + pub digit: bool, + pub digits: Option, + pub ratio: Option, + pub integer_failure: Option, + pub ratio_failure: Option, + pub text: Option, + pub items: Option>>, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct ConfigFacts { + pub global_entries: Vec, + pub local_entries: Vec, + pub global_failure: Option, + pub local_failure: Option, + pub order: Option>, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct Facts { + pub member: bool, + pub spend: Option, + pub limit: Option, + pub default_fraction: f64, + pub config: Option, +} + +#[derive(Clone, Debug, Serialize, PartialEq)] +pub struct Threshold { + pub key: String, + pub recipients: Vec, +} + +#[derive(Clone, Debug, Default, Serialize, PartialEq)] +pub struct Decision { + pub read_config: bool, + pub read_order: bool, + pub global_keys: Vec, + pub local_keys: Vec, + pub emit: bool, + pub thresholds: Option>, + pub failure: Option, +} + +fn member_entry(entry: &ThresholdFacts) -> Result { + if !entry.digit || entry.key.chars().count() > 3 { + return Ok(false); + } + if let Some(error) = &entry.integer_failure { + return Err(error.clone()); + } + Ok(entry + .digits + .as_deref() + .and_then(|digits| digits.parse::().ok()) + .is_some_and(|n| (1..=100).contains(&n))) +} + +fn eligible(entries: &[ThresholdFacts], member: bool) -> Result, Failure> { + entries.iter().try_fold(Vec::new(), |mut selected, entry| { + if !member || member_entry(entry)? { + selected.push(entry); + } + Ok(selected) + }) +} + +fn recipients(entry: &ThresholdFacts) -> Vec { + let strings: Vec<&str> = if let Some(text) = &entry.text { + text.split(',').collect() + } else { + entry + .items + .as_deref() + .unwrap_or_default() + .iter() + .filter_map(Option::as_deref) + .collect() + }; + strings + .into_iter() + .map(|s| s.trim_matches(|c: char| c.is_whitespace() || matches!(c, '\u{1c}'..='\u{1f}'))) + .filter(|s| !s.is_empty()) + .map(str::to_owned) + .collect() +} + +fn merge(key: &str, global: &[&ThresholdFacts], local: &[&ThresholdFacts]) -> Threshold { + let recipients = global + .iter() + .chain(local) + .filter(|entry| entry.key == key) + .flat_map(|entry| recipients(entry)) + .fold(Vec::new(), |mut merged, email| { + if !merged.contains(&email) { + merged.push(email); + } + merged + }); + Threshold { + key: key.to_owned(), + recipients, + } +} + +fn evaluate(facts: &Facts) -> Result { + if !facts.member && !(facts.limit.is_some() && facts.spend.is_some_and(|spend| spend > 0.0)) { + return Ok(Decision::default()); + } + let Some(config) = &facts.config else { + return Ok(Decision { + read_config: true, + ..Decision::default() + }); + }; + if let Some(failure) = &config.global_failure { + return Err(failure.clone()); + } + if !facts.member + && let Some(failure) = &config.local_failure + { + return Err(failure.clone()); + } + let global = eligible(&config.global_entries, false)?; + let local = if facts.member && config.local_failure.is_some() { + vec![] + } else { + eligible(&config.local_entries, facts.member)? + }; + if global.is_empty() && local.is_empty() { + let emit = !facts.member + && matches!((facts.spend, facts.limit), (Some(spend), Some(limit)) if spend >= limit * facts.default_fraction && spend < limit); + return Ok(Decision { + emit, + ..Decision::default() + }); + } + let Some(order) = &config.order else { + return Ok(Decision { + read_order: true, + global_keys: global.iter().map(|entry| entry.key.clone()).collect(), + local_keys: local.iter().map(|entry| entry.key.clone()).collect(), + ..Decision::default() + }); + }; + let thresholds: Vec = order + .iter() + .map(|key| merge(key, &global, &local)) + .collect(); + if facts.member && facts.spend.is_some_and(|spend| spend <= 0.0) { + return Ok(Decision::default()); + } + let mut minimum: Option<&ThresholdFacts> = None; + for key in order { + let Some(entry) = global.iter().chain(&local).find(|entry| entry.key == *key) else { + continue; + }; + if !entry.digit { + continue; + } + if let Some(failure) = &entry.integer_failure { + return Err(failure.clone()); + } + let Some(digits) = entry.digits.as_ref() else { + continue; + }; + if minimum + .and_then(|entry| entry.digits.as_ref()) + .is_none_or(|current| (digits.len(), digits) < (current.len(), current)) + { + minimum = Some(entry); + } + } + let Some(minimum) = minimum else { + return Ok(Decision::default()); + }; + if let Some(failure) = &minimum.ratio_failure { + return Err(failure.clone()); + } + let emit = matches!((facts.spend, facts.limit, minimum.ratio), (Some(spend), Some(limit), Some(ratio)) if spend.partial_cmp(&(limit * ratio)) != Some(std::cmp::Ordering::Less)); + Ok(Decision { + emit, + thresholds: Some(thresholds), + ..Decision::default() + }) +} + +pub fn plan(facts: &Facts) -> Decision { + evaluate(facts).unwrap_or_else(|failure| Decision { + failure: Some(failure), + ..Decision::default() + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::before(7.9, false)] + #[case::at(8.0, true)] + #[case::at_limit(10.0, false)] + #[case::nan(f64::NAN, false)] + #[case::zero(0.0, false)] + fn default_key_threshold(#[case] spend: f64, #[case] emit: bool) { + let facts = Facts { + member: false, + spend: Some(spend), + limit: Some(10.0), + default_fraction: 0.8, + config: Some(ConfigFacts { + global_entries: vec![], + local_entries: vec![], + global_failure: None, + local_failure: None, + order: None, + }), + }; + assert_eq!(plan(&facts).emit, emit); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/lib.rs b/litellm-rust/crates/request-auth-core/src/lib.rs index ea732a7146f..e74b05be20a 100644 --- a/litellm-rust/crates/request-auth-core/src/lib.rs +++ b/litellm-rust/crates/request-auth-core/src/lib.rs @@ -1,5 +1,6 @@ #![forbid(unsafe_code)] +pub mod budget_alerts; pub mod budget_pricing; pub mod common_budgets; pub mod credential; diff --git a/litellm/_auth/budget_alerts.py b/litellm/_auth/budget_alerts.py new file mode 100644 index 00000000000..d5546243790 --- /dev/null +++ b/litellm/_auth/budget_alerts.py @@ -0,0 +1,69 @@ +from dataclasses import dataclass +from typing import Final + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class Failure: + kind: str + message: str + + +@dataclass(frozen=True, slots=True) +class ThresholdFacts: + key: str + digit: bool + digits: str | None = None + ratio: float | None = None + integer_failure: Failure | None = None + ratio_failure: Failure | None = None + text: str | None = None + items: tuple[str | None, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class ConfigFacts: + global_entries: tuple[ThresholdFacts, ...] = () + local_entries: tuple[ThresholdFacts, ...] = () + global_failure: Failure | None = None + local_failure: Failure | None = None + order: tuple[str, ...] | None = None + + +@dataclass(frozen=True, slots=True) +class Facts: + member: bool + spend: float | None + limit: float | None + default_fraction: float + config: ConfigFacts | None = None + + +@dataclass(frozen=True, slots=True) +class Threshold: + key: str + recipients: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class Decision: + read_config: bool = False + read_order: bool = False + global_keys: tuple[str, ...] = () + local_keys: tuple[str, ...] = () + emit: bool = False + thresholds: tuple[Threshold, ...] | None = None + failure: Failure | None = None + + +_FACTS: Final = TypeAdapter(Facts) +_DECISION: Final = TypeAdapter(Decision) + + +@native_boundary +def plan(facts: Facts) -> Decision: + return _DECISION.validate_python(_native.request_auth_budget_alert_plan(_FACTS.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts diff --git a/litellm/proxy/auth/budget_alert_adapter.py b/litellm/proxy/auth/budget_alert_adapter.py new file mode 100644 index 00000000000..2537af0f8b8 --- /dev/null +++ b/litellm/proxy/auth/budget_alert_adapter.py @@ -0,0 +1,178 @@ +import asyncio +from collections.abc import Mapping +from dataclasses import replace +from typing import Final + +from pydantic import TypeAdapter, ValidationError + +import litellm +from litellm._auth import budget_alerts +from litellm._auth.budget_alerts import ConfigFacts, Decision, Facts, Failure, ThresholdFacts +from litellm._logging import verbose_proxy_logger +from litellm.constants import EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE +from litellm.proxy._types import CallInfo, Litellm_EntityType, LiteLLM_UserTable, UserAPIKeyAuth +from litellm.proxy.auth.common_budget_adapter import metadata +from litellm.proxy.utils import ProxyLogging + +_CONFIG: Final = TypeAdapter(Mapping[str, object]) +_OBJECT: Final = TypeAdapter(object) +_ITEMS: Final = TypeAdapter(tuple[object, ...]) + + +def _threshold(key: str, raw: object) -> ThresholdFacts: + text: Final = raw if isinstance(raw, str) else None + items: Final = ( + tuple(item if isinstance(item, str) else None for item in _ITEMS.validate_python(raw)) + if isinstance(raw, list) + else None + ) + initial: Final = ThresholdFacts(key=key, digit=key.isdigit(), text=text, items=items) + try: + integer: Final = int(key) + except ValueError as error: + return replace(initial, integer_failure=Failure("ValueError", str(error))) + parsed: Final = replace(initial, digits=str(integer)) + try: + return replace(parsed, ratio=integer / 100.0) + except OverflowError as error: + return replace(parsed, ratio_failure=Failure("OverflowError", str(error))) + + +def _config(raw: object) -> tuple[tuple[ThresholdFacts, ...], Failure | None]: + if not raw: + return (), None + try: + value: Final = _CONFIG.validate_python(raw) + except ValidationError: + return (), Failure("AttributeError", f"'{type(raw).__name__}' object has no attribute 'items'") + return tuple(_threshold(key, value) for key, value in value.items()), None + + +def read_alert_facts(facts: Facts, global_config: object, local_config: object) -> Facts: + decision: Final = budget_alerts.plan(facts) + if not decision.read_config: + return facts + global_entries, global_failure = _config(global_config) + local_entries, local_failure = _config(local_config) + loaded: Final = replace(facts, config=ConfigFacts(global_entries, local_entries, global_failure, local_failure)) + order: Final = budget_alerts.plan(loaded) + if not order.read_order: + return loaded + assert loaded.config is not None + return replace(loaded, config=replace(loaded.config, order=tuple(set(order.global_keys) | set(order.local_keys)))) + + +def read_key_budget_alert_facts(valid_token: UserAPIKeyAuth) -> Facts: + facts: Final = Facts( + False, valid_token.spend, valid_token.max_budget, EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE + ) + if not budget_alerts.plan(facts).read_config: + return facts + return read_alert_facts( + facts, + _OBJECT.validate_python(litellm.default_key_max_budget_alert_emails), # pyright: ignore[reportUnknownMemberType] # validate the legacy global config snapshot + metadata(valid_token.metadata).get("max_budget_alert_emails"), # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType] # legacy metadata is validated at the boundary + ) + + +def read_member_budget_alert_facts( + spend: float, + max_budget: float, + team_metadata: Mapping[str, object] | None, +) -> Facts: + facts: Final = Facts(True, spend, max_budget, EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE) + return read_alert_facts( + facts, None, team_metadata.get("team_member_max_budget_alert_emails") if team_metadata else None + ) + + +def raise_alert_failure(decision: Decision) -> None: + failure: Final = decision.failure + if failure is None: + return + match failure.kind: + case "AttributeError": + raise AttributeError(failure.message) + case "OverflowError": + raise OverflowError(failure.message) + case _: + raise ValueError(failure.message) + + +def schedule_budget_alert(decision: Decision, call_info: CallInfo, proxy_logging_obj: ProxyLogging) -> None: + raise_alert_failure(decision) + if decision.emit: + projected: Final = ( + call_info + if decision.thresholds is None + else call_info.model_copy(update={"max_budget_alert_emails": alert_emails(decision)}) + ) + asyncio.create_task(proxy_logging_obj.budget_alerts(type="max_budget_alert", user_info=projected)) + + +def alert_emails( + decision: Decision, +) -> dict[str, list[str]] | None: # mutable-ok: CallInfo requires the legacy dict of lists projection + return None if decision.thresholds is None else {entry.key: list(entry.recipients) for entry in decision.thresholds} + + +async def virtual_key_max_budget_alert_check( + valid_token: UserAPIKeyAuth, + proxy_logging_obj: ProxyLogging, + user_obj: LiteLLM_UserTable | None = None, +) -> None: + decision: Final = budget_alerts.plan(read_key_budget_alert_facts(valid_token)) + raise_alert_failure(decision) + if not decision.emit: + return + if decision.thresholds is None: + assert valid_token.max_budget is not None + verbose_proxy_logger.debug( + "Reached Max Budget Alert Threshold for token %s, spend %s, max_budget %s, alert_threshold %s", + valid_token.token, + valid_token.spend, + valid_token.max_budget, + valid_token.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE, + ) + call_info: Final = CallInfo( + token=valid_token.token, + spend=valid_token.spend, + max_budget=valid_token.max_budget, + soft_budget=valid_token.soft_budget, + user_id=valid_token.user_id, + team_id=valid_token.team_id, + team_alias=valid_token.team_alias, + organization_id=valid_token.org_id, + user_email=user_obj.user_email if user_obj else None, + key_alias=valid_token.key_alias, + event_group=Litellm_EntityType.KEY, + ) + schedule_budget_alert(decision, call_info, proxy_logging_obj) + + +def team_member_max_budget_alert_check( + team_id: str, + team_alias: str | None, + team_metadata: Mapping[str, object] | None, + organization_id: str | None, + user_id: str, + user_email: str | None, + proxy_logging_obj: ProxyLogging, + spend: float, + max_budget: float, +) -> None: + decision: Final = budget_alerts.plan(read_member_budget_alert_facts(spend, max_budget, team_metadata)) + raise_alert_failure(decision) + if not decision.emit: + return + call_info: Final = CallInfo( + spend=spend, + max_budget=max_budget, + user_id=user_id, + team_id=team_id, + team_alias=team_alias, + organization_id=organization_id, + user_email=user_email, + event_group=Litellm_EntityType.TEAM_MEMBER, + ) + schedule_budget_alert(decision, call_info, proxy_logging_obj) diff --git a/litellm/proxy/auth/budget_pricing_adapter.py b/litellm/proxy/auth/budget_pricing_adapter.py index 5baa51eaddd..841d174cae0 100644 --- a/litellm/proxy/auth/budget_pricing_adapter.py +++ b/litellm/proxy/auth/budget_pricing_adapter.py @@ -7,6 +7,7 @@ from pydantic import TypeAdapter import litellm from litellm._auth import budget_pricing from litellm._auth.budget_pricing import DeploymentPrice, ModelPriceFacts, RouteBudgetFacts +from litellm._auth.failure import NativeAuthFailure from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.decisions_pricing import decisions_pricing_fields, decisions_request_price from litellm.proxy.auth.route_checks import RouteChecks @@ -56,6 +57,8 @@ def read_model_price_facts(model: str, router: Router, cached: bool | None = Non input_price=group.input_cost_per_token if group is not None else None, output_price=group.output_cost_per_token if group is not None else None, ) + except NativeAuthFailure: + raise except Exception: # noqa: BLE001 # preserve lookup failure as non-exempt without caching return ModelPriceFacts(lookup_failed=True) if not budget_pricing.price(facts).read_deployments: @@ -64,6 +67,8 @@ def read_model_price_facts(model: str, router: Router, cached: bool | None = Non deployments: Final = tuple( _deployment_price(deployment) for deployment in _deployments_served_for(model, router) ) + except NativeAuthFailure: + raise except Exception: # noqa: BLE001 # preserve lookup failure as non-exempt without caching return ModelPriceFacts(lookup_failed=True) return replace(facts, deployments=deployments) diff --git a/litellm/proxy/auth/budget_windows_adapter.py b/litellm/proxy/auth/budget_windows_adapter.py index 250591621b2..1c60bc96179 100644 --- a/litellm/proxy/auth/budget_windows_adapter.py +++ b/litellm/proxy/auth/budget_windows_adapter.py @@ -6,6 +6,7 @@ from typing import Final from pydantic import TypeAdapter from litellm._auth.common_budgets import BudgetFacts, WindowFacts, window_start +from litellm._auth.failure import NativeAuthFailure from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy._types import Litellm_EntityType, LiteLLM_TeamTable, UserAPIKeyAuth from litellm.proxy.auth.common_budget_adapter import CounterRead, raise_budget, read_budget @@ -28,6 +29,8 @@ def _window(raw: object) -> BudgetWindow: def _seconds(duration: str) -> int | None: try: return duration_in_seconds(duration) + except NativeAuthFailure: + raise except Exception: # noqa: BLE001 # preserve invalid-duration skip behavior return None @@ -38,6 +41,8 @@ def _reset(value: object) -> datetime | None: if isinstance(value, str): try: parsed: Final = datetime.fromisoformat(value.replace("Z", "+00:00")) + except NativeAuthFailure: + raise except ValueError: return None return _reset(parsed) diff --git a/litellm/proxy/auth/common_budget_adapter.py b/litellm/proxy/auth/common_budget_adapter.py index 58adf767faa..db25792df30 100644 --- a/litellm/proxy/auth/common_budget_adapter.py +++ b/litellm/proxy/auth/common_budget_adapter.py @@ -212,7 +212,7 @@ async def user_max_budget_check( BudgetFacts( "User", user_object.max_budget, - team_key=team_object is not None, + team_key=team_object is not None and team_object.team_id is not None, # pyright: ignore[reportUnnecessaryComparison] # retain the legacy None guard for constructed cache rows apply_user_budget_to_team_keys=general_settings.get("apply_user_budget_to_team_keys") is True, ), CounterRead(f"spend:user:{user_object.user_id}", user_object.spend or 0.0, user_object.max_budget), diff --git a/litellm/proxy/auth/common_budget_lookups.py b/litellm/proxy/auth/common_budget_lookups.py index 8a58e43a163..7c9e473f5b4 100644 --- a/litellm/proxy/auth/common_budget_lookups.py +++ b/litellm/proxy/auth/common_budget_lookups.py @@ -4,6 +4,7 @@ from typing import Final from litellm._auth import common_budgets from litellm._auth.common_budgets import BudgetFacts +from litellm._auth.failure import NativeAuthFailure from litellm.proxy._types import ( Litellm_EntityType, LiteLLM_TagTable, @@ -53,6 +54,8 @@ async def organization_max_budget_check( proxy_logging_obj=proxy_logging_obj, include_budget_table=True, ) + except NativeAuthFailure: + raise except Exception: # noqa: BLE001 # preserve organization lookup failure handling return if org is None: @@ -165,10 +168,10 @@ async def check_team_member_budget( get_team_member_default_budget, get_team_membership, get_utc_datetime, - team_member_max_budget_alert_check, ) + from litellm.proxy.auth.budget_alert_adapter import team_member_max_budget_alert_check - if team_object is None or valid_token is None or valid_token.user_id is None: + if team_object is None or team_object.team_id is None or valid_token is None or valid_token.user_id is None: # pyright: ignore[reportUnnecessaryComparison] # retain the legacy None guard for constructed cache rows return membership: Final = ( team_membership diff --git a/litellm/proxy/auth/model_budget_adapter.py b/litellm/proxy/auth/model_budget_adapter.py index 10aae9547c5..50d804bb960 100644 --- a/litellm/proxy/auth/model_budget_adapter.py +++ b/litellm/proxy/auth/model_budget_adapter.py @@ -7,6 +7,7 @@ from pydantic import TypeAdapter import litellm from litellm._auth import model_budget_data +from litellm._auth.failure import NativeAuthFailure from litellm._auth.model_budget_data import ( ModelBudgetEntity, ModelBudgetEntry, @@ -40,6 +41,8 @@ def _entry(name: str, raw: object) -> ModelBudgetEntry: try: config: Final = BudgetConfig.model_validate(raw) seconds: Final = duration_in_seconds(config.budget_duration) if config.budget_duration is not None else None + except NativeAuthFailure: + raise except Exception: # noqa: BLE001 # preserve malformed entry fallthrough return ModelBudgetEntry(name, None, None, None) return ModelBudgetEntry(name, config.max_budget, config.budget_duration, seconds) @@ -48,6 +51,8 @@ def _entry(name: str, raw: object) -> ModelBudgetEntry: def _spend(raw: object) -> float: try: return _FLOAT.validate_python(raw or 0.0) + except NativeAuthFailure: + raise except ValueError: return 0.0 @@ -148,7 +153,9 @@ async def _read_snapshot( return await _read_snapshot(io, replace(facts, snapshot=legacy_snapshot), key, legacy_key, duration) case "Shared": now: Final = io.clock() - shared: Final = await io.shared_windows.read(key, duration) if duration is not None else None + shared: Final = ( + await io.shared_windows.read(spend_key=key, budget_duration=duration) if duration is not None else None + ) shared_snapshot: Final = replace( facts.snapshot, shared_spend=shared.spend if shared is not None else None, diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 1a63c9f06c1..2c889910359 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -810,3 +810,5 @@ class RouterCall(Generic[_S]): def next(self, outcome: object) -> _S: ... def request_auth_budget_exemption(facts: object) -> dict[str, object]: ... + +def request_auth_budget_alert_plan(facts: object) -> dict[str, object]: ... diff --git a/tests/unit/proxy/auth/budget_alert_reference.py b/tests/unit/proxy/auth/budget_alert_reference.py new file mode 100644 index 00000000000..a47fd364db6 --- /dev/null +++ b/tests/unit/proxy/auth/budget_alert_reference.py @@ -0,0 +1,42 @@ +from collections.abc import Mapping + +from pydantic import TypeAdapter, ValidationError + +from litellm._auth.budget_alerts import Facts + + +def reference_recipients(raw: object) -> list[str]: + if isinstance(raw, list): + return [email.strip() for email in raw if isinstance(email, str) and email.strip()] + if isinstance(raw, str): + return [email.strip() for email in raw.split(",") if email.strip()] + return [] + + +def reference_alert(facts: Facts, global_config: object, local_config: object) -> tuple[bool, object]: + if not facts.member and not (facts.limit is not None and facts.spend is not None and facts.spend > 0): + return False, None + if facts.member: + try: + config = TypeAdapter(Mapping[str, object]).validate_python(local_config) + except ValidationError: + config = {} + local_config = { + key: value for key, value in config.items() if key.isdigit() and len(key) <= 3 and 1 <= int(key) <= 100 + } + global_entries = {key: reference_recipients(value) for key, value in (global_config or {}).items()} + local_entries = {key: reference_recipients(value) for key, value in (local_config or {}).items()} + merged = { + key: list(dict.fromkeys(global_entries.get(key, []) + local_entries.get(key, []))) + for key in set(global_entries) | set(local_entries) + } + if merged: + if facts.member and facts.spend <= 0: + return False, None + minimum = min((int(key) for key in merged if key.isdigit()), default=None) + if minimum is None or facts.spend < facts.limit * (minimum / 100.0): + return False, None + return True, tuple((key, tuple(value)) for key, value in merged.items()) + return ( + not facts.member and facts.spend >= facts.limit * facts.default_fraction and facts.spend < facts.limit + ), None diff --git a/tests/unit/proxy/auth/test_budget_alert_adapter.py b/tests/unit/proxy/auth/test_budget_alert_adapter.py new file mode 100644 index 00000000000..8582074fc6a --- /dev/null +++ b/tests/unit/proxy/auth/test_budget_alert_adapter.py @@ -0,0 +1,125 @@ +import asyncio +from typing import Final +from unittest.mock import AsyncMock, Mock + +import pytest + +import litellm +from litellm._auth import budget_alerts +from litellm._auth.budget_alerts import Facts +from litellm._auth.failure import NativeAuthFailure, NativeAuthValueError +from litellm.constants import EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE +from litellm.proxy._types import CallInfo, Litellm_EntityType, UserAPIKeyAuth +from litellm.proxy.auth import budget_alert_adapter as native +from litellm.rust_bridge import _native +from tests.unit.proxy.auth.budget_alert_reference import reference_alert + + +def observed_reference(facts: Facts, global_config: object, local_config: object) -> object: + try: + return reference_alert(facts, global_config, local_config) + except (ValueError, OverflowError, AttributeError) as error: + return type(error).__name__, str(error) + + +def observed_native(facts: Facts, global_config: object, local_config: object) -> object: + try: + decision: Final = budget_alerts.plan(native.read_alert_facts(facts, global_config, local_config)) + native.raise_alert_failure(decision) + config: Final = native.alert_emails(decision) if decision.emit else None + return decision.emit, None if config is None else tuple((key, tuple(value)) for key, value in config.items()) + except (ValueError, OverflowError, AttributeError) as error: + return type(error).__name__, str(error) + + +@pytest.mark.parametrize("member", (False, True)) +@pytest.mark.parametrize("spend", (-1.0, 0.0, 7.9, 8.0, 9.0, 10.0, float("nan"), float("inf"))) +@pytest.mark.parametrize("limit", (-1.0, 0.0, 10.0, float("nan"), float("inf"))) +@pytest.mark.parametrize( + "configs", + ( + (None, None), + ({"80": " a@x.test,b@x.test "}, {"80": [" b@x.test ", "c@x.test", "", None]}), + (None, {"bad": "a@x.test"}), + (None, {"0": [], "101": None, "080": [" a@x.test "]}), + (None, {"80": ("a@x.test",)}), + (None, {"٨٠": "a@x.test"}), + (None, {"²": []}), + (None, {"9" * 400: [], "80": "a@x.test"}), + (None, {"9" * 400: []}), + (None, "invalid"), + ), +) +def test_native_threshold_merge_and_trigger_match_reference( + member: bool, spend: float, limit: float, configs: tuple[object, object] +) -> None: + facts: Final = Facts(member, spend, limit, EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE) + global_config, local_config = configs + global_value: Final = None if member else global_config + assert observed_native(facts, global_value, local_config) == observed_reference(facts, global_value, local_config) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("config", (None, {"80": " a@x.test, b@x.test "})) +async def test_key_alert_schedules_once_without_awaiting_and_preserves_projection( + monkeypatch: pytest.MonkeyPatch, config: object +) -> None: + monkeypatch.setattr(litellm, "default_key_max_budget_alert_emails", None) + key: Final = UserAPIKeyAuth( + token="token", + spend=8, + max_budget=10, + metadata={"max_budget_alert_emails": config}, + user_id="owner", + team_id="team", + ) + logger: Final = Mock(budget_alerts=AsyncMock()) + await native.virtual_key_max_budget_alert_check(key, logger) + assert logger.budget_alerts.call_count == 1 + assert logger.budget_alerts.await_count == 0 + call_info: Final = logger.budget_alerts.call_args.kwargs["user_info"] + expected: Final = CallInfo( + token="token", + spend=8, + max_budget=10, + soft_budget=None, + user_id="owner", + team_id="team", + team_alias=None, + organization_id=None, + user_email=None, + key_alias=None, + event_group=Litellm_EntityType.KEY, + ) + projected: Final = ( + expected + if config is None + else expected.model_copy(update={"max_budget_alert_emails": {"80": ["a@x.test", "b@x.test"]}}) + ) + assert call_info.model_dump() == projected.model_dump() + assert call_info.model_fields_set == projected.model_fields_set + await asyncio.sleep(0) + assert logger.budget_alerts.await_count == 1 + + +def test_native_failure_is_terminal_before_alert_config_reads(monkeypatch: pytest.MonkeyPatch) -> None: + failure: Final = ValueError("native alert decoder failed") + backend: Final = Mock(side_effect=failure) + monkeypatch.setattr(_native, "request_auth_budget_alert_plan", backend) + with pytest.raises(NativeAuthFailure, match="native alert decoder failed") as error: + native.read_alert_facts(Facts(False, 8.0, 10.0, 0.8), None, {"80": []}) + assert isinstance(error.value, ValueError) + assert backend.call_count == 1 + + +def test_pricing_native_failure_never_becomes_lookup_failure() -> None: + from litellm.proxy.auth.budget_pricing_adapter import read_model_price_facts + + failure: Final = NativeAuthValueError("native registry failed") + registry: Final = Mock(group=Mock(side_effect=failure)) + router: Final = Mock(registry=Mock(return_value=registry)) + with pytest.raises(NativeAuthFailure) as error: + read_model_price_facts("model", router) + assert error.value is failure + registry.group.assert_called_once() + registry.deployments.assert_not_called() diff --git a/tests/unit/proxy/auth/test_common_budget_lookups.py b/tests/unit/proxy/auth/test_common_budget_lookups.py new file mode 100644 index 00000000000..b034dafcf47 --- /dev/null +++ b/tests/unit/proxy/auth/test_common_budget_lookups.py @@ -0,0 +1,44 @@ +from typing import Final +from unittest.mock import AsyncMock + +import pytest + +from litellm._auth import common_budgets +from litellm.proxy import proxy_server +from litellm.proxy.auth.common_budget_lookups import _read_access_group_rows +from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("limit", "reads", "denied"), ((None, 0, False), (0, 0, False), (-1, 0, False), (1, 1, True), (2, 2, False)) +) +async def test_raw_access_group_snapshot_stops_reads_at_native_denial( + monkeypatch: pytest.MonkeyPatch, + limit: float | None, + reads: int, + denied: bool, +) -> None: + counter: Final = AsyncMock(return_value=1.0) + monkeypatch.setattr(proxy_server, "get_current_spend", counter) + rows: Final = { + "first": ModelAccessGroupBudget(access_group_name="first", max_budget=limit), + "second": ModelAccessGroupBudget(access_group_name="second", max_budget=limit), + } + snapshots: Final = await _read_access_group_rows(("first", "second"), rows) + assert common_budgets.check(snapshots[0].facts).denied is denied + assert len(snapshots) == (1 if denied else 2) + assert snapshots[0].name == "first" + assert counter.await_count == reads + + +@pytest.mark.asyncio +async def test_access_group_counter_failure_is_terminal_before_next_group(monkeypatch: pytest.MonkeyPatch) -> None: + failure: Final = RuntimeError("counter unavailable") + counter: Final = AsyncMock(side_effect=failure) + monkeypatch.setattr(proxy_server, "get_current_spend", counter) + row: Final = ModelAccessGroupBudget(access_group_name="first", max_budget=1) + with pytest.raises(RuntimeError, match="counter unavailable") as raised: + await _read_access_group_rows(("first", "second"), {"first": row, "second": row}) + assert raised.value is failure + assert counter.await_count == 1