scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

Select model budget scope reads and wire enterprise route policy

Browse source
This commit is contained in:
Devin AI 2026-11-09 00:53:25 +00:00
parent ca42107a42
commit a7d198f673
11 changed files with 270 additions and 45 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

28
litellm-rust/crates/python-bridge/src/lib.rs
View file

@ -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",

19
litellm-rust/crates/python-bridge/src/request_auth.rs
View file

@ -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<Py<PyAny>> {
pub(crate) fn request_auth_model_projection(
py: Python<'_>,
facts: &Bound<'_, PyAny>,
reads: &Bound<'_, PyAny>,
) -> PyResult<Py<PyAny>> {
let facts: core::model_access::Facts = from_py(facts)?;
let reads: Vec<core::model_projection::Reading> = 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<Py<PyAny>> {
pub(crate) fn request_auth_route_enable(
py: Python<'_>,
facts: &Bound<'_, PyAny>,
) -> PyResult<Py<PyAny>> {
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)
}

102
litellm-rust/crates/request-auth-core/src/common_flow.rs
View file

@ -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<bool>,
budget_present: Option<bool>,
mapping: Option<bool>,
team_present: Option<bool>,
},
FallbackAllowed {
router: bool,
disabled: Option<bool>,
},
BudgetSkip {
skip: bool,
exempt: Option<bool>,
@ -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<bool>,
budget: Option<bool>,
mapping: Option<bool>,
team: Option<bool>,
) -> 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)]

4
litellm-rust/crates/request-auth-core/src/lib.rs
View file

@ -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;

1
litellm-rust/crates/request-auth-core/src/model_budgets.rs
View file

@ -275,6 +275,7 @@ mod tests {
team: None,
embedded_team: None,
assigned_models: None,
assigned_scope: None,
member: None,
project: None,
team_present: false,

25
litellm/_auth/common_flow.py
View file

@ -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)))

9
litellm/proxy/auth/route_checks.py
View file

@ -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(

12
litellm/proxy/auth/route_enable_adapter.py
View file

@ -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)

51
litellm/proxy/auth/user_api_key_auth.py
View file

@ -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:

38
tests/unit/_auth/test_common_flow.py
View file

@ -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) == (

26
tests/unit/proxy/auth/test_route_enable_adapter.py
View file

@ -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"),
(