diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 9bc7b42a37b..6c4c6d90c88 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -81,10 +81,12 @@ mod _native { request_auth_litellm_header_token, request_auth_llm_route, request_auth_managed, request_auth_managed_model_grants, request_auth_managed_models, request_auth_member_limit, request_auth_member_models, request_auth_model_access, request_auth_model_attribution, - request_auth_model_budget_check, request_auth_model_budget_query, - request_auth_model_budget_resolve, request_auth_model_groups, request_auth_model_helpers, - request_auth_model_projection, request_auth_org_admin, request_auth_org_admin_raw, - request_auth_org_route, request_auth_owner_teams, request_auth_personal_budget, + request_auth_model_budget_aliases, request_auth_model_budget_check, + request_auth_model_budget_query, request_auth_model_budget_resolve, + request_auth_model_budget_scope, request_auth_model_budget_team_exclusions, + request_auth_model_groups, request_auth_model_helpers, request_auth_model_projection, + request_auth_org_admin, request_auth_org_admin_raw, request_auth_org_route, + request_auth_owner_teams, request_auth_personal_budget, request_auth_personal_model_listing, request_auth_personal_model_request, request_auth_personal_models, request_auth_personal_owner, request_auth_reservation, request_auth_resolved_models, request_auth_route_access, request_auth_route_enable, @@ -205,6 +207,9 @@ mod tests { "request_auth_model_access", "request_auth_model_attribution", "request_auth_model_budget_check", + "request_auth_model_budget_aliases", + "request_auth_model_budget_scope", + "request_auth_model_budget_team_exclusions", "request_auth_model_budget_query", "request_auth_model_budget_resolve", "request_auth_model_groups", diff --git a/litellm-rust/crates/python-bridge/src/request_auth.rs b/litellm-rust/crates/python-bridge/src/request_auth.rs index 9d07c977f29..b9ea5225262 100644 --- a/litellm-rust/crates/python-bridge/src/request_auth.rs +++ b/litellm-rust/crates/python-bridge/src/request_auth.rs @@ -200,6 +200,42 @@ pub(crate) fn request_auth_model_budget_resolve( ) } +#[pyfunction] +pub(crate) fn request_auth_model_budget_aliases( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::model_budget_aliases::Facts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::model_budget_aliases::decide(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_model_budget_scope( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::model_budget_aliases::ScopeFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::model_budget_aliases::scope(&facts)), + ) +} + +#[pyfunction] +pub(crate) fn request_auth_model_budget_team_exclusions( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: Vec = from_py(facts)?; + to_py( + py, + &py.detach(move || core::model_budget_data::team_exclusions(&facts)), + ) +} + #[pyfunction] pub(crate) fn request_auth_model_budget_query( py: Python<'_>, diff --git a/litellm-rust/crates/request-auth-core/src/lib.rs b/litellm-rust/crates/request-auth-core/src/lib.rs index 592e7268c52..e0959931df4 100644 --- a/litellm-rust/crates/request-auth-core/src/lib.rs +++ b/litellm-rust/crates/request-auth-core/src/lib.rs @@ -24,6 +24,7 @@ pub mod managed_models; pub mod member_models; pub mod model_access; pub mod model_attribution; +pub mod model_budget_aliases; pub mod model_budget_data; mod model_budgets; pub mod model_groups; diff --git a/litellm-rust/crates/request-auth-core/src/model_budget_aliases.rs b/litellm-rust/crates/request-auth-core/src/model_budget_aliases.rs new file mode 100644 index 00000000000..dd1884f62b8 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/model_budget_aliases.rs @@ -0,0 +1,250 @@ +use serde::{Deserialize, Serialize}; + +use crate::model_budget_data::{Snapshot, check}; + +#[derive(Clone, Debug, Deserialize)] +pub struct Row { + pub key: String, + pub names: Vec, + pub prefix_known: Option, + pub provider_read: bool, + pub provider: Option, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct Facts { + pub model: String, + pub rows: Vec, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, strum::VariantArray)] +pub enum Read { + Entry, + Prefix, + Provider, +} + +#[derive(Clone, Debug, Serialize)] +pub struct Decision { + pub read: Option, + pub key: String, + pub models: Vec, +} + +fn reading(read: Read, key: &str) -> Decision { + Decision { + read: Some(read), + key: key.to_owned(), + models: vec![], + } +} + +fn row<'a>(facts: &'a Facts, key: &str) -> Result<&'a Row, Decision> { + facts + .rows + .iter() + .find(|row| row.key == key) + .ok_or_else(|| reading(Read::Entry, key)) +} + +fn first(row: &Row) -> Option<&str> { + row.names + .first() + .filter(|name| **name != row.key) + .map(String::as_str) +} + +fn provider(row: &Row) -> Result, Decision> { + if let Some((prefix, _)) = row.key.split_once('/') { + match row.prefix_known { + Some(true) => return Ok(Some(prefix)), + Some(false) => (), + None => return Err(reading(Read::Prefix, &row.key)), + } + } + if !row.provider_read { + return Err(reading(Read::Provider, &row.key)); + } + Ok(row.provider.as_deref()) +} + +fn target(facts: &Facts, key: &str, requested_provider: &str) -> Result, Decision> { + let row = row(facts, key)?; + let Some(name) = first(row) else { + return Ok(None); + }; + if provider(row)? != Some(requested_provider) { + return Ok(None); + } + Ok(Some( + name.strip_prefix(&format!("{requested_provider}/")) + .unwrap_or(name) + .to_owned(), + )) +} + +fn models(facts: &Facts) -> Result, Decision> { + let Some((provider, name)) = facts.model.split_once('/') else { + return Ok(match first(row(facts, &facts.model)?) { + Some(target) => vec![facts.model.clone(), target.to_owned()], + None => vec![facts.model.clone()], + }); + }; + let direct = target(facts, &facts.model, provider)?.filter(|name| !name.is_empty()); + let resolved = match direct { + Some(target) => Some(target), + None => target(facts, name, provider)?.filter(|name| !name.is_empty()), + }; + Ok(match resolved { + Some(target) if target != name => vec![facts.model.clone(), format!("{provider}/{target}")], + _ => vec![facts.model.clone()], + }) +} + +pub fn decide(facts: &Facts) -> Decision { + match models(facts) { + Ok(models) => Decision { + read: None, + key: facts.model.clone(), + models, + }, + Err(read) => read, + } +} + +#[derive(Clone, Debug, Deserialize)] +pub struct ScopeFacts { + pub models: Vec, + pub snapshots: Vec, + pub budget_present: bool, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize)] +pub struct ScopeDecision { + pub read_index: Option, + pub result_index: Option, +} + +pub fn scope(facts: &ScopeFacts) -> ScopeDecision { + if let Some(index) = facts.snapshots.iter().position(|value| check(value).denied) { + return ScopeDecision { + read_index: None, + result_index: Some(index), + }; + } + if !facts.budget_present { + return ScopeDecision { + read_index: None, + result_index: None, + }; + } + if facts.snapshots.len() < facts.models.len() { + return ScopeDecision { + read_index: Some(facts.snapshots.len()), + result_index: None, + }; + } + ScopeDecision { + read_index: None, + result_index: facts.snapshots.len().checked_sub(1), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::{fixture, rstest}; + + #[fixture] + fn alias() -> Row { + Row { + key: "provider/latest".into(), + names: vec!["provider/target".into()], + prefix_known: Some(true), + provider_read: false, + provider: None, + } + } + + #[rstest] + #[case::catalog_first(vec![], Some(Read::Entry), "provider/latest", vec![])] + #[case::prefix_before_metadata(vec![Row { prefix_known: None, ..alias() }], Some(Read::Prefix), "provider/latest", vec![])] + #[case::metadata_for_unknown_provider(vec![Row { prefix_known: Some(false), ..alias() }], Some(Read::Provider), "provider/latest", vec![])] + #[case::known_prefix_ignores_metadata(vec![alias()], None, "provider/latest", vec!["provider/latest", "provider/target"])] + #[case::absent_qualified_reads_bare(vec![Row { names: vec![], ..alias() }], Some(Read::Entry), "latest", vec![])] + #[case::matching_metadata(vec![Row { prefix_known: Some(false), provider_read: true, provider: Some("provider".into()), ..alias() }], None, "provider/latest", vec!["provider/latest", "provider/target"])] + #[case::wrong_provider_reads_bare(vec![Row { prefix_known: Some(false), provider_read: true, provider: Some("other".into()), ..alias() }], Some(Read::Entry), "latest", vec![])] + fn lazy_alias_reads( + #[case] rows: Vec, + #[case] read: Option, + #[case] key: &str, + #[case] models: Vec<&str>, + ) { + let result = decide(&Facts { + model: "provider/latest".into(), + rows, + }); + assert_eq!( + (result.read, result.key.as_str(), result.models), + (read, key, models.into_iter().map(str::to_owned).collect()) + ); + } + + #[rstest] + #[case::bare_alias("latest", vec!["target"], vec!["latest", "target"])] + #[case::self("latest", vec!["latest", "other"], vec!["latest"])] + #[case::empty_row("latest", vec![], vec!["latest"])] + #[case::empty_bare_target("latest", vec![""], vec!["latest", ""])] + fn bare_alias(#[case] model: &str, #[case] names: Vec<&str>, #[case] expected: Vec<&str>) { + let result = decide(&Facts { + model: model.into(), + rows: vec![Row { + key: model.into(), + names: names.into_iter().map(str::to_owned).collect(), + ..alias() + }], + }); + assert_eq!(result.read, None); + assert_eq!(result.models, expected); + } + + fn snapshot(limit: f64, spend: f64) -> Snapshot { + Snapshot { + max_budget: Some(limit), + local_spend: Some(spend), + legacy_spend: None, + shared_spend: None, + shared_started_at: None, + shared_window_seconds: None, + now: 0.0, + } + } + + #[rstest] + #[case::first(vec![], true, Some(0), None)] + #[case::target_after_allowed(vec![snapshot(1.0, 0.0)], true, Some(1), None)] + #[case::alias_denial(vec![snapshot(1.0, 1.0)], true, None, Some(0))] + #[case::target_denial(vec![snapshot(1.0, 0.0), snapshot(2.0, 3.0)], true, None, Some(1))] + #[case::all_allowed(vec![snapshot(1.0, 0.0), snapshot(2.0, 0.0)], true, None, Some(1))] + #[case::no_budgets(vec![], false, None, None)] + #[case::retain_denial(vec![snapshot(1.0, 1.0)], false, None, Some(0))] + #[case::nan_continues(vec![snapshot(f64::NAN, 1.0)], true, Some(1), None)] + fn scope_order( + #[case] snapshots: Vec, + #[case] budget_present: bool, + #[case] read_index: Option, + #[case] result_index: Option, + ) { + assert_eq!( + scope(&ScopeFacts { + models: vec!["alias".into(), "target".into()], + snapshots, + budget_present + }), + ScopeDecision { + read_index, + result_index + } + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/model_budget_data.rs b/litellm-rust/crates/request-auth-core/src/model_budget_data.rs index 51ebd2fc53c..17e7fe3b0b6 100644 --- a/litellm-rust/crates/request-auth-core/src/model_budget_data.rs +++ b/litellm-rust/crates/request-auth-core/src/model_budget_data.rs @@ -25,6 +25,8 @@ pub struct ResolveFacts { pub key_entries: Vec, pub team: bool, pub access_groups: Option>, + #[serde(default)] + pub excluded_entries: Vec, } #[derive(Clone, Debug, Serialize)] @@ -60,6 +62,18 @@ fn exact<'a>(entries: &'a [Entry], names: &[&str]) -> Option<&'a Entry> { } pub fn resolve(facts: &ResolveFacts) -> Resolution { + if !facts.excluded_entries.is_empty() { + return resolve(&ResolveFacts { + entries: facts + .entries + .iter() + .filter(|entry| !facts.excluded_entries.contains(&entry.name)) + .cloned() + .collect(), + excluded_entries: vec![], + ..facts.clone() + }); + } let names = candidates(facts); let key_replaces = facts.team && exact(&facts.key_entries, &names) @@ -125,6 +139,24 @@ pub fn resolve(facts: &ResolveFacts) -> Resolution { } } +pub fn team_exclusions(facts: &[ResolveFacts]) -> Vec { + facts + .iter() + .flat_map(|facts| { + let names = candidates(facts); + if exact(&facts.key_entries, &names) + .is_some_and(|entry| entry.max_budget.is_some_and(|limit| limit >= 0.0)) + { + names.into_iter().map(str::to_owned).collect::>() + } else { + vec![] + } + }) + .collect::>() + .into_iter() + .collect() +} + fn glob(pattern: &str, name: &str) -> bool { let parts: Vec<&str> = pattern.split('*').collect(); if parts.len() == 1 { @@ -260,4 +292,40 @@ mod tests { fn wildcard(#[case] pattern: &str, #[case] name: &str, #[case] matches: bool) { assert_eq!(glob(pattern, name), matches); } + + #[rstest] + #[case::alias_zero("alias", Some(0.0), vec!["alias"])] + #[case::target_positive("target", Some(1.0), vec!["target"])] + #[case::target_infinite("target", Some(f64::INFINITY), vec!["target"])] + #[case::negative("target", Some(-1.0), vec![])] + #[case::nan("target", Some(f64::NAN), vec![])] + #[case::missing("alias", None, vec![])] + fn team_alias_overrides( + #[case] key: &str, + #[case] limit: Option, + #[case] expected: Vec<&str>, + ) { + let entry = Entry { + name: key.into(), + max_budget: limit, + duration: Some("1h".into()), + duration_seconds: Some(3600.0), + }; + let alias = ResolveFacts { + model: "alias".into(), + bedrock_base: "alias".into(), + bedrock_provider: String::new(), + entries: vec![], + key_entries: vec![entry], + team: true, + access_groups: None, + excluded_entries: vec![], + }; + let target = ResolveFacts { + model: "target".into(), + bedrock_base: "target".into(), + ..alias.clone() + }; + assert_eq!(team_exclusions(&[alias, target]), expected); + } } diff --git a/litellm/_auth/model_budget_aliases.py b/litellm/_auth/model_budget_aliases.py new file mode 100644 index 00000000000..b9fedf9f091 --- /dev/null +++ b/litellm/_auth/model_budget_aliases.py @@ -0,0 +1,59 @@ +from dataclasses import dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.model_budget_data import ModelBudgetSnapshot +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class Row: + key: str + names: tuple[str, ...] + prefix_known: bool | None = None + provider_read: bool = False + provider: str | None = None + + +@dataclass(frozen=True, slots=True) +class Facts: + model: str + rows: tuple[Row, ...] = () + + +@dataclass(frozen=True, slots=True) +class Decision: + read: Literal["Entry", "Prefix", "Provider"] | None + key: str + models: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class ScopeFacts: + models: tuple[str, ...] + snapshots: tuple[ModelBudgetSnapshot, ...] = () + budget_present: bool = True + + +@dataclass(frozen=True, slots=True) +class ScopeDecision: + read_index: int | None + result_index: int | None + + +_FACTS: Final = TypeAdapter(Facts) +_DECISION: Final = TypeAdapter(Decision) +_SCOPE: Final = TypeAdapter(ScopeFacts) +_SCOPE_DECISION: Final = TypeAdapter(ScopeDecision) + + +@native_boundary +def decide(facts: Facts) -> Decision: + return _DECISION.validate_python(_native.request_auth_model_budget_aliases(_FACTS.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes typed native facts + + +@native_boundary +def scope(facts: ScopeFacts) -> ScopeDecision: + return _SCOPE_DECISION.validate_python(_native.request_auth_model_budget_scope(_SCOPE.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes typed native facts diff --git a/litellm/_auth/model_budget_data.py b/litellm/_auth/model_budget_data.py index b0b6f95c770..59e39ce8da6 100644 --- a/litellm/_auth/model_budget_data.py +++ b/litellm/_auth/model_budget_data.py @@ -26,6 +26,7 @@ class ModelBudgetResolveFacts: key_entries: tuple[ModelBudgetEntry, ...] = () team: bool = False access_groups: tuple[str, ...] | None = None + excluded_entries: tuple[str, ...] = () @dataclass(frozen=True, slots=True) @@ -69,6 +70,13 @@ _RESOLUTION: Final = TypeAdapter(ModelBudgetResolution) _QUERY: Final = TypeAdapter(ModelBudgetQueryFacts) _SNAPSHOT: Final = TypeAdapter(ModelBudgetSnapshot) _DECISION: Final = TypeAdapter(ModelBudgetDecision) +_TEAM_FACTS: Final = TypeAdapter(tuple[ModelBudgetResolveFacts, ...]) +_NAMES: Final = TypeAdapter(tuple[str, ...]) + + +@native_boundary +def team_exclusions(facts: tuple[ModelBudgetResolveFacts, ...]) -> tuple[str, ...]: + return _NAMES.validate_python(_native.request_auth_model_budget_team_exclusions(_TEAM_FACTS.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes typed native facts @native_boundary diff --git a/litellm/proxy/auth/common_budget_adapter.py b/litellm/proxy/auth/common_budget_adapter.py index e1894685d02..cc4d90b26f3 100644 --- a/litellm/proxy/auth/common_budget_adapter.py +++ b/litellm/proxy/auth/common_budget_adapter.py @@ -22,6 +22,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.utils import ProxyLogging +from litellm.redis.stores.spend_counters import end_user_counter, key_counter, team_counter, user_counter _METADATA: Final = TypeAdapter(Mapping[str, object]) AlertKind: TypeAlias = Literal[ @@ -136,7 +137,7 @@ async def team_max_budget_check( return facts, result = await read_budget( BudgetFacts("Team", team_object.max_budget, token_present=valid_token is not None), - CounterRead(f"spend:team:{team_object.team_id}", team_object.spend or 0.0, team_object.max_budget), + CounterRead(team_counter(team_object.team_id), team_object.spend or 0.0, team_object.max_budget), ) if result.alert is not None and valid_token is not None: emit_alert( @@ -175,7 +176,7 @@ async def virtual_key_max_budget_check( tpm_limit=valid_token.tpm_limit, rpm_limit=valid_token.rpm_limit, ), - CounterRead(f"spend:key:{valid_token.token}", valid_token.spend or 0.0, valid_token.max_budget), + CounterRead(key_counter(str(valid_token.token)), valid_token.spend or 0.0, valid_token.max_budget), ) if result.alert is not None and facts.spend is not None: send_token_budget_alert(valid_token, facts.spend, proxy_logging_obj, user_obj) @@ -231,7 +232,7 @@ async def user_max_budget_check( 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), + CounterRead(user_counter(user_object.user_id), user_object.spend or 0.0, user_object.max_budget), ) if result.alert is not None: emit_alert( @@ -262,7 +263,7 @@ async def check_end_user_budget(end_user_obj: LiteLLM_EndUserTable, route: str) ) facts, result = await read_budget( BudgetFacts("EndUser", limit, info_route=RouteChecks.is_info_route(route)), - CounterRead(f"spend:end_user:{end_user_obj.user_id}", end_user_obj.spend or 0.0, limit, authoritative=True), + CounterRead(end_user_counter(end_user_obj.user_id), end_user_obj.spend or 0.0, limit, authoritative=True), ) if result.denied: raise_budget( diff --git a/litellm/proxy/auth/common_budget_lookups.py b/litellm/proxy/auth/common_budget_lookups.py index 071892cfa17..e7b24ab8d70 100644 --- a/litellm/proxy/auth/common_budget_lookups.py +++ b/litellm/proxy/auth/common_budget_lookups.py @@ -22,9 +22,14 @@ from litellm.proxy.auth.common_budget_adapter import ( resolve_effective_team_member_budget, ) from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache -from litellm.redis.stores.spend_counters import model_access_group_counter from litellm.proxy.spend_tracking.carried_budget_state import carry_organization_budget_state from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.redis.stores.spend_counters import ( + model_access_group_counter, + org_counter, + tag_counter, + team_member_counter, +) from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget @@ -47,7 +52,7 @@ async def organization_max_budget_check( carry_organization_budget_state(valid_token=valid_token, org_table=org) limit: Final = org.litellm_budget_table.max_budget if org.litellm_budget_table is not None else None facts, result = await read_budget( - BudgetFacts("Organization", limit), CounterRead(f"spend:org:{org_id}", org.spend or 0.0, limit) + BudgetFacts("Organization", limit), CounterRead(org_counter(org_id), org.spend or 0.0, limit) ) if result.alert is not None: info: Final = key_alert_info(valid_token, facts.spend, limit, Litellm_EntityType.ORGANIZATION) @@ -89,7 +94,7 @@ async def check_tag_row(name: str, row: LiteLLM_TagTable | None) -> None: ) facts, result = await read_budget( BudgetFacts("Tag", limit), - CounterRead(f"spend:tag:{name}", (row.spend or 0.0) if row is not None else 0.0, limit, authoritative=True), + CounterRead(tag_counter(name), (row.spend or 0.0) if row is not None else 0.0, limit, authoritative=True), ) if result.denied: raise_budget( @@ -164,7 +169,7 @@ async def check_team_member_budget( facts, result = await read_budget( BudgetFacts("TeamMember", limit), CounterRead( - f"spend:team_member:{valid_token.user_id}:{team_object.team_id}", + team_member_counter(valid_token.user_id, team_object.team_id), (membership.spend or 0.0) if membership is not None else 0.0, limit, ), @@ -200,9 +205,7 @@ async def read_access_group_budget_row(name: str, row: ModelAccessGroupBudget | limit: Final = row.max_budget if row is not None else None facts, _decision = await read_budget( BudgetFacts("AccessGroup", limit), - CounterRead( - model_access_group_counter(name), row.spend if row is not None else 0.0, limit, authoritative=True - ), + CounterRead(model_access_group_counter(name), row.spend if row is not None else 0.0, limit, authoritative=True), ) return AccessGroupBudgetReading(name, facts) diff --git a/litellm/proxy/auth/model_budget_adapter.py b/litellm/proxy/auth/model_budget_adapter.py index 50d804bb960..9c1833f2919 100644 --- a/litellm/proxy/auth/model_budget_adapter.py +++ b/litellm/proxy/auth/model_budget_adapter.py @@ -6,7 +6,7 @@ from typing import Final from pydantic import TypeAdapter import litellm -from litellm._auth import model_budget_data +from litellm._auth import model_budget_aliases, model_budget_data from litellm._auth.failure import NativeAuthFailure from litellm._auth.model_budget_data import ( ModelBudgetEntity, @@ -19,6 +19,7 @@ from litellm._auth.model_budget_data import ( from litellm.caching.caching import DualCache from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.duration_parser import duration_in_seconds +from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry from litellm.llms.bedrock.common_utils import get_bedrock_base_model from litellm.proxy._types import Litellm_EntityType from litellm.proxy.hooks.model_budget_shared_windows import SharedWindows @@ -57,6 +58,10 @@ def _spend(raw: object) -> float: return 0.0 +def _price_entry(model: str) -> object: + return _OBJECT.validate_python(litellm.model_cost.get(model)) # pyright: ignore[reportUnknownMemberType] # raw catalog IO + + @dataclass(frozen=True, slots=True) class ModelBudgetReading: snapshot: ModelBudgetSnapshot @@ -72,6 +77,7 @@ class ModelBudgetIO: process_started_at: float clock: Callable[[], float] = time.time monotonic: Callable[[], float] = time.monotonic + price_entry: Callable[[str], object] = _price_entry async def counter(self, key: str) -> float: redis: Final = self.cache.redis_cache @@ -81,6 +87,70 @@ class ModelBudgetIO: return _spend(raw) +@dataclass(frozen=True, slots=True) +class AliasReading: + row: model_budget_aliases.Row + source: AliasedCostEntry | None + + +def read_alias_models(model: str, io: ModelBudgetIO, readings: tuple[AliasReading, ...] = ()) -> tuple[str, ...]: + decision: Final = model_budget_aliases.decide( + model_budget_aliases.Facts(model, tuple(reading.row for reading in readings)) + ) + match decision.read: + case None: + return decision.models + case "Entry": + raw: Final = io.price_entry(decision.key) + source: Final = raw if isinstance(raw, AliasedCostEntry) else None + entry: Final = AliasReading( + model_budget_aliases.Row(decision.key, source.names if source is not None else ()), source + ) + return read_alias_models(model, io, (*readings, entry)) + case "Prefix": + prefix_known: Final = decision.key.partition("/")[0] in litellm.provider_list + updated: Final = tuple( + replace(reading, row=replace(reading.row, prefix_known=prefix_known)) + if reading.row.key == decision.key + else reading + for reading in readings + ) + return read_alias_models(model, io, updated) + case "Provider": + selected: Final = next(reading for reading in readings if reading.row.key == decision.key) + assert selected.source is not None + provider: Final = selected.source.get("litellm_provider") + provided: Final = replace( + selected, + row=replace(selected.row, provider_read=True, provider=provider if isinstance(provider, str) else None), + ) + return read_alias_models( + model, io, tuple(provided if reading.row.key == decision.key else reading for reading in readings) + ) + + +def _resolve_facts( + model: str, + budgets: Mapping[str, object], + io: ModelBudgetIO, + key_budgets: Mapping[str, object] | None = None, + *, + team: bool = False, + excluded_entries: tuple[str, ...] = (), +) -> ModelBudgetResolveFacts: + base: Final = get_bedrock_base_model(model) + raw_price: Final = as_str_mapping(io.price_entry(base)) or {} + return ModelBudgetResolveFacts( + model=model, + bedrock_base=base, + bedrock_provider=str(raw_price.get("litellm_provider") or ""), + entries=tuple(_entry(name, config) for name, config in budgets.items()), + key_entries=tuple(_entry(name, config) for name, config in (key_budgets or {}).items()), + team=team, + excluded_entries=excluded_entries, + ) + + def resolve_model_budget( model: str, budgets: Mapping[str, object], @@ -89,17 +159,9 @@ def resolve_model_budget( key_budgets: Mapping[str, object] | None = None, *, team: bool = False, + excluded_entries: tuple[str, ...] = (), ) -> ModelBudgetResolution: - base: Final = get_bedrock_base_model(model) - raw_price: Final = as_str_mapping(_OBJECT.validate_python(litellm.model_cost.get(base))) or {} # pyright: ignore[reportUnknownMemberType] # project the legacy pricing map value - facts: Final = ModelBudgetResolveFacts( - model=model, - bedrock_base=base, - bedrock_provider=str(raw_price.get("litellm_provider") or ""), - entries=tuple(_entry(name, config) for name, config in budgets.items()), - key_entries=tuple(_entry(name, config) for name, config in (key_budgets or {}).items()), - team=team, - ) + facts: Final = _resolve_facts(model, budgets, io, key_budgets, team=team, excluded_entries=excluded_entries) initial: Final = model_budget_data.resolve(facts) if not initial.read_groups: return initial @@ -116,7 +178,55 @@ async def read_model_budget( budgets: Mapping[str, object] | None, key_budgets: Mapping[str, object] | None = None, ) -> ModelBudgetReading: - resolution: Final = resolve_model_budget(model, budgets or {}, io, team_id, key_budgets, team=entity == "Team") + models: Final = read_alias_models(model, io) + exclusions: Final = ( + model_budget_data.team_exclusions( + tuple(_resolve_facts(candidate, {}, io, key_budgets, team=True) for candidate in models) + ) + if entity == "Team" + else () + ) + scope_models: Final = read_alias_models(model, io) if entity == "Team" else models + return await _read_scope(io, entity, entity_id, team_id, scope_models, budgets, exclusions) + + +async def _read_scope( + io: ModelBudgetIO, + entity: ModelBudgetEntity, + entity_id: str | None, + team_id: str | None, + models: tuple[str, ...], + budgets: Mapping[str, object] | None, + excluded_entries: tuple[str, ...], + readings: tuple[ModelBudgetReading, ...] = (), +) -> ModelBudgetReading: + decision: Final = model_budget_aliases.scope( + model_budget_aliases.ScopeFacts(models, tuple(reading.snapshot for reading in readings), bool(budgets)) + ) + if decision.read_index is not None: + reading: Final = await _read_single( + io, entity, entity_id, team_id, models[decision.read_index], budgets, excluded_entries + ) + return await _read_scope( + io, entity, entity_id, team_id, models, budgets, excluded_entries, (*readings, reading) + ) + return ( + readings[decision.result_index] + if decision.result_index is not None + else ModelBudgetReading(ModelBudgetSnapshot(), None, False) + ) + + +async def _read_single( + io: ModelBudgetIO, + entity: ModelBudgetEntity, + entity_id: str | None, + team_id: str | None, + model: str, + budgets: Mapping[str, object] | None, + excluded_entries: tuple[str, ...], +) -> ModelBudgetReading: + resolution: Final = resolve_model_budget(model, budgets or {}, io, team_id, excluded_entries=excluded_entries) entry: Final = resolution.entry if entry is None: return ModelBudgetReading(ModelBudgetSnapshot(), None, False) diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index e030a85b0bf..f2317e49c9d 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -1046,28 +1046,6 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): model: str, exceeded_message: str, key_model_max_budget: Mapping[str, object] | None = None, - ) -> bool: - for candidate in expand_model_latest_alias(model): - await self._is_single_entity_within_model_budget( - entity_type=entity_type, - entity_id=entity_id, - team_id=team_id, - model_max_budget=model_max_budget, - model=candidate, - exceeded_message=exceeded_message, - key_model_max_budget=key_model_max_budget, - ) - return True - - async def _is_single_entity_within_model_budget( - self, - entity_type: Litellm_EntityType, - entity_id: str | None, - team_id: str | None, - model_max_budget: Mapping[str, object] | None, - model: str, - exceeded_message: str, - key_model_max_budget: Mapping[str, object] | None = None, ) -> bool: entities: Final[Mapping[Litellm_EntityType, ModelBudgetEntity]] = { Litellm_EntityType.KEY: "Key", diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index e264363e463..b7c948d7c0d 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -26,7 +26,9 @@ def native_redis_breaker_states() -> list[str]: ... class NativeAuthContinuation: def resume(self, answer: object) -> tuple[dict[str, object], NativeAuthContinuation | None]: ... -def request_auth_start(facts: object, settings: object, now: object, resolved: bool) -> tuple[dict[str, object], NativeAuthContinuation | None]: ... +def request_auth_start( + facts: object, settings: object, now: object, resolved: bool +) -> tuple[dict[str, object], NativeAuthContinuation | None]: ... def request_auth_budgets(entry: str) -> tuple[dict[str, object], NativeAuthContinuation | None]: ... def request_auth_access(facts: object) -> dict[str, object]: ... def request_auth_agent_access(facts: object) -> str: ... @@ -51,6 +53,9 @@ def request_auth_budget_route(facts: object) -> str: ... def request_auth_model_budget_resolve(facts: object) -> dict[str, object]: ... def request_auth_model_budget_query(facts: object) -> dict[str, object]: ... def request_auth_model_budget_check(facts: object) -> dict[str, object]: ... +def request_auth_model_budget_aliases(facts: object) -> dict[str, object]: ... +def request_auth_model_budget_scope(facts: object) -> dict[str, object]: ... +def request_auth_model_budget_team_exclusions(facts: object) -> object: ... def request_auth_select_api_key(keys: object) -> dict[str, str] | None: ... def request_auth_select_websocket_api_key(keys: object) -> dict[str, str]: ... def request_auth_scheme_credential(value: str) -> str | None: ... @@ -196,9 +201,7 @@ class NativeRouterUsage: def attach(self, client: NativeRedisClient) -> bool: ... @property def has_redis(self) -> bool: ... - def increment( - self, counter: _UsageCounter, value: float, ttl_seconds: int - ) -> Future[float]: ... + def increment(self, counter: _UsageCounter, value: float, ttl_seconds: int) -> Future[float]: ... def increment_deployment( self, deployment_id: str, @@ -207,12 +210,8 @@ class NativeRouterUsage: tokens: float, requests: int, ) -> Future[None]: ... - def get_many( - self, counters: Sequence[_UsageCounter] - ) -> Future[list[float | None]]: ... - def increment_sync( - self, counter: _UsageCounter, value: float, ttl_seconds: int - ) -> float: ... + def get_many(self, counters: Sequence[_UsageCounter]) -> Future[list[float | None]]: ... + def increment_sync(self, counter: _UsageCounter, value: float, ttl_seconds: int) -> float: ... def increment_deployment_sync( self, deployment_id: str, @@ -221,9 +220,7 @@ class NativeRouterUsage: tokens: float, requests: int, ) -> None: ... - def get_many_sync( - self, counters: Sequence[_UsageCounter] - ) -> list[float | None]: ... + def get_many_sync(self, counters: Sequence[_UsageCounter]) -> list[float | None]: ... def flush_local_counters(self) -> None: ... def snapshot(self) -> dict[str, float]: ... @@ -399,7 +396,9 @@ class NativeAuthCache: def get_sync(self, key: str) -> object: ... def set(self, key: str, value: object, ttl_seconds: float | None = None) -> Future[None]: ... def set_sync(self, key: str, value: object, ttl_seconds: float | None = None) -> None: ... - def declare_set(self, key: str, value: object, ttl_seconds: float | None, send: NativeSend) -> NativePending[None]: ... + def declare_set( + self, key: str, value: object, ttl_seconds: float | None, send: NativeSend + ) -> NativePending[None]: ... def declare_delete(self, key: str, send: NativeSend) -> NativePending[None]: ... def set_many(self, entries: Sequence[tuple[str, object]], ttl_seconds: float | None = None) -> Future[None]: ... def write_back(self, entries: Sequence[tuple[str, object, float | None]], send: NativeSend) -> Future[None]: ... @@ -1137,7 +1136,9 @@ class RouterCoreHandle: callbacks: Mapping[str, bool] | None = None, ) -> RouterCall[_SuccessReportSteps]: ... # Count a success against any id, driven like `report_failure`; `done` carries `key`. - def count_success(self, id: str, callbacks: Mapping[str, bool] | None = None) -> RouterCall[_SuccessReportSteps]: ... + def count_success( + self, id: str, callbacks: Mapping[str, bool] | None = None + ) -> RouterCall[_SuccessReportSteps]: ... # [(id, tpm_key, rpm_key)] for this minute, or None for a group with no deployment. def usage_keys(self, model_group: str) -> list[tuple[str, str, str]] | None: ... # Every deployment as a `model_list` entry, in config order. @@ -1236,52 +1237,29 @@ 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]: ... - 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: ... - def request_auth_end_user_normalize_id(facts: object) -> object: ... - def request_auth_end_user_default_id(facts: object) -> object: ... - def request_auth_end_user_apply(facts: object) -> object: ... - def request_auth_end_user_registry(facts: object) -> object: ... - def request_auth_end_user_assignments(facts: object) -> object: ... - def request_auth_end_user_lookup(facts: object) -> object: ... - def request_auth_end_user_validation(facts: object) -> object: ... - def request_auth_budget_lookup_organization(facts: object) -> object: ... def request_auth_budget_lookup_member(facts: object) -> object: ... def request_auth_budget_lookup_emails(facts: object) -> object: ... def request_auth_budget_lookup_soft(facts: object) -> bool: ... - def request_auth_route_role(facts: object) -> object: ... - def request_auth_route_role_api(facts: object) -> object: ... - def request_auth_route_role_jwt(facts: object) -> object: ... - def request_auth_route_role_inside(facts: object) -> object: ... - def request_auth_route_role_expand(facts: object) -> object: ... - def request_auth_route_role_listed(facts: object) -> object: ... - def request_auth_org_admin_raw(facts: object) -> object: ... diff --git a/tests/unit/_auth/model_alias_reference.py b/tests/unit/_auth/model_alias_reference.py new file mode 100644 index 00000000000..b656699e894 --- /dev/null +++ b/tests/unit/_auth/model_alias_reference.py @@ -0,0 +1,38 @@ +from collections.abc import Mapping, Sequence +from typing import Final + +from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry + + +def _entry(key: str, catalog: Mapping[str, object]) -> AliasedCostEntry | None: + value: Final = catalog.get(key) + return value if isinstance(value, AliasedCostEntry) and value.names[:1] not in ((), (key,)) else None + + +def _provider(key: str, value: AliasedCostEntry, providers: Sequence[str]) -> str | None: + prefix, separator, _ = key.partition("/") + if separator and prefix in providers: + return prefix + provider: Final = value.get("litellm_provider") + return provider if isinstance(provider, str) else None + + +def _target(key: str, provider: str, catalog: Mapping[str, object], providers: Sequence[str]) -> str | None: + value: Final = _entry(key, catalog) + if value is None or _provider(key, value, providers) != provider: + return None + return value.names[0].removeprefix(f"{provider}/") + + +def expand(model: str, catalog: Mapping[str, object], providers: Sequence[str]) -> tuple[str, ...]: + """Alias expansion frozen from main 2f430873aead937b9112e683703956eb1ab179af.""" + provider, separator, name = model.partition("/") + if not separator: + value: Final = _entry(model, catalog) + return (model,) if value is None else (model, value.names[0]) + target: Final = ( + _target(f"{provider}/{name}", provider, catalog, providers) + or _target(name, provider, catalog, providers) + or name + ) + return (model, f"{provider}/{target}") if target != name else (model,) diff --git a/tests/unit/_auth/test_model_budget_aliases.py b/tests/unit/_auth/test_model_budget_aliases.py new file mode 100644 index 00000000000..0d1e1b5bd60 --- /dev/null +++ b/tests/unit/_auth/test_model_budget_aliases.py @@ -0,0 +1,66 @@ +from collections.abc import Mapping +from typing import Final +from unittest.mock import Mock, call + +import pytest + +import litellm +from litellm._auth import model_budget_aliases as native +from litellm._auth.model_budget_data import ModelBudgetSnapshot +from litellm.caching.caching import DualCache +from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry +from litellm.proxy.auth.model_budget_adapter import ModelBudgetIO, read_alias_models +from litellm.proxy.hooks.model_budget_shared_windows import SharedWindows + +from .model_alias_reference import expand + + +@pytest.mark.parametrize("model", ("latest", "openai/latest", "unknown/latest", "bedrock/latest")) +@pytest.mark.parametrize("names", ((), ("latest",), ("openai/latest",), ("target",), ("",), ("openai/target",))) +@pytest.mark.parametrize("provider", (None, "openai", "unknown", "bedrock_converse", 42)) +def test_alias_lookup_matches_frozen_main(model: str, names: tuple[str, ...], provider: object) -> None: + row: Final = AliasedCostEntry({"litellm_provider": provider}, names) + catalog: Final[Mapping[str, object]] = {model: row, "latest": row} + expected: Final = expand(model, catalog, litellm.provider_list) + lookup: Final = Mock(side_effect=catalog.get) + io: Final = ModelBudgetIO( + DualCache(), lambda _model, _team: frozenset(), Mock(spec=SharedWindows), 0, price_entry=lookup + ) + assert read_alias_models(model, io) == expected + assert lookup.call_args_list[0].args == (model,) + + +@pytest.mark.parametrize("raw", (None, {}, {"alias_of": "target"}, {"names": ["target"]})) +def test_plain_catalog_rows_do_not_become_aliases(raw: object) -> None: + io: Final = ModelBudgetIO( + DualCache(), lambda _model, _team: frozenset(), Mock(spec=SharedWindows), 0, price_entry=lambda _: raw + ) + assert read_alias_models("latest", io) == ("latest",) + + +@pytest.mark.parametrize("direct", (True, False)) +def test_catalog_reads_bare_name_only_when_qualified_alias_misses(direct: bool) -> None: + row: Final = AliasedCostEntry({"litellm_provider": "openai"}, ("target", "latest")) + catalog: Final = {"openai/latest": row if direct else None, "latest": row} + lookup: Final = Mock(side_effect=catalog.get) + io: Final = ModelBudgetIO( + DualCache(), lambda _model, _team: frozenset(), Mock(spec=SharedWindows), 0, price_entry=lookup + ) + assert read_alias_models("openai/latest", io) == ("openai/latest", "openai/target") + assert lookup.call_args_list == ([call("openai/latest")] if direct else [call("openai/latest"), call("latest")]) + + +@pytest.mark.parametrize( + ("snapshots", "expected"), + ( + ((), native.ScopeDecision(0, None)), + ((ModelBudgetSnapshot(1, 0),), native.ScopeDecision(1, None)), + ((ModelBudgetSnapshot(1, 1),), native.ScopeDecision(None, 0)), + ((ModelBudgetSnapshot(1, 0), ModelBudgetSnapshot(2, 2)), native.ScopeDecision(None, 1)), + ((ModelBudgetSnapshot(1, 0), ModelBudgetSnapshot(2, 0)), native.ScopeDecision(None, 1)), + ), +) +def test_scope_reads_only_until_first_denial( + snapshots: tuple[ModelBudgetSnapshot, ...], expected: native.ScopeDecision +) -> None: + assert native.scope(native.ScopeFacts(("alias", "target"), snapshots)) == expected diff --git a/tests/unit/_auth/test_model_budget_data.py b/tests/unit/_auth/test_model_budget_data.py index 6fb053d2d4b..314a287681c 100644 --- a/tests/unit/_auth/test_model_budget_data.py +++ b/tests/unit/_auth/test_model_budget_data.py @@ -55,6 +55,27 @@ def entry(name: str, limit: float | None = 1, duration: str | None = "1h") -> Mo return ModelBudgetEntry(name, limit, duration, 3600 if duration is not None else None) +@pytest.mark.parametrize("limit", (None, -1.0, 0.0, 1.0, inf, nan)) +@pytest.mark.parametrize("overrides", (("alias",), ("target",), ("alias", "target"))) +def test_team_exact_entries_are_removed_across_all_alias_names(limit: float | None, overrides: tuple[str, ...]) -> None: + key: Final = tuple(entry(name, limit) for name in overrides) + facts: Final = tuple(ModelBudgetResolveFacts(model, model, "", (), key, True) for model in ("alias", "target")) + excluded: Final = native.team_exclusions(facts) + assert excluded == (tuple(sorted(overrides)) if limit is not None and limit >= 0 else ()) + resolution: Final = native.resolve( + ModelBudgetResolveFacts( + "alias", + "alias", + "", + (entry("alias"), entry("target"), entry("*")), + access_groups=("alias", "target"), + excluded_entries=excluded, + ) + ) + expected: Final = "alias" if "alias" not in excluded else "target" if "target" not in excluded else "*" + assert resolution.entry is not None and resolution.entry.name == expected + + @pytest.mark.parametrize( ("entries", "groups", "expected"), ( diff --git a/tests/unit/proxy/auth/test_model_budget_adapter.py b/tests/unit/proxy/auth/test_model_budget_adapter.py index 0393b018e5d..de22d4c78dc 100644 --- a/tests/unit/proxy/auth/test_model_budget_adapter.py +++ b/tests/unit/proxy/auth/test_model_budget_adapter.py @@ -1,14 +1,22 @@ from collections.abc import Awaitable, Mapping from dataclasses import dataclass, field from typing import Final -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, call import pytest import litellm from litellm._auth import model_budget_data +from litellm._auth.budget_pricing import ModelPriceFacts +from litellm._auth.fallback_candidate import ( + CandidateBudgetRead, + CandidateBudgetScope, + FallbackCandidateFacts, + fallback_candidate, +) from litellm._auth.model_budget_data import ModelBudgetEntity from litellm.caching.caching import DualCache +from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry from litellm.proxy._types import Litellm_EntityType from litellm.proxy.auth.model_budget_adapter import ( ModelBudgetIO, @@ -136,3 +144,164 @@ async def test_shared_read_failure_is_terminal_and_not_replayed() -> None: await read_model_budget(io, "User", "id", None, "model", {"model": {"max_budget": 1, "budget_duration": "1h"}}) assert raised.value is failure assert windows.read.await_count == 1 + + +@pytest.fixture +def alias_catalog(monkeypatch: pytest.MonkeyPatch) -> Mapping[str, object]: + catalog: Final = {"budget-latest": AliasedCostEntry({}, ("budget-target", "budget-latest"))} + monkeypatch.setattr(litellm, "model_cost", catalog) + return catalog + + +@pytest.mark.asyncio +@pytest.mark.parametrize("entity", ("Key", "User", "Team", "TeamMember", "EndUser")) +@pytest.mark.parametrize( + ("alias_spend", "target_spend", "expected"), + ((1.0, 2.0, "budget-latest"), (0.0, 2.0, "budget-target"), (0.0, 0.0, None)), +) +async def test_raw_alias_scope_preserves_first_denial_and_counter_order( + alias_catalog: Mapping[str, object], + entity: ModelBudgetEntity, + alias_spend: float, + target_spend: float, + expected: str | None, +) -> None: + cache: Final = DualCache() + scope: Final = { + "Key": Litellm_EntityType.KEY, + "User": Litellm_EntityType.USER, + "Team": Litellm_EntityType.TEAM, + "TeamMember": Litellm_EntityType.TEAM_MEMBER, + "EndUser": Litellm_EntityType.END_USER, + }[entity] + alias_key: Final = reference.model_budget_spend_cache_key(scope, "id", "budget-latest", "1h") + target_key: Final = reference.model_budget_spend_cache_key(scope, "id", "budget-target", "1h") + cache.in_memory_cache.set_cache(alias_key, alias_spend) # pyright: ignore[reportUnknownMemberType] # legacy cache setter has no annotations + cache.in_memory_cache.set_cache(target_key, target_spend) # pyright: ignore[reportUnknownMemberType] # legacy cache setter has no annotations + shared_read: Final = AsyncMock(return_value=None) + windows: Final = Windows(read=shared_read) + io: Final = ModelBudgetIO( + cache, groups, windows, 0, clock=lambda: 100, monotonic=lambda: 0, price_entry=alias_catalog.get + ) + reading: Final = await read_model_budget( + io, + entity, + "id", + None, + "budget-latest", + { + "budget-latest": {"max_budget": 1, "budget_duration": "1h"}, + "budget-target": {"max_budget": 2, "budget_duration": "1h"}, + }, + ) + decision: Final = model_budget_data.check(reading.snapshot) + assert decision.denied is (expected is not None) + assert reading.budget_model == (expected or "budget-target") + assert decision.current_spend == (alias_spend if expected == "budget-latest" else target_spend) + candidate: Final = fallback_candidate( + FallbackCandidateFacts( + True, + True, + True, + membership="Loaded", + price=ModelPriceFacts(cached=False), + budget_models=("budget-latest",), + budget_scopes=(CandidateBudgetScope(entity, True, True),), + budgets=(CandidateBudgetRead(0, entity, reading.snapshot),), + ) + ) + assert candidate.kind == ("Complete" if expected is not None else "ReadAuthorization") + assert candidate.usable is False + expected_shared: Final = ( + [] if entity == "EndUser" or alias_spend == 1 else [alias_key] if target_spend == 2 else [alias_key, target_key] + ) + assert shared_read.call_args_list == [call(spend_key=key, budget_duration="1h") for key in expected_shared] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("overrides", (("budget-latest",), ("budget-target",), ("budget-latest", "budget-target"))) +async def test_team_overrides_apply_across_alias_and_target( + alias_catalog: Mapping[str, object], overrides: tuple[str, ...] +) -> None: + io: Final = ModelBudgetIO(DualCache(), groups, Windows(), 0, price_entry=alias_catalog.get) + budgets: Final = {name: {"max_budget": 0, "budget_duration": "1h"} for name in ("budget-latest", "budget-target")} + key: Final = {name: {"max_budget": 1, "budget_duration": "1h"} for name in overrides} + reading: Final = await read_model_budget(io, "Team", "id", "id", "budget-latest", budgets, key) + expected: Final = next((name for name in budgets if name not in overrides), None) + assert model_budget_data.check(reading.snapshot).denied is (expected is not None) + assert reading.budget_model == expected + + +@pytest.mark.asyncio +async def test_shared_native_model_flow_enforces_target_before_user_scope(alias_catalog: Mapping[str, object]) -> None: + from starlette.requests import Request + + from litellm._auth.decision import Deny, ModelBudgetExceeded + from litellm._auth.driver import drive_model_budgets + from litellm._auth.rust_core import RustCore + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import ( + _ModelBudgetEffects, # pyright: ignore[reportPrivateUsage] # exercise the production raw effect boundary + ) + + token: Final = UserAPIKeyAuth( + token="id", + user_id="owner", + key_alias="alias-key", + model_max_budget={"budget-target": {"max_budget": 0, "budget_duration": "1h"}}, + user_model_max_budget={"budget-latest": {"max_budget": 0, "budget_duration": "1h"}}, + ) + effects: Final = _ModelBudgetEffects( + request=Request({"type": "http", "headers": [], "path": "/chat/completions"}), + request_data={"model": "budget-latest"}, + valid_token=token, + route="/chat/completions", + limiter=reference._PROXY_VirtualKeyModelMaxBudgetLimiter(DualCache(), clock=lambda: 100), # pyright: ignore[reportPrivateUsage] # concrete limiter integration + llm_router=None, + llm_model_list=None, + ) + outcome: Final = await drive_model_budgets(RustCore().check_model_budgets(), effects) + assert isinstance(outcome, Deny) and isinstance(outcome.denial, ModelBudgetExceeded) + assert (outcome.denial.entity_type, outcome.denial.entity_id, outcome.denial.max_budget) == ("key", "id", 0.0) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("authorization_denied", (False, True)) +async def test_fallback_probe_retains_original_denial_without_authorization_retries( + alias_catalog: Mapping[str, object], authorization_denied: bool +) -> None: + from starlette.requests import Request + + from litellm._auth.decision import Deny, ModelBudgetExceeded + from litellm._auth.driver import drive_model_budgets + from litellm._auth.rust_core import RustCore + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import ( + _ModelBudgetEffects, # pyright: ignore[reportPrivateUsage] # exercise the production raw effect boundary + ) + + data: Final[dict[str, object]] = {"model": "initial"} + token: Final = UserAPIKeyAuth( + token="id", + models=["initial", "later"], + model_max_budget={ + "initial": {"max_budget": 0, "budget_duration": "1h"}, + "budget-target": {"max_budget": 1 if authorization_denied else 0, "budget_duration": "1h"}, + }, + budget_fallbacks={"initial": ["budget-latest", "later"] if authorization_denied else ["budget-latest"]}, + ) + effects: Final = _ModelBudgetEffects( + request=Request({"type": "http", "headers": [], "path": "/chat/completions"}), + request_data=data, + valid_token=token, + route="/chat/completions", + limiter=reference._PROXY_VirtualKeyModelMaxBudgetLimiter( # pyright: ignore[reportPrivateUsage] # concrete limiter integration + DualCache(), clock=lambda: 100, shared_windows=Windows(read=AsyncMock(return_value=None)) + ), + llm_router=None, + llm_model_list=None, + ) + outcome: Final = await drive_model_budgets(RustCore().check_model_budgets(), effects) + assert isinstance(outcome, Deny) and isinstance(outcome.denial, ModelBudgetExceeded) + assert "model=initial" in outcome.denial.message + assert data == {"model": "initial"}