diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 3941a1329e9..03db8cc0aea 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -58,20 +58,18 @@ mod _native { 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_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_role, request_auth_route_role_api, - request_auth_route_role_expand, request_auth_route_role_inside, - request_auth_route_role_jwt, request_auth_route_role_listed, - request_auth_scheme_credential, request_auth_select_api_key, - request_auth_select_websocket_api_key, request_auth_spend_counter_policy, - request_auth_start, request_auth_team_aliases, request_auth_team_override, - request_auth_unpriced_model, - request_auth_model_projection, - request_auth_route_enable, + 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, + request_auth_route_role, request_auth_route_role_api, request_auth_route_role_expand, + request_auth_route_role_inside, request_auth_route_role_jwt, + request_auth_route_role_listed, request_auth_scheme_credential, + request_auth_select_api_key, request_auth_select_websocket_api_key, + request_auth_spend_counter_policy, request_auth_start, request_auth_team_aliases, + request_auth_team_override, request_auth_unpriced_model, }; #[pymodule_export] use crate::router::{RouterCall, RouterCoreHandle}; @@ -181,6 +179,7 @@ mod tests { "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", @@ -193,6 +192,7 @@ mod tests { "request_auth_reservation", "request_auth_resolved_models", "request_auth_route_access", + "request_auth_route_enable", "request_auth_route_role", "request_auth_route_role_api", "request_auth_route_role_expand", diff --git a/litellm-rust/crates/python-bridge/src/request_auth.rs b/litellm-rust/crates/python-bridge/src/request_auth.rs index f47e8f546da..b1604b240e5 100644 --- a/litellm-rust/crates/python-bridge/src/request_auth.rs +++ b/litellm-rust/crates/python-bridge/src/request_auth.rs @@ -902,16 +902,27 @@ pub(crate) fn request_auth_org_admin_raw( } #[pyfunction] -pub(crate) fn request_auth_model_projection(py: Python<'_>, facts: &Bound<'_, PyAny>, reads: &Bound<'_, PyAny>) -> PyResult> { +pub(crate) fn request_auth_model_projection( + py: Python<'_>, + facts: &Bound<'_, PyAny>, + reads: &Bound<'_, PyAny>, +) -> PyResult> { let facts: core::model_access::Facts = from_py(facts)?; let reads: Vec = from_py(reads)?; - let decision = py.detach(move || core::model_projection::decide(&facts, &reads)).map_err(map_error)?; + let decision = py + .detach(move || core::model_projection::decide(&facts, &reads)) + .map_err(map_error)?; to_py(py, &decision) } #[pyfunction] -pub(crate) fn request_auth_route_enable(py: Python<'_>, facts: &Bound<'_, PyAny>) -> PyResult> { +pub(crate) fn request_auth_route_enable( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { let facts: core::route_enable::Facts = from_py(facts)?; - let decision = py.detach(move || core::route_enable::decide(&facts)).map_err(map_error)?; + let decision = py + .detach(move || core::route_enable::decide(&facts)) + .map_err(map_error)?; to_py(py, &decision) } diff --git a/litellm-rust/crates/request-auth-core/src/common_flow.rs b/litellm-rust/crates/request-auth-core/src/common_flow.rs index 6e4d8900c45..2f6a4e9873d 100644 --- a/litellm-rust/crates/request-auth-core/src/common_flow.rs +++ b/litellm-rust/crates/request-auth-core/src/common_flow.rs @@ -1,8 +1,26 @@ use serde::{Deserialize, Serialize}; +#[derive(Deserialize, strum::VariantArray)] +pub enum ModelBudgetScope { + Key, + User, + Member, +} + #[derive(Deserialize)] #[serde(tag = "kind")] pub enum ReadFacts { + ModelBudgets { + scope: ModelBudgetScope, + entity_present: Option, + budget_present: Option, + mapping: Option, + team_present: Option, + }, + FallbackAllowed { + router: bool, + disabled: Option, + }, BudgetSkip { skip: bool, exempt: Option, @@ -23,6 +41,14 @@ pub enum ReadFacts { #[derive(Debug, PartialEq, Eq, Serialize, strum::VariantArray)] pub enum ReadPlan { + CheckModelBudgets, + ReadModelBudgetPresence, + ReadModelBudgetShape, + ReadModelBudgetToken, + ReadModelBudgetUser, + ReadModelBudgetTeam, + AllowFallback, + ReadFallbackDisabled, ReadExemption, SkipBudgets, CheckBudgets, @@ -37,6 +63,22 @@ pub enum ReadPlan { pub fn read(facts: ReadFacts) -> ReadPlan { match facts { + ReadFacts::ModelBudgets { + scope, + entity_present, + budget_present, + mapping, + team_present, + } => model_budgets(scope, entity_present, budget_present, mapping, team_present), + ReadFacts::FallbackAllowed { + router: true, + disabled: Some(false), + } => ReadPlan::AllowFallback, + ReadFacts::FallbackAllowed { + router: true, + disabled: None, + } => ReadPlan::ReadFallbackDisabled, + ReadFacts::FallbackAllowed { .. } => ReadPlan::Continue, ReadFacts::BudgetSkip { skip: true, .. } | ReadFacts::BudgetSkip { exempt: Some(true), .. @@ -80,6 +122,41 @@ pub fn read(facts: ReadFacts) -> ReadPlan { } } +fn model_budgets( + scope: ModelBudgetScope, + entity: Option, + budget: Option, + mapping: Option, + team: Option, +) -> ReadPlan { + use ReadPlan as P; + let stages = match scope { + ModelBudgetScope::Key => vec![ + (budget, P::ReadModelBudgetPresence), + (entity, P::ReadModelBudgetToken), + ], + ModelBudgetScope::User => vec![ + (entity, P::ReadModelBudgetUser), + (mapping, P::ReadModelBudgetShape), + (budget, P::ReadModelBudgetPresence), + ], + ModelBudgetScope::Member => vec![ + (entity, P::ReadModelBudgetUser), + (team, P::ReadModelBudgetTeam), + (mapping, P::ReadModelBudgetShape), + (budget, P::ReadModelBudgetPresence), + ], + }; + for (present, missing) in stages { + match present { + None => return missing, + Some(false) => return P::Continue, + Some(true) => {} + } + } + P::CheckModelBudgets +} + #[derive(Deserialize)] pub struct Facts { pub model_present: bool, @@ -192,6 +269,31 @@ mod tests { use super::*; use rstest::rstest; + #[rstest] + #[case::available(true, true, true, true, ReadPlan::CheckModelBudgets)] + #[case::missing_identity(false, true, true, true, ReadPlan::Continue)] + #[case::empty_budget(true, false, true, true, ReadPlan::Continue)] + #[case::malformed_member_budget(true, true, false, true, ReadPlan::Continue)] + #[case::no_team(true, true, true, false, ReadPlan::Continue)] + fn model_budget_availability( + #[case] entity_present: bool, + #[case] budget_present: bool, + #[case] mapping: bool, + #[case] team_present: bool, + #[case] expected: ReadPlan, + ) { + assert_eq!( + read(ReadFacts::ModelBudgets { + scope: ModelBudgetScope::Member, + entity_present: Some(entity_present), + budget_present: Some(budget_present), + mapping: Some(mapping), + team_present: Some(team_present) + }), + expected + ); + } + #[rstest] #[case::explicit_skip(ReadFacts::BudgetSkip {skip:true, exempt:None}, ReadPlan::SkipBudgets)] #[case::exempt(ReadFacts::BudgetSkip {skip:false, exempt:Some(true)}, ReadPlan::SkipBudgets)] diff --git a/litellm-rust/crates/request-auth-core/src/lib.rs b/litellm-rust/crates/request-auth-core/src/lib.rs index e4bb53e6d76..c52896cfa10 100644 --- a/litellm-rust/crates/request-auth-core/src/lib.rs +++ b/litellm-rust/crates/request-auth-core/src/lib.rs @@ -22,13 +22,12 @@ pub mod managed; pub mod managed_models; pub mod member_models; pub mod model_access; -pub mod model_projection; -pub mod route_enable; pub mod model_attribution; pub mod model_budget_data; mod model_budgets; pub mod model_groups; pub mod model_helpers; +pub mod model_projection; pub mod org_routes; pub mod owner_teams; pub mod patterns; @@ -41,6 +40,7 @@ pub mod protocol; pub mod reservation; pub mod resolved_models; pub mod route_access; +pub mod route_enable; pub mod route_roles; pub mod spend_counter_policy; pub mod team_aliases; diff --git a/litellm-rust/crates/request-auth-core/src/model_budgets.rs b/litellm-rust/crates/request-auth-core/src/model_budgets.rs index c2484af6475..709a13c8a43 100644 --- a/litellm-rust/crates/request-auth-core/src/model_budgets.rs +++ b/litellm-rust/crates/request-auth-core/src/model_budgets.rs @@ -275,6 +275,7 @@ mod tests { team: None, embedded_team: None, assigned_models: None, + assigned_scope: None, member: None, project: None, team_present: false, diff --git a/litellm/_auth/common_flow.py b/litellm/_auth/common_flow.py index 238718dfce2..208f68cd253 100644 --- a/litellm/_auth/common_flow.py +++ b/litellm/_auth/common_flow.py @@ -32,6 +32,21 @@ class MemberFallback: resource_bound: bool | None = None +@dataclass(frozen=True, slots=True) +class ModelBudgets: + scope: Literal["Key", "User", "Member"] + entity_present: bool | None = None + budget_present: bool | None = None + mapping: bool | None = None + team_present: bool | None = None + + +@dataclass(frozen=True, slots=True) +class FallbackAllowed: + router: bool + disabled: bool | None = None + + ReadPlan: TypeAlias = Literal[ "ReadExemption", "SkipBudgets", @@ -42,13 +57,21 @@ ReadPlan: TypeAlias = Literal[ "ReadContentType", "ReadResourceBinding", "ReadFallbackSettings", + "CheckModelBudgets", + "ReadModelBudgetPresence", + "ReadModelBudgetShape", + "ReadModelBudgetToken", + "ReadModelBudgetUser", + "ReadModelBudgetTeam", + "AllowFallback", + "ReadFallbackDisabled", "Continue", ] _READ: Final = TypeAdapter(ReadPlan) @native_boundary -def read(facts: BudgetSkip | Pricing | Unpriced | MemberFallback) -> ReadPlan: +def read(facts: BudgetSkip | Pricing | Unpriced | MemberFallback | ModelBudgets | FallbackAllowed) -> ReadPlan: return _READ.validate_python(_native.request_auth_common_read(encode_facts(facts))) diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index 48abd8d7252..a5842da07c1 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -114,14 +114,9 @@ _TOKEN_COUNTING_ACTION: Final = "countTokens" class RouteChecks: @staticmethod def check_route_enabled(route: str) -> None: - try: - from litellm_enterprise.proxy.auth.route_checks import EnterpriseRouteChecks + from litellm.proxy.auth.route_enable_adapter import check_route_enabled - EnterpriseRouteChecks.should_call_route(route=route) - except HTTPException as e: - raise e - except Exception: - pass + check_route_enabled(route) @staticmethod def should_call_route( diff --git a/litellm/proxy/auth/route_enable_adapter.py b/litellm/proxy/auth/route_enable_adapter.py index abdeb50001f..9e5e413fca9 100644 --- a/litellm/proxy/auth/route_enable_adapter.py +++ b/litellm/proxy/auth/route_enable_adapter.py @@ -4,9 +4,9 @@ from dataclasses import dataclass, replace from typing import Final, assert_never from fastapi import HTTPException -from pydantic import BaseModel +from pydantic import TypeAdapter -from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError, native_boundary +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError from litellm._auth.route_enable import RouteEnableDecision, RouteEnableFacts, RouteFeature, RouteSetting, route_enable from litellm.proxy._types import CommonProxyErrors from litellm.proxy.auth.route_policy_adapter import route_context @@ -14,8 +14,7 @@ from litellm.proxy.auth.route_policy_adapter import route_context _KEYS: Final = {"Admin": "DISABLE_ADMIN_ENDPOINTS", "Llm": "DISABLE_LLM_API_ENDPOINTS"} -class _Exemptions(BaseModel, frozen=True): - LLM_API_EXEMPT_ROUTES: tuple[str, ...] +_EXEMPTIONS: Final = TypeAdapter(tuple[str, ...]) def read_enterprise() -> object: @@ -39,9 +38,10 @@ def read_secret(feature: RouteFeature) -> bool: return get_secret_bool(_KEYS[feature]) is True -@native_boundary def read_exemptions(enterprise: object) -> tuple[str, ...]: - return _Exemptions.model_validate(enterprise, from_attributes=True).LLM_API_EXEMPT_ROUTES + return _EXEMPTIONS.validate_python( + getattr(enterprise, "LLM_API_EXEMPT_ROUTES") # noqa: B009 # optional class attributes are validated at this boundary + ) @dataclass(frozen=True, slots=True) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index b4a32046f1c..bbc0e51e6a5 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -14,7 +14,7 @@ from dataclasses import replace from datetime import datetime, timezone from enum import Enum from types import MappingProxyType -from typing import Final, NamedTuple, Protocol, Union, assert_never, cast +from typing import Final, Literal, NamedTuple, Protocol, Union, assert_never, cast import fastapi from fastapi import HTTPException, Request, WebSocket, status @@ -58,7 +58,7 @@ from litellm._auth.decision import ( WithinModelBudgets, ) from litellm._auth.driver import drive, drive_model_budgets, drive_team_member_model_budgets -from litellm._auth.failure import NativeAuthFailure +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError from litellm._auth.fallback_access import FallbackAuthorizationFacts from litellm._auth.fallback_candidate import FallbackCandidateFacts from litellm._auth.key_model_access import KeyAndFallbackModelAccessFacts, key_and_fallback_model_access @@ -574,6 +574,38 @@ def _model_budget_read( ) +def _model_budget_scope_enabled(scope: Literal["Key", "User", "Member"], token: UserAPIKeyAuth, budget: object) -> bool: + facts = common_flow.ModelBudgets(scope) # rebind-ok: native-selected raw field read stages + while True: + match common_flow.read(facts): + case "ReadModelBudgetPresence": + facts = replace(facts, budget_present=bool(budget)) + case "ReadModelBudgetShape": + facts = replace(facts, mapping=isinstance(budget, Mapping)) + case "ReadModelBudgetToken": + facts = replace(facts, entity_present=token.token is not None) + case "ReadModelBudgetUser": + facts = replace(facts, entity_present=token.user_id is not None) + case "ReadModelBudgetTeam": + facts = replace(facts, team_present=token.team_id is not None) + case "CheckModelBudgets": + return True + case "Continue": + return False + case _: + raise NativeAuthRuntimeError("Unexpected model-budget scope read") + + +def _model_budget_fallback_enabled(router_present: bool, data: Mapping[str, object]) -> bool: + initial: Final = common_flow.FallbackAllowed(router_present) + facts: Final = ( + replace(initial, disabled=fallbacks_disabled(data)) + if common_flow.read(initial) == "ReadFallbackDisabled" + else initial + ) + return common_flow.read(facts) == "AllowFallback" + + class _ModelBudgetEffects: def __init__( self, @@ -619,11 +651,11 @@ class _ModelBudgetEffects: ) ), rewritable_model=requested if isinstance(requested, str) else None, - key_has_model_budgets=bool(self._valid_token.model_max_budget) and self._valid_token.token is not None, + key_has_model_budgets=_model_budget_scope_enabled( + "Key", self._valid_token, self._valid_token.model_max_budget + ), key_budget_fallbacks=MappingProxyType({model: tuple(chain) for model, chain in budget_fallbacks.items()}), - user_has_model_budgets=self._valid_token.user_id is not None - and isinstance(user_model_max_budget, Mapping) - and bool(user_model_max_budget), + user_has_model_budgets=_model_budget_scope_enabled("User", self._valid_token, user_model_max_budget), ) async def check_key_model_budget(self, model: str) -> ModelBudgetRead: @@ -738,11 +770,8 @@ class _TeamMemberModelBudgetEffects: return TeamMemberModelBudgetScope( models=self._models, rewritable_model=requested if isinstance(requested, str) else None, - has_model_budgets=self._valid_token.user_id is not None - and self._valid_token.team_id is not None - and isinstance(budget, Mapping) - and bool(budget), - can_fall_back=self._llm_router is not None and not fallbacks_disabled(self._request_data), + has_model_budgets=_model_budget_scope_enabled("Member", self._valid_token, budget), + can_fall_back=_model_budget_fallback_enabled(self._llm_router is not None, self._request_data), ) async def check_team_member_model_budget(self, model: str) -> ModelBudgetRead: diff --git a/tests/unit/_auth/test_common_flow.py b/tests/unit/_auth/test_common_flow.py index 6f2598197e1..365a10e654c 100644 --- a/tests/unit/_auth/test_common_flow.py +++ b/tests/unit/_auth/test_common_flow.py @@ -46,6 +46,44 @@ def test_common_model_plan_matches_reference(facts: common_flow.Facts) -> None: ) +@pytest.mark.parametrize("entity", (False, True)) +@pytest.mark.parametrize("budget", (False, True)) +@pytest.mark.parametrize("mapping", (False, True)) +@pytest.mark.parametrize("team", (False, True)) +def test_model_budget_scope_matches_reference(entity: bool, budget: bool, mapping: bool, team: bool) -> None: + expected: Final = "CheckModelBudgets" if entity and budget and mapping and team else "Continue" + assert common_flow.read(common_flow.ModelBudgets("Member", entity, budget, mapping, team)) == expected + + +@pytest.mark.parametrize("router", (False, True)) +@pytest.mark.parametrize("disabled", (False, True)) +def test_fallback_eligibility_matches_reference(router: bool, disabled: bool) -> None: + expected: Final = "AllowFallback" if router and not disabled else "Continue" + assert common_flow.read(common_flow.FallbackAllowed(router, disabled)) == expected + + +@pytest.mark.parametrize( + "facts, expected", + ( + (common_flow.ModelBudgets("Key"), "ReadModelBudgetPresence"), + (common_flow.ModelBudgets("Key", budget_present=False), "Continue"), + (common_flow.ModelBudgets("Key", budget_present=True), "ReadModelBudgetToken"), + (common_flow.ModelBudgets("User"), "ReadModelBudgetUser"), + (common_flow.ModelBudgets("User", entity_present=False), "Continue"), + (common_flow.ModelBudgets("User", entity_present=True), "ReadModelBudgetShape"), + (common_flow.ModelBudgets("User", entity_present=True, mapping=False), "Continue"), + (common_flow.ModelBudgets("Member", entity_present=True), "ReadModelBudgetTeam"), + (common_flow.ModelBudgets("Member", entity_present=True, team_present=False), "Continue"), + (common_flow.FallbackAllowed(False), "Continue"), + (common_flow.FallbackAllowed(True), "ReadFallbackDisabled"), + ), +) +def test_scope_reads_preserve_reference_short_circuit( + facts: common_flow.ModelBudgets | common_flow.FallbackAllowed, expected: common_flow.ReadPlan +) -> None: + assert common_flow.read(facts) == expected + + @pytest.mark.parametrize("facts", _BUDGETS) def test_budget_scopes_preserve_reference_priority(facts: common_flow.BudgetFacts) -> None: assert common_flow.budgets(facts) == ( diff --git a/tests/unit/proxy/auth/test_route_enable_adapter.py b/tests/unit/proxy/auth/test_route_enable_adapter.py index 56f1bc96ee5..7f356c5bd78 100644 --- a/tests/unit/proxy/auth/test_route_enable_adapter.py +++ b/tests/unit/proxy/auth/test_route_enable_adapter.py @@ -10,6 +10,32 @@ from litellm.proxy._types import CommonProxyErrors from litellm.proxy.auth.route_enable_adapter import RouteEnableReaders, check_route_enabled +@pytest.mark.parametrize("route", ("/models", "/chat/completions")) +def test_enterprise_class_exemptions_preserve_enabled_route_policy(route: str) -> None: + class EnterpriseRoutes: + LLM_API_EXEMPT_ROUTES: Final = ("/models",) + + readers: Final = RouteEnableReaders( + enterprise=lambda: EnterpriseRoutes, + setting=lambda _: RouteSetting(True, True), + secret=lambda _: True, + ) + if route == "/models": + assert check_route_enabled(route, readers=readers) is None + else: + with pytest.raises(HTTPException) as error: + check_route_enabled(route, readers=readers) + assert (error.value.status_code, error.value.detail) == ( + 403, + "LLM API routes are disabled for this instance.", + ) + + +def test_missing_optional_class_exemptions_preserve_ignored_read_failure() -> None: + readers: Final = RouteEnableReaders(enterprise=lambda: object()) + assert check_route_enabled("/chat/completions", readers=readers) is None + + @pytest.mark.parametrize( ("route", "expected"), (