diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 8592abc06ae..33cf423106a 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -39,32 +39,32 @@ mod _native { use crate::model_registry::bundled_model_cost_map; #[pymodule_export] use crate::request_auth::{ - NativeAuthContinuation, - request_auth_access, - request_auth_admission, - request_auth_budget_alert, - request_auth_budget_price, - request_auth_budget_route, - request_auth_budget_window_start, - request_auth_budgets, - request_auth_common_budget, - request_auth_fallback_access, - request_auth_litellm_header_token, - request_auth_managed, - request_auth_member_limit, - request_auth_model_access, - request_auth_model_budget_check, - request_auth_model_budget_query, - request_auth_model_budget_resolve, - request_auth_model_groups, - request_auth_org_admin, - request_auth_org_route, - request_auth_route_access, - request_auth_scheme_credential, - request_auth_select_api_key, - request_auth_select_websocket_api_key, - request_auth_start, - ||||||| parent of da364c75e2 (refactor(auth): port model and route policy to native decisions) NativeAuthContinuation, + NativeAuthContinuation, + request_auth_access, + request_auth_admission, + 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_fallback_access, + request_auth_litellm_header_token, + request_auth_managed, + request_auth_member_limit, + request_auth_model_access, + request_auth_model_budget_check, + request_auth_model_budget_query, + request_auth_model_budget_resolve, + request_auth_model_groups, + request_auth_org_admin, + request_auth_org_route, + request_auth_route_access, + 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 23f682f1be4..cf7151048e8 100644 --- a/litellm-rust/crates/python-bridge/src/request_auth.rs +++ b/litellm-rust/crates/python-bridge/src/request_auth.rs @@ -277,3 +277,15 @@ pub(crate) fn request_auth_fallback_access(py: Python<'_>, facts: &Bound<'_, PyA let decision = py.detach(move || core::fallback_access::decide(facts)).map_err(map_error)?; to_py(py, &decision) } + +#[pyfunction] +pub(crate) fn request_auth_budget_exemption( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::budget_pricing::ExemptionFacts = from_py(facts)?; + to_py( + py, + &py.detach(move || core::budget_pricing::exemption(&facts)), + ) +} diff --git a/litellm-rust/crates/request-auth-core/src/budget_pricing.rs b/litellm-rust/crates/request-auth-core/src/budget_pricing.rs index cc33266a38f..29329d5c0ea 100644 --- a/litellm-rust/crates/request-auth-core/src/budget_pricing.rs +++ b/litellm-rust/crates/request-auth-core/src/budget_pricing.rs @@ -117,6 +117,63 @@ pub fn route(facts: &RouteBudgetFacts) -> RouteBudgetDecision { if mcp_discovery { Skip } else { ReadPricing } } +#[derive(Clone, Debug, Deserialize)] +pub struct ExemptionFacts { + pub route: RouteBudgetFacts, + pub model_input_present: bool, + pub router_present: bool, + pub model_count: usize, + #[serde(default)] + pub cache_read: bool, + pub prices: Vec, +} + +#[derive(Clone, Debug, PartialEq, Serialize)] +pub struct ExemptionDecision { + pub read_pass_through: bool, + pub read_price_cache: bool, + pub read_model: Option, + pub exempt: bool, +} + +pub fn exemption(facts: &ExemptionFacts) -> ExemptionDecision { + let done = |exempt| ExemptionDecision { + read_pass_through: false, + read_price_cache: false, + read_model: None, + exempt, + }; + match route(&facts.route) { + RouteBudgetDecision::ReadPassThrough => { + return ExemptionDecision { + read_pass_through: true, + ..done(false) + }; + } + RouteBudgetDecision::Skip => return done(true), + RouteBudgetDecision::ReadPricing => (), + } + if !facts.model_input_present || !facts.router_present { + return done(false); + } + if !facts.cache_read { + return ExemptionDecision { + read_price_cache: true, + ..done(false) + }; + } + if facts.prices.iter().any(|facts| !price(facts).zero_cost) { + return done(false); + } + if facts.prices.len() < facts.model_count { + return ExemptionDecision { + read_model: Some(facts.prices.len()), + ..done(false) + }; + } + done(true) +} + #[cfg(test)] mod tests { use super::*; @@ -154,4 +211,51 @@ mod tests { } ); } + #[rstest] + #[case::read_first(vec![], Some(0), false)] + #[case::read_second(vec![true], Some(1), false)] + #[case::priced_first(vec![false], None, false)] + #[case::all_free(vec![true, true], None, true)] + #[case::priced_second(vec![true, false], None, false)] + fn exemption_reads_models_in_order( + #[case] cached: Vec, + #[case] read: Option, + #[case] exempt: bool, + ) { + let facts = ExemptionFacts { + route: RouteBudgetFacts { + side_effect_route: false, + discovery_route: false, + llm_route: true, + route: "/chat/completions".to_owned(), + method: None, + body_present: false, + registered_pass_through: None, + }, + model_input_present: true, + router_present: true, + model_count: 2, + cache_read: true, + prices: cached + .into_iter() + .map(|value| ModelPriceFacts { + cached: Some(value), + lookup_failed: false, + group_found: false, + input_price: None, + output_price: None, + deployments: None, + }) + .collect(), + }; + assert_eq!( + exemption(&facts), + ExemptionDecision { + read_pass_through: false, + read_price_cache: false, + read_model: read, + exempt + } + ); + } } diff --git a/litellm/_auth/budget_pricing.py b/litellm/_auth/budget_pricing.py index 6dd560f68e8..b63b5ea699a 100644 --- a/litellm/_auth/budget_pricing.py +++ b/litellm/_auth/budget_pricing.py @@ -58,3 +58,29 @@ def price(facts: ModelPriceFacts) -> PriceDecision: def route(facts: RouteBudgetFacts) -> RouteBudgetDecision: return _ROUTE_DECISION.validate_python(_native.request_auth_budget_route(_ROUTE.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes the validated frozen facts for the native boundary + + +@dataclass(frozen=True, slots=True) +class ExemptionFacts: + route: RouteBudgetFacts + model_input_present: bool + router_present: bool + model_count: int + prices: tuple[ModelPriceFacts, ...] = () + cache_read: bool = False + + +@dataclass(frozen=True, slots=True) +class ExemptionDecision: + read_pass_through: bool + read_price_cache: bool + read_model: int | None + exempt: bool + + +_EXEMPTION: Final = TypeAdapter(ExemptionFacts) +_EXEMPTION_DECISION: Final = TypeAdapter(ExemptionDecision) + + +def exemption(facts: ExemptionFacts) -> ExemptionDecision: + return _EXEMPTION_DECISION.validate_python(_native.request_auth_budget_exemption(_EXEMPTION.dump_python(facts))) # pyright: ignore[reportAny] # Pydantic serializes frozen facts diff --git a/litellm/proxy/auth/budget_pricing_adapter.py b/litellm/proxy/auth/budget_pricing_adapter.py index c42cd05d32a..5baa51eaddd 100644 --- a/litellm/proxy/auth/budget_pricing_adapter.py +++ b/litellm/proxy/auth/budget_pricing_adapter.py @@ -56,7 +56,7 @@ 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 Exception: + 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: return facts @@ -64,18 +64,22 @@ 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 Exception: + except Exception: # noqa: BLE001 # preserve lookup failure as non-exempt without caching return ModelPriceFacts(lookup_failed=True) return replace(facts, deployments=deployments) -def _model_zero(model: str, router: Router, cache: Mapping[str, bool] | None) -> bool: - cached: Final = cache.get(model) if cache is not None else None +def _read_cached_model_price(model: str, router: Router, cache: Mapping[str, bool] | None) -> ModelPriceFacts: + cached: Final = cache.get(model) if isinstance(cache, dict) else None facts: Final = read_model_price_facts(model, router, cached) decision: Final = budget_pricing.price(facts) if isinstance(cache, dict) and decision.cache_value is not None: - cache[model] = decision.cache_value # rebind-ok: preserve registry cache writes selected by native policy - return decision.zero_cost + cache[model] = decision.cache_value # rebind-ok: native-selected registry cache write + return facts + + +def read_model_cost_facts(model: str, router: Router) -> ModelPriceFacts: + return _read_cached_model_price(model, router, router.registry().zero_cost()) def is_model_cost_zero(model: str | Sequence[str] | None, llm_router: Router | None) -> bool: @@ -83,7 +87,7 @@ def is_model_cost_zero(model: str | Sequence[str] | None, llm_router: Router | N return budget_pricing.price(ModelPriceFacts()).zero_cost names: Final = (model,) if isinstance(model, str) else tuple(model) cache: Final = llm_router.registry().zero_cost() - return all(_model_zero(name, llm_router, cache) for name in names) + return all(budget_pricing.price(_read_cached_model_price(name, llm_router, cache)).zero_cost for name in names) def read_route_budget_facts(route: str, body: Mapping[str, object] | None = None) -> RouteBudgetFacts: @@ -115,6 +119,51 @@ def route_skips_budget_checks(route: str) -> bool: return _route_decision(route) == "Skip" +def read_budget_exemption_facts( + route: str, + model: str | Sequence[str] | None, + llm_router: Router | None, + *, + request_body: Mapping[str, object] | None = None, +) -> budget_pricing.ExemptionFacts: + names: Final = (model,) if isinstance(model, str) else tuple(model or ()) + initial: Final = budget_pricing.ExemptionFacts( + read_route_budget_facts(route, request_body), + model is not None, + llm_router is not None, + len(names), + ) + return _read_exemption(initial, names, llm_router, None) + + +def _read_exemption( + facts: budget_pricing.ExemptionFacts, + models: tuple[str, ...], + router: Router | None, + cache: Mapping[str, bool] | None, +) -> budget_pricing.ExemptionFacts: + current = facts # rebind-ok: native-selected IO driver state + selected_cache = cache # rebind-ok: cache handle acquired at the first model read + while True: + decision = budget_pricing.exemption(current) + if decision.read_pass_through: + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import InitPassThroughEndpointHelpers + + registered = InitPassThroughEndpointHelpers.is_registered_pass_through_route(route=current.route.route) + current = replace(current, route=replace(current.route, registered_pass_through=registered)) + continue + if decision.read_price_cache: + assert router is not None + selected_cache = router.registry().zero_cost() + current = replace(current, cache_read=True) + continue + if decision.read_model is None: + return current + assert router is not None + price = _read_cached_model_price(models[decision.read_model], router, selected_cache) + current = replace(current, prices=(*current.prices, price)) + + def request_skips_budget_checks( route: str, model: str | Sequence[str] | None, @@ -122,6 +171,5 @@ def request_skips_budget_checks( *, request_body: Mapping[str, object] | None = None, ) -> bool: - if _route_decision(route, request_body) == "Skip": - return True - return is_model_cost_zero(model, llm_router) + facts: Final = read_budget_exemption_facts(route, model, llm_router, request_body=request_body) + return budget_pricing.exemption(facts).exempt diff --git a/litellm/proxy/auth/budget_windows_adapter.py b/litellm/proxy/auth/budget_windows_adapter.py index 9b593817442..250591621b2 100644 --- a/litellm/proxy/auth/budget_windows_adapter.py +++ b/litellm/proxy/auth/budget_windows_adapter.py @@ -28,7 +28,7 @@ def _window(raw: object) -> BudgetWindow: def _seconds(duration: str) -> int | None: try: return duration_in_seconds(duration) - except Exception: + except Exception: # noqa: BLE001 # preserve invalid-duration skip behavior return None diff --git a/litellm/proxy/auth/common_budget_lookups.py b/litellm/proxy/auth/common_budget_lookups.py index 393b7a24e42..8a58e43a163 100644 --- a/litellm/proxy/auth/common_budget_lookups.py +++ b/litellm/proxy/auth/common_budget_lookups.py @@ -1,6 +1,8 @@ -from collections.abc import Sequence +from collections.abc import Mapping, Sequence +from dataclasses import dataclass from typing import Final +from litellm._auth import common_budgets from litellm._auth.common_budgets import BudgetFacts from litellm.proxy._types import ( Litellm_EntityType, @@ -51,7 +53,7 @@ async def organization_max_budget_check( proxy_logging_obj=proxy_logging_obj, include_budget_table=True, ) - except Exception: + except Exception: # noqa: BLE001 # preserve organization lookup failure handling return if org is None: return @@ -218,3 +220,50 @@ async def check_team_member_budget( f"{valid_token.user_id}:{team_object.team_id}", f"Budget has been exceeded! User={valid_token.user_id} in Team={team_object.team_id} Current cost: {facts.spend}, Max budget: {limit}", ) + + +@dataclass(frozen=True, slots=True) +class AccessGroupBudgetReading: + name: str + facts: BudgetFacts + + +async def read_access_group_budget_row(name: str, row: ModelAccessGroupBudget | None) -> AccessGroupBudgetReading: + limit: Final = row.max_budget if row is not None else None + facts, _decision = await read_budget( + BudgetFacts("AccessGroup", limit), + CounterRead( + model_access_group_spend_counter_key(name), row.spend if row is not None else 0.0, limit, authoritative=True + ), + ) + return AccessGroupBudgetReading(name, facts) + + +async def read_access_group_budget_facts( + groups: Sequence[str], + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, +) -> tuple[AccessGroupBudgetReading, ...]: + from litellm.proxy.auth.auth_checks import get_model_access_group_budgets_batch + + if prisma_client is None or not groups: + return () + rows: Final = await get_model_access_group_budgets_batch( + access_group_names=groups, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + return await _read_access_group_rows(tuple(groups), rows) + + +async def _read_access_group_rows( + groups: tuple[str, ...], + rows: Mapping[str, ModelAccessGroupBudget], +) -> tuple[AccessGroupBudgetReading, ...]: + readings: tuple[AccessGroupBudgetReading, ...] = () # rebind-ok: native-selected IO driver accumulator + for name in groups: + reading = await read_access_group_budget_row(name, rows.get(name)) + readings = (*readings, reading) + if common_budgets.check(reading.facts).denied: + return readings + return readings diff --git a/litellm/proxy/auth/model_budget_adapter.py b/litellm/proxy/auth/model_budget_adapter.py index 73e230d98ec..10aae9547c5 100644 --- a/litellm/proxy/auth/model_budget_adapter.py +++ b/litellm/proxy/auth/model_budget_adapter.py @@ -40,7 +40,7 @@ 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 Exception: + except Exception: # noqa: BLE001 # preserve malformed entry fallthrough return ModelBudgetEntry(name, None, None, None) return ModelBudgetEntry(name, config.max_budget, config.budget_duration, seconds) diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index cc2e91cd93e..af4da69d948 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -15,6 +15,8 @@ from typing import ( ) import litellm +from litellm._auth import model_budget_data +from litellm._auth.model_budget_data import ModelBudgetEntity from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache @@ -40,6 +42,7 @@ from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import BudgetConfig, LlmProviders, StandardLoggingPayload if TYPE_CHECKING: + from litellm.proxy.auth.model_budget_adapter import ModelBudgetReading from litellm.router import Router VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX: Final = "virtual_key_spend" @@ -947,9 +950,8 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): entity_type=Litellm_EntityType.TEAM, entity_id=team_id, team_id=team_id, - model_max_budget=team_model_budget_entries( - model=model, team_model_max_budget=team_model_max_budget, key_model_max_budget=key_model_max_budget - ), + model_max_budget=team_model_max_budget, + key_model_max_budget=key_model_max_budget, model=model, exceeded_message=f"LiteLLM Team: {team_id}, exceeded budget for model={model}", ) @@ -973,6 +975,27 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): ), ) + async def read_model_budget_raw( + self, + entity: ModelBudgetEntity, + entity_id: str | None, + team_id: str | None, + model: str, + budgets: Mapping[str, object] | None, + key_budgets: Mapping[str, object] | None = None, + ) -> "ModelBudgetReading": + from litellm.proxy.auth.model_budget_adapter import ModelBudgetIO, read_model_budget + + io: Final = ModelBudgetIO( + self.dual_cache, + self._model_access_groups, + self._shared_windows, + _PROCESS_STARTED_AT, + clock=self._clock, + monotonic=time.monotonic, + ) + return await read_model_budget(io, entity, entity_id, team_id, model, budgets, key_budgets) + async def _is_entity_within_model_budget( self, entity_type: Litellm_EntityType, @@ -981,45 +1004,32 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): model_max_budget: Mapping[str, object] | None, model: str, exceeded_message: str, + key_model_max_budget: Mapping[str, object] | None = None, ) -> bool: - if not model_max_budget: - return True - resolved: Final = resolve_model_budget( - model=model, - model_max_budget=model_max_budget, - model_access_groups=self._model_access_groups, - team_id=team_id, + entities: Final[Mapping[Litellm_EntityType, ModelBudgetEntity]] = { + Litellm_EntityType.KEY: "Key", + Litellm_EntityType.USER: "User", + Litellm_EntityType.TEAM: "Team", + Litellm_EntityType.TEAM_MEMBER: "TeamMember", + Litellm_EntityType.END_USER: "EndUser", + } + reading: Final = await self.read_model_budget_raw( + entities[entity_type], + entity_id, + team_id, + model, + model_max_budget, + key_model_max_budget, ) - if resolved is None: - verbose_proxy_logger.debug("Model %s not found in %s model_max_budget", model, entity_type.value) - return True - - max_budget: Final = resolved.budget_config.max_budget - if max_budget is None or max_budget < 0: - return True - - local_spend: Final = await self._get_spend_for_model_budget( - entity_type=entity_type, - entity_id=entity_id, - model=model, - resolved=resolved, - ) - current_spend: Final = ( - local_spend - if local_spend >= max_budget - else await self._with_shared_spend( - entity_type=entity_type, entity_id=entity_id, resolved=resolved, local_spend=local_spend - ) - ) - if current_spend >= max_budget: + decision: Final = model_budget_data.check(reading.snapshot) + if decision.denied: + assert reading.snapshot.max_budget is not None raise litellm.BudgetExceededError( - message=( - exceeded_message - if resolved.budget_model in _budget_model_candidates(model) - else f"{exceeded_message}, budget_model={resolved.budget_model}" - ), - current_cost=current_spend, - max_budget=max_budget, + message=exceeded_message + if reading.named_model + else f"{exceeded_message}, budget_model={reading.budget_model}", + current_cost=decision.current_spend, + max_budget=reading.snapshot.max_budget, entity_type=entity_type.value, entity_id=entity_id, ) diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 4cc7bf4f761..2b621b25ade 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -817,3 +817,5 @@ class RouterCall(Generic[_S]): # candidate ids, pick first; pre_route: see _PreRouteStep; cache: see _CacheStep; secrets, build and model_info: # see those steps; first call: None. def next(self, outcome: object) -> _S: ... + +def request_auth_budget_exemption(facts: object) -> dict[str, object]: ... diff --git a/tests/unit/proxy/auth/budget_reference.py b/tests/unit/proxy/auth/budget_reference.py new file mode 100644 index 00000000000..a04a997ab46 --- /dev/null +++ b/tests/unit/proxy/auth/budget_reference.py @@ -0,0 +1,44 @@ +from collections.abc import Mapping +from typing import Final + +import litellm +from litellm.proxy._types import Litellm_EntityType +from litellm.proxy.hooks.model_max_budget_limiter import ( + _PROXY_VirtualKeyModelMaxBudgetLimiter, + entry_names_model, + resolve_model_budget, +) + + +async def model_allowed( + limiter: _PROXY_VirtualKeyModelMaxBudgetLimiter, + entity_type: Litellm_EntityType, + entity_id: str | None, + team_id: str | None, + budgets: Mapping[str, object] | None, + model: str, + message: str, +) -> bool: + if not budgets: + return True + resolved: Final = resolve_model_budget(model, budgets, limiter._model_access_groups, team_id) + if resolved is None: + return True + limit: Final = resolved.budget_config.max_budget + if limit is None or limit < 0: + return True + local: Final = await limiter._get_spend_for_model_budget(entity_type, entity_id, model, resolved) + current: Final = ( + local if local >= limit else await limiter._with_shared_spend(entity_type, entity_id, resolved, local) + ) + if current >= limit: + raise litellm.BudgetExceededError( + message=message + if entry_names_model(resolved.budget_model, model) + else f"{message}, budget_model={resolved.budget_model}", + current_cost=current, + max_budget=limit, + entity_type=entity_type.value, + entity_id=entity_id, + ) + return True diff --git a/tests/unit/proxy/auth/test_budget_pricing_adapter.py b/tests/unit/proxy/auth/test_budget_pricing_adapter.py new file mode 100644 index 00000000000..b0dacfdba2a --- /dev/null +++ b/tests/unit/proxy/auth/test_budget_pricing_adapter.py @@ -0,0 +1,121 @@ +from types import SimpleNamespace +from typing import Final +from unittest.mock import Mock + +import pytest + +import litellm +from litellm._auth.budget_pricing import ModelPriceFacts +from litellm.proxy.auth import auth_checks as reference +from litellm.proxy.auth import budget_pricing_adapter as native + + +def router(input_price: float | None, output_price: float | None, deployments: tuple[object, ...] = ()) -> Mock: + view: Final = Mock() + view.group.return_value = SimpleNamespace(input_cost_per_token=input_price, output_cost_per_token=output_price) + view.zero_cost.return_value = {} + view.deployments.return_value = deployments + return Mock(registry=Mock(return_value=view)) + + +@pytest.mark.parametrize(("input_price", "output_price"), ((None, 0), (0, None), (1, 0), (0, 1))) +def test_priced_or_undefined_group_does_not_read_deployments( + input_price: float | None, output_price: float | None +) -> None: + fake: Final = router(input_price, output_price) + facts: Final = native.read_model_price_facts("model", fake) + assert facts == ModelPriceFacts(group_found=True, input_price=input_price, output_price=output_price) + fake.registry().deployments.assert_not_called() + + +@pytest.mark.parametrize( + "route", + ( + "/v1/models", + "/model/info", + "/health/services", + "/chat/completions", + "/mcp", + "/mcp-rest/tools/list", + "/mcp/tools/call", + "/key/info", + ), +) +@pytest.mark.parametrize("body", (None, {}, {"method": "tools/list"}, {"method": "tools/call"}, {"method": 1})) +def test_route_and_mcp_decision_matches_production_reference(route: str, body: dict[str, object] | None) -> None: + assert native.request_skips_budget_checks( + route, None, None, request_body=body + ) == reference.request_skips_budget_checks(route, None, None, request_body=body) + assert native.route_skips_budget_checks(route) == reference.route_skips_budget_checks(route) + + +def test_missing_models_and_empty_model_list_preserve_exemption_behavior() -> None: + fake: Final = router(0, 0) + assert native.is_model_cost_zero(None, fake) is False + assert native.is_model_cost_zero("model", None) is False + assert native.is_model_cost_zero([], fake) is True + + +@pytest.mark.parametrize( + ("prices", "reads", "exempt"), (((0, 1, 0), 2, False), ((0, 0, 0), 3, True), ((1, 0, 0), 1, False)) +) +def test_raw_exemption_prices_only_native_selected_models_and_preserves_cache_writes( + monkeypatch: pytest.MonkeyPatch, prices: tuple[float, ...], reads: int, exempt: bool +) -> None: + from litellm._auth import budget_pricing + + old: Final = router(0, 0) + new: Final = router(0, 0) + monkeypatch.setitem( + litellm.model_cost, "budget-zero-fixture", {"input_cost_per_token": 0, "output_cost_per_token": 0} + ) + deployment: Final = { + "model_name": "model", + "litellm_params": {"model": "openai/fixture"}, + "model_info": {"id": "budget-zero-fixture"}, + } + old.registry().deployments.return_value = (deployment,) + new.registry().deployments.return_value = (deployment,) + old.registry().resolve.side_effect = lambda model: model + new.registry().resolve.side_effect = lambda model: model + old.registry().group.side_effect = tuple( + SimpleNamespace(input_cost_per_token=p, output_cost_per_token=0) for p in prices + ) + new.registry().group.side_effect = tuple( + SimpleNamespace(input_cost_per_token=p, output_cost_per_token=0) for p in prices + ) + old_cache: Final[dict[str, bool]] = {} + new_cache: Final[dict[str, bool]] = {} + old.registry().zero_cost.return_value = old_cache + new.registry().zero_cost.return_value = new_cache + models: Final = ["first", "second", "third"] + before: Final = reference.request_skips_budget_checks("/chat/completions", models, old) + facts: Final = native.read_budget_exemption_facts("/chat/completions", models, new) + assert budget_pricing.exemption(facts).exempt is before is exempt + assert len(facts.prices) == reads + assert old.registry().group.call_count == new.registry().group.call_count == reads + assert old_cache == new_cache + + +def test_raw_exemption_route_skip_does_not_touch_model_registry() -> None: + fake: Final = router(0, 0) + facts: Final = native.read_budget_exemption_facts("/v1/models", ["first", "second"], fake) + assert facts.prices == () + fake.registry.assert_not_called() + + +def test_raw_cached_price_does_not_read_model_group_or_write_again() -> None: + fake: Final = router(0, 0) + fake.registry().zero_cost.return_value = {"model": False} + facts: Final = native.read_model_cost_facts("model", fake) + assert facts.cached is False + fake.registry().group.assert_not_called() + + +def test_empty_model_list_still_reads_zero_cost_cache() -> None: + fake: Final = router(0, 0) + facts: Final = native.read_budget_exemption_facts("/chat/completions", [], fake) + assert facts.prices == () + assert facts.cache_read is True + fake.registry().zero_cost.assert_called_once_with() + fake.registry().group.assert_not_called() diff --git a/tests/unit/proxy/auth/test_common_budget_adapter.py b/tests/unit/proxy/auth/test_common_budget_adapter.py new file mode 100644 index 00000000000..01d3952e041 --- /dev/null +++ b/tests/unit/proxy/auth/test_common_budget_adapter.py @@ -0,0 +1,132 @@ +import asyncio +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from datetime import datetime, timezone +from math import inf +from typing import Final +from unittest.mock import AsyncMock, Mock + +import pytest + +import litellm +from litellm._auth.common_budgets import BudgetFacts +from litellm.models.budget import LiteLLM_BudgetTable +from litellm.proxy import proxy_server +from litellm.proxy._types import LiteLLM_EndUserTable, LiteLLM_TeamTable, UserAPIKeyAuth +from litellm.proxy.auth import auth_checks as reference +from litellm.proxy.auth import common_budget_adapter as native +from litellm.proxy.utils import ProxyLogging + + +@dataclass(frozen=True, slots=True) +class Observed: + error: tuple[object, ...] | None + alerts: tuple[object, ...] + throttle: float | None + + +async def observe(run: Callable[[], Awaitable[None]], logging: Mock, key: UserAPIKeyAuth) -> Observed: + try: + await run() + except litellm.BudgetExceededError as error: + failure: Final = (error.message, error.current_cost, error.max_budget, error.entity_type, error.entity_id) + await asyncio.sleep(0) + return Observed(failure, tuple(logging.budget_alerts.call_args_list), key.budget_throttle_pct) + await asyncio.sleep(0) + return Observed(None, tuple(logging.budget_alerts.call_args_list), key.budget_throttle_pct) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("limit", "spend"), ((None, 2), (1, 0), (1, 1), (1, 2), (0, 0), (-1, 0), (inf, 2))) +async def test_team_error_envelope_and_alert_projection_match_reference( + monkeypatch: pytest.MonkeyPatch, limit: float | None, spend: float +) -> None: + reader: Final = AsyncMock(return_value=spend) + monkeypatch.setattr(proxy_server, "get_current_spend", reader) + team: Final = LiteLLM_TeamTable(team_id="team", max_budget=limit, spend=0) + key: Final = UserAPIKeyAuth(token="hash", team_id="team") + old_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + new_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + before: Final = await observe(lambda: reference.team_max_budget_check(team, key, old_log), old_log, key) + after: Final = await observe(lambda: native.team_max_budget_check(team, key, new_log), new_log, key) + assert after == before + assert reader.await_count == (0 if limit is None else 2) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("limit", "spend"), ((None, 2), (1, 0), (1, 1), (1, 2), (0, 0), (-1, 0), (inf, 2))) +@pytest.mark.parametrize("throttle", (None, 0.5, 2)) +async def test_key_error_envelope_alert_and_throttle_mutation_match_reference( + monkeypatch: pytest.MonkeyPatch, limit: float | None, spend: float, throttle: float | None +) -> None: + monkeypatch.setattr(proxy_server, "get_current_spend", AsyncMock(return_value=spend)) + monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", throttle) + old: Final = UserAPIKeyAuth( + token="hash", key_alias="alias", max_budget=limit, tpm_limit=10, metadata={"throttle_on_budget_exceeded": True} + ) + new: Final = old.model_copy(deep=True) + old_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + new_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + before: Final = await observe(lambda: reference.virtual_key_max_budget_check(old, old_log), old_log, old) + after: Final = await observe(lambda: native.virtual_key_max_budget_check(new, new_log), new_log, new) + assert after == before + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ("/chat/completions", "/key/info")) +@pytest.mark.parametrize(("limit", "spend"), ((None, 2), (1, 1), (1, 2), (0, 0), (-1, 0))) +async def test_end_user_info_skip_and_strict_ceiling_match_reference( + monkeypatch: pytest.MonkeyPatch, route: str, limit: float | None, spend: float +) -> None: + reader: Final = AsyncMock(return_value=spend) + monkeypatch.setattr(proxy_server, "get_current_spend", reader) + user: Final = LiteLLM_EndUserTable( + user_id="end", blocked=False, litellm_budget_table=LiteLLM_BudgetTable(max_budget=limit) + ) + key: Final = UserAPIKeyAuth() + old_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + new_log: Final = Mock(spec=ProxyLogging, budget_alerts=AsyncMock()) + before: Final = await observe(lambda: reference.check_end_user_budget(user, route), old_log, key) + after: Final = await observe(lambda: native.check_end_user_budget(user, route), new_log, key) + assert after == before + + +@pytest.mark.asyncio +async def test_spend_io_failure_is_terminal_and_does_not_replay() -> None: + failure: Final = RuntimeError("counter unavailable") + reader: Final = AsyncMock(side_effect=failure) + with pytest.raises(RuntimeError, match="counter unavailable") as raised: + await native.read_budget(BudgetFacts("Key", 1), native.CounterRead("spend:key:hash", 0, 1), reader) + assert raised.value is failure + assert reader.await_count == 1 + + +@pytest.mark.parametrize( + ("member", "default", "increase", "expiry"), + ( + (1, 5, 2, datetime(2025, 1, 2)), + (None, 5, 2, datetime(2025, 1, 2)), + (1, 5, 2, datetime(2025, 1, 1, tzinfo=timezone.utc)), + (None, 0, 2, datetime(2025, 1, 2)), + (None, -1, 2, None), + (0, 5, None, None), + ), +) +def test_member_budget_selection_and_temporary_expiry_match_reference( + member: float | None, default: float | None, increase: float | None, expiry: datetime | None +) -> None: + now: Final = datetime(2025, 1, 1, tzinfo=timezone.utc) + row: Final = LiteLLM_BudgetTable(max_budget=member, temp_budget_increase=increase, temp_budget_expiry=expiry) + team: Final = LiteLLM_BudgetTable(max_budget=default) + assert native.resolve_effective_team_member_budget( + row, team, now + ) == reference.resolve_effective_team_member_budget(row, team, now) + + +@pytest.mark.parametrize( + "policy", (reference.resolve_effective_team_member_budget, native.resolve_effective_team_member_budget) +) +def test_naive_now_preserves_datetime_error(policy: Callable[..., float | None]) -> None: + row: Final = LiteLLM_BudgetTable(max_budget=1, temp_budget_increase=2, temp_budget_expiry=datetime(2025, 1, 2)) + with pytest.raises(TypeError, match="offset-naive"): + policy(row, None, datetime(2025, 1, 1)) diff --git a/tests/unit/proxy/auth/test_model_budget_adapter.py b/tests/unit/proxy/auth/test_model_budget_adapter.py new file mode 100644 index 00000000000..0393b018e5d --- /dev/null +++ b/tests/unit/proxy/auth/test_model_budget_adapter.py @@ -0,0 +1,138 @@ +from collections.abc import Awaitable, Mapping +from dataclasses import dataclass, field +from typing import Final +from unittest.mock import AsyncMock + +import pytest + +import litellm +from litellm._auth import model_budget_data +from litellm._auth.model_budget_data import ModelBudgetEntity +from litellm.caching.caching import DualCache +from litellm.proxy._types import Litellm_EntityType +from litellm.proxy.auth.model_budget_adapter import ( + ModelBudgetIO, + enforce_model_budget, + read_model_budget, + resolve_model_budget, +) +from litellm.proxy.hooks import model_max_budget_limiter as reference +from litellm.proxy.hooks.model_budget_shared_windows import SharedWindow + +from .budget_reference import model_allowed + + +@dataclass(frozen=True, slots=True) +class Windows: + value: SharedWindow | None = None + read: AsyncMock = field(default_factory=AsyncMock) + + async def open(self, spend_key: str, budget_duration: str, now: float) -> float | None: + raise AssertionError("request auth must not open a window") + + async def add(self, spend_key: str, budget_duration: str, started_at: float, spend: float) -> None: + raise AssertionError("request auth must not write a counter") + + +def groups(model: str, team_id: str | None) -> frozenset[str]: + return frozenset(("group",)) + + +@pytest.mark.parametrize( + ("model", "budgets"), + ( + ( + "provider/model", + {"model": {"max_budget": 1, "budget_duration": "1h"}, "*": {"max_budget": 2, "budget_duration": "1h"}}, + ), + ("model", {"model": {"budget_duration": "bad"}, "*": {"max_budget": 2, "budget_duration": "1h"}}), + ("model", {"model": {}, "group": {"max_budget": 2, "budget_duration": "1h"}}), + ( + "model", + {"m*": {"max_budget": 2, "budget_duration": "1h"}, "group": {"max_budget": 1, "budget_duration": "1h"}}, + ), + ( + "mødél", + {"m*é*": {"max_budget": 2, "budget_duration": "1h"}, "*": {"max_budget": 1, "budget_duration": "1h"}}, + ), + ( + "model", + {"model": {"max_budget": None, "budget_duration": "1h"}, "*": {"max_budget": 1, "budget_duration": "1h"}}, + ), + ), +) +def test_resolution_matches_production_reference(model: str, budgets: Mapping[str, object]) -> None: + io: Final = ModelBudgetIO(DualCache(), groups, Windows(), 0) + old: Final = reference.resolve_model_budget(model, budgets, groups) + result: Final = resolve_model_budget(model, budgets, io, None) + assert result.read_groups is False + assert (result.entry.name if result.entry else None) == (old.budget_model if old else None) + assert (result.entry.max_budget if result.entry else None) == (old.budget_config.max_budget if old else None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("entity", "scope"), + ( + ("Key", Litellm_EntityType.KEY), + ("User", Litellm_EntityType.USER), + ("Team", Litellm_EntityType.TEAM), + ("TeamMember", Litellm_EntityType.TEAM_MEMBER), + ("EndUser", Litellm_EntityType.END_USER), + ), +) +@pytest.mark.parametrize( + ("local", "shared", "now", "limit"), + ((0, 2, 100, 1), (1, 0, 100, 1), (0, 2, 3700, 1), (0, 0, 100, 0), (0, 2, 100, -1), (0, 2, 100, None)), +) +async def test_counter_and_shared_window_match_production_reference( + entity: ModelBudgetEntity, scope: Litellm_EntityType, local: float, shared: float, now: float, limit: float | None +) -> None: + cache: Final = DualCache() + key: Final = reference.model_budget_spend_cache_key(scope, "id", "model", "1h") + cache.set_cache(key, local) + windows: Final = Windows(read=AsyncMock(return_value=SharedWindow(100, shared))) + budgets: Final = {"model": {"max_budget": limit, "budget_duration": "1h"}} + old: Final = reference._PROXY_VirtualKeyModelMaxBudgetLimiter( + cache, clock=lambda: now, model_access_groups=groups, shared_windows=windows + ) + io: Final = ModelBudgetIO(cache, groups, windows, 0, clock=lambda: now, monotonic=lambda: 0) + before: Final = await observed_model_result(model_allowed(old, scope, "id", None, budgets, "model", "exceeded")) + after: Final = await observed_model_result( + enforce_model_budget(io, entity, "id", None, "model", budgets, "exceeded") + ) + assert after == before + + +async def observed_model_result(call: Awaitable[bool]) -> tuple[object, ...]: + try: + return (await call,) + except litellm.BudgetExceededError as error: + return (error.message, error.current_cost, error.max_budget, error.entity_type, error.entity_id) + + +@pytest.mark.asyncio +async def test_legacy_counter_is_added_to_the_current_counter_before_shared_read() -> None: + cache: Final = DualCache() + cache.set_cache(reference.model_budget_spend_cache_key(Litellm_EntityType.KEY, "id", "model", "1h"), 0.4) + cache.set_cache(reference.model_budget_spend_cache_key(Litellm_EntityType.KEY, "id", "provider/model", "1h"), 0.6) + windows: Final = Windows(read=AsyncMock(side_effect=AssertionError("already over local limit"))) + io: Final = ModelBudgetIO(cache, groups, windows, 0, monotonic=lambda: 1) + reading: Final = await read_model_budget( + io, "Key", "id", None, "provider/model", {"model": {"max_budget": 1, "budget_duration": "1h"}} + ) + assert reading.snapshot.local_spend == 0.4 + assert reading.snapshot.legacy_spend == 0.6 + assert model_budget_data.check(reading.snapshot).denied is True + windows.read.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_shared_read_failure_is_terminal_and_not_replayed() -> None: + failure: Final = RuntimeError("database unavailable") + windows: Final = Windows(read=AsyncMock(side_effect=failure)) + io: Final = ModelBudgetIO(DualCache(), groups, windows, 0, monotonic=lambda: 0) + with pytest.raises(RuntimeError, match="database unavailable") as raised: + 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