diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 3c82335495a..a00057e0013 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -52,10 +52,12 @@ mod _native { request_auth_litellm_header_token, request_auth_managed, request_auth_member_limit, 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_org_admin, request_auth_org_route, - request_auth_owner_teams, request_auth_personal_budget, request_auth_personal_models, - request_auth_reservation, request_auth_route_access, request_auth_scheme_credential, - request_auth_select_api_key, request_auth_select_websocket_api_key, request_auth_start, + request_auth_model_groups, request_auth_model_helpers, request_auth_org_admin, + request_auth_org_route, request_auth_owner_teams, request_auth_personal_budget, + request_auth_personal_model_request, request_auth_personal_models, + request_auth_personal_owner, request_auth_reservation, request_auth_route_access, + request_auth_scheme_credential, request_auth_select_api_key, + request_auth_select_websocket_api_key, request_auth_start, request_auth_team_override, request_auth_unpriced_model, }; #[pymodule_export] diff --git a/litellm-rust/crates/python-bridge/src/request_auth.rs b/litellm-rust/crates/python-bridge/src/request_auth.rs index 456c52d2e5c..9fe009e647b 100644 --- a/litellm-rust/crates/python-bridge/src/request_auth.rs +++ b/litellm-rust/crates/python-bridge/src/request_auth.rs @@ -574,3 +574,51 @@ pub(crate) fn request_auth_context_percentage( let facts: core::budget_context::NumericFacts = from_py(facts)?; Ok(py.detach(move || core::budget_context::percentage(&facts))) } + +#[pyfunction] +pub(crate) fn request_auth_team_override( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::team_override::Facts = from_py(facts)?; + let decision = py + .detach(move || core::team_override::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_model_helpers( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::model_helpers::Facts = from_py(facts)?; + let decision = py + .detach(move || core::model_helpers::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_personal_owner( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::personal_owner::Facts = from_py(facts)?; + let decision = py + .detach(move || core::personal_owner::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_personal_model_request( + py: Python<'_>, + facts: &Bound<'_, PyAny>, +) -> PyResult> { + let facts: core::personal_model_request::Facts = from_py(facts)?; + let decision = py + .detach(move || core::personal_model_request::decide(&facts)) + .map_err(map_error)?; + to_py(py, &decision) +} diff --git a/litellm-rust/crates/request-auth-core/src/lib.rs b/litellm-rust/crates/request-auth-core/src/lib.rs index f83fdb66b31..c3320a1a1d8 100644 --- a/litellm-rust/crates/request-auth-core/src/lib.rs +++ b/litellm-rust/crates/request-auth-core/src/lib.rs @@ -22,14 +22,18 @@ pub mod model_attribution; pub mod model_budget_data; mod model_budgets; pub mod model_groups; +pub mod model_helpers; pub mod org_routes; pub mod owner_teams; pub mod patterns; pub mod personal_budget; +pub mod personal_model_request; pub mod personal_models; +pub mod personal_owner; pub mod protocol; pub mod reservation; pub mod route_access; +pub mod team_override; pub mod unpriced_models; pub use error::Error; diff --git a/litellm-rust/crates/request-auth-core/src/model_helpers.rs b/litellm-rust/crates/request-auth-core/src/model_helpers.rs new file mode 100644 index 00000000000..3e5b1a75c2e --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/model_helpers.rs @@ -0,0 +1,141 @@ +use crate::{Error, model_access::Provider, patterns}; +use serde::{Deserialize, Serialize}; +use std::collections::BTreeSet; + +#[derive(Clone, Copy, Deserialize, PartialEq, Eq, strum::VariantArray)] +pub enum Kind { + Group, + Helper, +} + +#[derive(Clone, Deserialize)] +pub struct Facts { + pub kind: Kind, + pub model: String, + pub grants: Option>, + pub team_id: Option, + pub team_aliases: Vec, + pub router_present: bool, + pub groups: Option>, + pub names: Option>, + pub provider_loaded: bool, + pub provider: Option, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Complete { allowed: bool }, + ReadGroups, + ReadNames, + ReadProvider, +} + +pub fn decide(facts: &Facts) -> Result { + let done = |allowed| Ok(Decision::Complete { allowed }); + if facts.kind == Kind::Group + && facts + .grants + .as_ref() + .is_none_or(|grants| grants.contains(&facts.model)) + { + return done(true); + } + let grants = facts.grants.as_ref().ok_or(Error::InvalidAnswer)?; + if facts.router_present && facts.groups.is_none() { + return Ok(Decision::ReadGroups); + } + let groups = facts.groups.as_deref().unwrap_or_default(); + if facts.kind == Kind::Group { + return done(grants.iter().any(|grant| groups.contains(grant))); + } + let expanded = if facts.router_present + && facts.team_id.is_some() + && grants.iter().any(|grant| grant == "all-team-models") + { + let Some(names) = &facts.names else { + return Ok(Decision::ReadNames); + }; + let mut seen = BTreeSet::new(); + let values: Vec<_> = grants + .iter() + .filter(|grant| *grant != "all-team-models") + .chain(names) + .filter(|grant| seen.insert(*grant)) + .cloned() + .collect(); + if values.is_empty() { + grants.clone() + } else { + values + } + } else { + grants.clone() + }; + if expanded.iter().any(|grant| groups.contains(grant)) + || facts.team_aliases.contains(&facts.model) + { + return done(true); + } + let filtered: Vec<_> = expanded + .iter() + .filter(|grant| !groups.contains(grant)) + .collect(); + let wildcards: Vec<_> = filtered + .iter() + .filter(|grant| grant.contains('*')) + .collect(); + for pattern in &wildcards { + if patterns::model_pattern(&facts.model, pattern)? { + return done(true); + } + } + if !wildcards.is_empty() { + if !facts.provider_loaded { + return Ok(Decision::ReadProvider); + } + if let Some(provider) = &facts.provider + && (provider.stripped_model != facts.model || !facts.model.contains('/')) + { + for pattern in wildcards { + if patterns::model_pattern( + &format!("{}/{}", provider.provider, provider.stripped_model), + pattern, + )? { + return done(true); + } + } + } + } + done( + (filtered.is_empty() && expanded.is_empty()) + || filtered.iter().any(|grant| { + matches!(grant.as_str(), "*" | "all-proxy-models") || **grant == facts.model + }), + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::group(Kind::Group, Decision::Complete { allowed: true })] + #[case::helper(Kind::Helper, Decision::ReadGroups)] + fn direct_match_read_order(#[case] kind: Kind, #[case] expected: Decision) { + let facts = Facts { + kind, + model: "model".into(), + grants: Some(vec!["model".into()]), + team_id: None, + team_aliases: vec![], + router_present: true, + groups: None, + names: None, + provider_loaded: false, + provider: None, + }; + assert_eq!(decide(&facts).unwrap(), expected); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/personal_model_request.rs b/litellm-rust/crates/request-auth-core/src/personal_model_request.rs new file mode 100644 index 00000000000..9db1e2a03cb --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/personal_model_request.rs @@ -0,0 +1,86 @@ +use crate::{Error, owner_teams, personal_models}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Deserialize)] +pub struct Loaded { + pub present: Vec, + pub scopes: Vec, +} + +#[derive(Clone, Deserialize)] +pub struct Facts { + pub requested: Vec, + pub owner_present: bool, + pub teams: Option, + pub models: Vec, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Allow, + Deny { model: String }, + ReadTeams, + ReadModel { index: usize, teams: Vec }, +} + +pub fn decide(facts: &Facts) -> Result { + let teams = if facts.owner_present { + let Some(teams) = &facts.teams else { + return Ok(Decision::ReadTeams); + }; + match owner_teams::decide(teams)? { + owner_teams::Decision::Complete { selected, .. } => selected, + _ => return Err(Error::InvalidAnswer), + } + } else { + vec![] + }; + for (index, model) in facts.requested.iter().enumerate() { + if !facts.owner_present { + return Ok(Decision::Deny { + model: model.clone(), + }); + } + let Some(loaded) = facts.models.get(index) else { + return Ok(Decision::ReadModel { index, teams }); + }; + match personal_models::decide(personal_models::Request::Loaded { + present: loaded.present.clone(), + scopes: loaded.scopes.clone(), + })? { + personal_models::Decision::Complete { allowed: true } => (), + personal_models::Decision::Complete { allowed: false } => { + return Ok(Decision::Deny { + model: model.clone(), + }); + } + _ => return Err(Error::InvalidAnswer), + } + } + Ok(Decision::Allow) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::empty(false, vec![], Decision::Allow)] + #[case::owner_lookup(true, vec![], Decision::ReadTeams)] + #[case::absent(false, vec!["model".into()], Decision::Deny { model: "model".into() })] + fn owner_and_empty_requests( + #[case] owner_present: bool, + #[case] requested: Vec, + #[case] expected: Decision, + ) { + let facts = Facts { + requested, + owner_present, + teams: None, + models: vec![], + }; + assert_eq!(decide(&facts).unwrap(), expected); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/personal_owner.rs b/litellm-rust/crates/request-auth-core/src/personal_owner.rs new file mode 100644 index 00000000000..bc821d2d509 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/personal_owner.rs @@ -0,0 +1,87 @@ +use crate::{Error, owner_teams}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Deserialize)] +pub struct Facts { + pub user_present: bool, + pub database_present: bool, + pub cache_present: bool, + pub logging_present: bool, + pub owner_loaded: bool, + pub owner_present: bool, + pub owner_failed: bool, + pub teams: Option, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + ReadOwner, + ReadTeams, + Complete { + selected: Vec, + withheld: Vec, + unreadable: bool, + }, +} + +pub fn decide(facts: &Facts) -> Result { + let empty = || { + Ok(Decision::Complete { + selected: vec![], + withheld: vec![], + unreadable: false, + }) + }; + if !facts.user_present + || !facts.database_present + || !facts.cache_present + || !facts.logging_present + { + return empty(); + } + if !facts.owner_loaded { + return Ok(Decision::ReadOwner); + } + if facts.owner_failed || !facts.owner_present { + return empty(); + } + let Some(teams) = &facts.teams else { + return Ok(Decision::ReadTeams); + }; + match owner_teams::decide(teams)? { + owner_teams::Decision::Complete { + selected, + withheld, + unreadable, + } => Ok(Decision::Complete { + selected, + withheld, + unreadable, + }), + _ => Err(Error::InvalidAnswer), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::no_database(false, Decision::Complete { selected: vec![], withheld: vec![], unreadable: false })] + #[case::lookup(true, Decision::ReadOwner)] + fn prerequisites(#[case] database_present: bool, #[case] expected: Decision) { + let facts = Facts { + user_present: true, + database_present, + cache_present: true, + logging_present: true, + owner_loaded: false, + owner_present: false, + owner_failed: false, + teams: None, + }; + assert_eq!(decide(&facts).unwrap(), expected); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/team_override.rs b/litellm-rust/crates/request-auth-core/src/team_override.rs new file mode 100644 index 00000000000..ca69868c4cc --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/team_override.rs @@ -0,0 +1,68 @@ +use crate::{ + Error, + fallback_access::{self, GroupedScope}, + model_access::{self, Decision as ModelDecision, Facts as ModelFacts, ObjectType}, +}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Default, Deserialize)] +pub struct Facts { + pub team: Option, + pub assigned_models: Option>, + pub assigned_scope: Option, +} + +#[derive(Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + Allow, + Deny { denial: ModelDecision }, + ReadTeam, + ReadTeamGroups, + ReadAssignedGroups, + ReadAssignedScope, +} + +pub fn decide(facts: &Facts) -> Result { + let Some(team) = &facts.team else { + return Ok(Decision::ReadTeam); + }; + let denial = fallback_access::grouped(team.clone(), true)?; + match denial { + ModelDecision::Allow => return Ok(Decision::Allow), + ModelDecision::ReadGroups => return Ok(Decision::ReadTeamGroups), + ModelDecision::Deny { + object_type: ObjectType::Team, + .. + } => (), + _ => return Ok(Decision::Deny { denial }), + } + let Some(assigned) = &facts.assigned_models else { + return Ok(Decision::ReadAssignedGroups); + }; + if assigned.is_empty() { + return Ok(Decision::Deny { denial }); + } + let Some(scope) = &facts.assigned_scope else { + return Ok(Decision::ReadAssignedScope); + }; + match model_access::check(scope)? { + ModelDecision::Allow => Ok(Decision::Allow), + ModelDecision::Deny { .. } => Ok(Decision::Deny { denial }), + error => Ok(Decision::Deny { denial: error }), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + fn team_is_read_before_override_groups() { + assert!(matches!( + decide(&Facts::default()).unwrap(), + Decision::ReadTeam + )); + } +} diff --git a/litellm/_auth/model_helpers.py b/litellm/_auth/model_helpers.py new file mode 100644 index 00000000000..8ba323a1e42 --- /dev/null +++ b/litellm/_auth/model_helpers.py @@ -0,0 +1,36 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.model_access import ModelProvider +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class ModelHelperFacts: + kind: Literal["Group", "Helper"] + model: str + grants: tuple[str, ...] | None + team_id: str | None = None + team_aliases: tuple[str, ...] = () + router_present: bool = False + groups: tuple[str, ...] | None = None + names: tuple[str, ...] | None = None + provider_loaded: bool = False + provider: ModelProvider | None = None + + +@dataclass(frozen=True, slots=True) +class ModelHelperDecision: + kind: Literal["Complete", "ReadGroups", "ReadNames", "ReadProvider"] + allowed: bool = False + + +_DECISION: Final = TypeAdapter(ModelHelperDecision) + + +@native_boundary +def model_helper(facts: ModelHelperFacts) -> ModelHelperDecision: + return _DECISION.validate_python(_native.request_auth_model_helpers(asdict(facts))) diff --git a/litellm/_auth/personal_model_request.py b/litellm/_auth/personal_model_request.py new file mode 100644 index 00000000000..aa885980426 --- /dev/null +++ b/litellm/_auth/personal_model_request.py @@ -0,0 +1,33 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.owner_teams import OwnerTeamFacts +from litellm._auth.personal_models import LoadedTeamModelFacts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class PersonalModelRequestFacts: + requested: tuple[str, ...] + owner_present: bool + teams: OwnerTeamFacts | None = None + models: tuple[LoadedTeamModelFacts, ...] = () + + +@dataclass(frozen=True, slots=True) +class PersonalModelRequestDecision: + kind: Literal["Allow", "Deny", "ReadTeams", "ReadModel"] + model: str = "" + index: int = 0 + teams: tuple[int, ...] = () + + +_DECISION: Final = TypeAdapter(PersonalModelRequestDecision) + + +@native_boundary +def personal_model_request(facts: PersonalModelRequestFacts) -> PersonalModelRequestDecision: + return _DECISION.validate_python(_native.request_auth_personal_model_request(asdict(facts))) diff --git a/litellm/_auth/personal_owner.py b/litellm/_auth/personal_owner.py new file mode 100644 index 00000000000..7ec4c9a16c8 --- /dev/null +++ b/litellm/_auth/personal_owner.py @@ -0,0 +1,36 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.owner_teams import OwnerTeamFacts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class PersonalOwnerFacts: + user_present: bool + database_present: bool + cache_present: bool + logging_present: bool + owner_loaded: bool = False + owner_present: bool = False + owner_failed: bool = False + teams: OwnerTeamFacts | None = None + + +@dataclass(frozen=True, slots=True) +class PersonalOwnerDecision: + kind: Literal["ReadOwner", "ReadTeams", "Complete"] + selected: tuple[int, ...] = () + withheld: tuple[int, ...] = () + unreadable: bool = False + + +_DECISION: Final = TypeAdapter(PersonalOwnerDecision) + + +@native_boundary +def personal_owner(facts: PersonalOwnerFacts) -> PersonalOwnerDecision: + return _DECISION.validate_python(_native.request_auth_personal_owner(asdict(facts))) diff --git a/litellm/_auth/team_override.py b/litellm/_auth/team_override.py new file mode 100644 index 00000000000..504f5cc83b7 --- /dev/null +++ b/litellm/_auth/team_override.py @@ -0,0 +1,30 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.model_access import ModelAccessDecision, ModelAccessFacts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class TeamOverrideFacts: + team: GroupedModelScope | None = None + assigned_models: tuple[str, ...] | None = None + assigned_scope: ModelAccessFacts | None = None + + +@dataclass(frozen=True, slots=True) +class TeamOverrideDecision: + kind: Literal["Allow", "Deny", "ReadTeam", "ReadTeamGroups", "ReadAssignedGroups", "ReadAssignedScope"] + denial: ModelAccessDecision | None = None + + +_DECISION: Final = TypeAdapter(TeamOverrideDecision) + + +@native_boundary +def team_override(facts: TeamOverrideFacts) -> TeamOverrideDecision: + return _DECISION.validate_python(_native.request_auth_team_override(asdict(facts))) diff --git a/litellm/proxy/auth/model_helper_adapter.py b/litellm/proxy/auth/model_helper_adapter.py new file mode 100644 index 00000000000..274efc648d6 --- /dev/null +++ b/litellm/proxy/auth/model_helper_adapter.py @@ -0,0 +1,50 @@ +from collections.abc import Mapping, Sequence +from dataclasses import replace +from typing import Final + +from litellm._auth.failure import NativeAuthRuntimeError +from litellm._auth.model_helpers import ModelHelperFacts, model_helper +from litellm.proxy.auth.model_access_adapter import ( + _provider, # pyright: ignore[reportPrivateUsage] # shared raw provider projection preserves terminal failures +) +from litellm.router_api import RegistryReader + + +def read_model_helper_facts(facts: ModelHelperFacts, router: RegistryReader | None) -> ModelHelperFacts: + decision: Final = model_helper(facts) + match decision.kind: + case "Complete": + return facts + case "ReadGroups": + if router is None: + raise NativeAuthRuntimeError("model helper registry read without router") + groups: Final = tuple(router.registry().access_groups(model_name=facts.model, team_id=facts.team_id)) + return read_model_helper_facts(replace(facts, groups=groups), router) + case "ReadNames": + if router is None: + raise NativeAuthRuntimeError("model helper names read without router") + return read_model_helper_facts(replace(facts, names=tuple(router.registry().names())), router) + case "ReadProvider": + return read_model_helper_facts( + replace(facts, provider_loaded=True, provider=_provider(facts.model)), router + ) + + +def check_model_access_helper( + model: str, + llm_router: RegistryReader | None, + models: Sequence[str], + team_model_aliases: Mapping[str, str] | None = None, + team_id: str | None = None, +) -> bool: + facts: Final = ModelHelperFacts( + "Helper", model, tuple(models), team_id, tuple(team_model_aliases or ()), bool(llm_router) + ) + return model_helper(read_model_helper_facts(facts, llm_router)).allowed + + +def model_in_access_group(model: str, team_models: Sequence[str] | None, llm_router: RegistryReader | None) -> bool: + facts: Final = ModelHelperFacts( + "Group", model, tuple(team_models) if team_models is not None else None, router_present=bool(llm_router) + ) + return model_helper(read_model_helper_facts(facts, llm_router)).allowed diff --git a/litellm/proxy/auth/personal_model_request_adapter.py b/litellm/proxy/auth/personal_model_request_adapter.py new file mode 100644 index 00000000000..f96da4834b1 --- /dev/null +++ b/litellm/proxy/auth/personal_model_request_adapter.py @@ -0,0 +1,93 @@ +from collections.abc import Sequence +from dataclasses import dataclass, replace +from typing import Final, Literal + +from litellm._auth.failure import NativeAuthRuntimeError +from litellm._auth.model_access import ModelAccessDecision +from litellm._auth.personal_model_request import PersonalModelRequestFacts, personal_model_request +from litellm.proxy._types import LiteLLM_TeamTableCachedObj, LiteLLM_UserTable, UserAPIKeyAuth +from litellm.proxy.auth.model_access_adapter import raise_model_denial +from litellm.proxy.auth.owner_team_adapter import read_owner_team_facts +from litellm.proxy.auth.personal_model_adapter import read_loaded_team_model_facts +from litellm.proxy.auth.personal_owner_adapter import selected_team +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.router_api import RegistryReader + + +@dataclass(frozen=True, slots=True) +class _Source: + owner: LiteLLM_UserTable | None + token: UserAPIKeyAuth + router: RegistryReader | None + prisma: PrismaClient | None + cache: UserApiKeyCache + logging: ProxyLogging + teams: tuple[LiteLLM_TeamTableCachedObj | None, ...] = () + + +async def read_personal_model_request_facts( + model: str | Sequence[str], + user_object: LiteLLM_UserTable | None, + valid_token: UserAPIKeyAuth, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> PersonalModelRequestFacts: + requested: Final = (model,) if isinstance(model, str) else tuple(model) + return await _read( + PersonalModelRequestFacts(requested, user_object is not None), + _Source(user_object, valid_token, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj), + ) + + +async def can_personal_key_call_model_via_teams( + model: str | Sequence[str], + user_object: LiteLLM_UserTable | None, + valid_token: UserAPIKeyAuth, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> Literal[True]: + facts: Final = await read_personal_model_request_facts( + model, user_object, valid_token, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj + ) + decision: Final = personal_model_request(facts) + if decision.kind == "Allow": + return True + if decision.kind != "Deny": + raise NativeAuthRuntimeError("incomplete personal model request facts") + return raise_model_denial( + ModelAccessDecision("PersonalTeam", model=decision.model), + model, + user_id=valid_token.user_id, + teams=user_object.teams if user_object is not None else None, + ) + + +async def _read(facts: PersonalModelRequestFacts, source: _Source) -> PersonalModelRequestFacts: + decision: Final = personal_model_request(facts) + match decision.kind: + case "Allow" | "Deny": + return facts + case "ReadTeams": + if source.owner is None: + raise NativeAuthRuntimeError("personal model team read without owner") + reading: Final = await read_owner_team_facts( + source.owner, source.token, source.prisma, source.cache, source.logging, purpose="Models" + ) + return await _read(replace(facts, teams=reading.facts), replace(source, teams=reading.teams)) + case "ReadModel": + teams: Final = tuple(selected_team(source.teams, index) for index in decision.teams) + model: Final = await read_loaded_team_model_facts( + facts.requested[decision.index], + teams, + source.token, + source.router, + source.prisma, + source.cache, + source.logging, + ) + return await _read(replace(facts, models=(*facts.models, model)), source) diff --git a/litellm/proxy/auth/personal_owner_adapter.py b/litellm/proxy/auth/personal_owner_adapter.py new file mode 100644 index 00000000000..20add242af3 --- /dev/null +++ b/litellm/proxy/auth/personal_owner_adapter.py @@ -0,0 +1,127 @@ +from __future__ import annotations + +from dataclasses import dataclass, replace +from typing import TYPE_CHECKING, Final + +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError +from litellm._auth.personal_owner import PersonalOwnerFacts, personal_owner +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import LiteLLM_TeamTableCachedObj, LiteLLM_UserTable, UserAPIKeyAuth +from litellm.proxy.auth.owner_team_adapter import read_owner_team_facts +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging + +if TYPE_CHECKING: + from litellm.proxy.auth.auth_checks import PersonalKeyOwnerTeams + + +@dataclass(frozen=True, slots=True) +class PersonalOwnerReading: + facts: PersonalOwnerFacts + teams: tuple[LiteLLM_TeamTableCachedObj | None, ...] = () + + +@dataclass(frozen=True, slots=True) +class _Source: + token: UserAPIKeyAuth + prisma: PrismaClient | None + cache: UserApiKeyCache | None + logging: ProxyLogging | None + check_db_only: bool + owner: LiteLLM_UserTable | None = None + + +async def read_personal_owner_facts( + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache | None, + proxy_logging_obj: ProxyLogging | None, + *, + check_db_only: bool = False, +) -> PersonalOwnerReading: + facts: Final = PersonalOwnerFacts( + valid_token.user_id is not None, + prisma_client is not None, + user_api_key_cache is not None, + proxy_logging_obj is not None, + ) + return await _read( + PersonalOwnerReading(facts), + _Source(valid_token, prisma_client, user_api_key_cache, proxy_logging_obj, check_db_only), + ) + + +async def personal_key_owner_teams_with_status( + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache | None, + proxy_logging_obj: ProxyLogging | None, + *, + check_db_only: bool = False, +) -> PersonalKeyOwnerTeams: + from litellm.proxy.auth.auth_checks import PersonalKeyOwnerTeams + + reading: Final = await read_personal_owner_facts( + valid_token, prisma_client, user_api_key_cache, proxy_logging_obj, check_db_only=check_db_only + ) + decision: Final = personal_owner(reading.facts) + if decision.kind != "Complete": + raise NativeAuthRuntimeError("incomplete personal owner team facts") + return PersonalKeyOwnerTeams( + teams=tuple(selected_team(reading.teams, index) for index in decision.selected), + withheld=tuple(selected_team(reading.teams, index) for index in decision.withheld), + unreadable=decision.unreadable, + ) + + +def selected_team(teams: tuple[LiteLLM_TeamTableCachedObj | None, ...], index: int) -> LiteLLM_TeamTableCachedObj: + team: Final = teams[index] + if team is None: + raise NativeAuthRuntimeError("missing selected owner team row") + return team + + +async def _read(reading: PersonalOwnerReading, source: _Source) -> PersonalOwnerReading: + from litellm.proxy.auth.auth_checks import get_user_object + + decision: Final = personal_owner(reading.facts) + if decision.kind == "Complete": + return reading + if source.token.user_id is None or source.prisma is None or source.cache is None or source.logging is None: + raise NativeAuthRuntimeError("owner read requested without required dependencies") + match decision.kind: + case "ReadOwner": + try: + owner: Final = await get_user_object( + user_id=source.token.user_id, + prisma_client=source.prisma, + user_api_key_cache=source.cache, + user_id_upsert=False, + parent_otel_span=source.token.parent_otel_span, + proxy_logging_obj=source.logging, + check_db_only=source.check_db_only, + ) + except NativeAuthFailure: + raise + except Exception as error: # noqa: BLE001 # preserve legacy owner lookup failure as raw data + verbose_proxy_logger.warning( + "Personal key team model access: owner lookup failed, granting nothing: %s", error + ) + return PersonalOwnerReading(replace(reading.facts, owner_loaded=True, owner_failed=True)) + return await _read( + PersonalOwnerReading(replace(reading.facts, owner_loaded=True, owner_present=owner is not None)), + replace(source, owner=owner), + ) + case "ReadTeams": + if source.owner is None: + raise NativeAuthRuntimeError("personal owner team read without owner row") + teams: Final = await read_owner_team_facts( + source.owner, + source.token, + source.prisma, + source.cache, + source.logging, + purpose="Grants", + check_db_only=source.check_db_only, + ) + return PersonalOwnerReading(replace(reading.facts, teams=teams.facts), teams.teams) diff --git a/litellm/proxy/auth/team_override_adapter.py b/litellm/proxy/auth/team_override_adapter.py new file mode 100644 index 00000000000..bc3fc6ffed6 --- /dev/null +++ b/litellm/proxy/auth/team_override_adapter.py @@ -0,0 +1,102 @@ +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, replace +from typing import Final, Literal + +from litellm._auth.failure import NativeAuthRuntimeError +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.team_override import TeamOverrideFacts, team_override +from litellm.proxy._types import LiteLLM_TeamTable, UserAPIKeyAuth +from litellm.proxy.auth.model_access_adapter import object_model_grants, project_model_facts, raise_model_denial +from litellm.proxy.auth.model_group_adapter import get_authorized_resources_from_key_access_groups +from litellm.repositories.prisma_protocols import DatabaseClient +from litellm.router_api import RegistryReader + + +@dataclass(frozen=True, slots=True) +class _Source: + model: str | Sequence[str] + token: UserAPIKeyAuth + team: LiteLLM_TeamTable | None + router: RegistryReader | None + team_aliases: Mapping[str, str] | None + key_aliases: Mapping[str, str] | None + prisma: DatabaseClient | None + + +async def read_team_override_facts( + model: str | Sequence[str], + valid_token: UserAPIKeyAuth, + team_object: LiteLLM_TeamTable | None, + llm_router: RegistryReader | None, + team_model_aliases: Mapping[str, str] | None = None, + key_model_aliases: Mapping[str, str] | None = None, + prisma_client: DatabaseClient | None = None, +) -> TeamOverrideFacts: + return await _read( + TeamOverrideFacts(), + _Source(model, valid_token, team_object, llm_router, team_model_aliases, key_model_aliases, prisma_client), + ) + + +async def can_team_call_model_with_override( + model: str | Sequence[str], + valid_token: UserAPIKeyAuth, + team_object: LiteLLM_TeamTable | None, + llm_router: RegistryReader | None, + team_model_aliases: Mapping[str, str] | None = None, + key_model_aliases: Mapping[str, str] | None = None, + prisma_client: DatabaseClient | None = None, +) -> Literal[True]: + facts: Final = await read_team_override_facts( + model, valid_token, team_object, llm_router, team_model_aliases, key_model_aliases, prisma_client + ) + decision: Final = team_override(facts) + if decision.kind == "Allow": + return True + if decision.kind != "Deny" or decision.denial is None: + raise NativeAuthRuntimeError("incomplete team override authorization") + return raise_model_denial(decision.denial, model) + + +async def _read(facts: TeamOverrideFacts, source: _Source) -> TeamOverrideFacts: + from litellm.proxy.auth.auth_checks import get_models_from_access_groups + + decision: Final = team_override(facts) + match decision.kind: + case "Allow" | "Deny": + return facts + case "ReadTeam": + scope: Final = project_model_facts( + source.model, + source.router, + object_model_grants(source.team) if source.team is not None else (), + source.team_aliases, + source.team.team_id if source.team is not None else None, + source.key_aliases, + "team", + ) + groups: Final = tuple(source.team.access_group_ids or ()) if source.team is not None else () + return await _read(replace(facts, team=GroupedModelScope(scope, groups)), source) + case "ReadTeamGroups": + if facts.team is None: + raise NativeAuthRuntimeError("missing team override scope") + models: Final = await get_models_from_access_groups(facts.team.group_ids, prisma_client=source.prisma) + return await _read(replace(facts, team=replace(facts.team, group_models=tuple(models))), source) + case "ReadAssignedGroups": + assigned: Final = await get_authorized_resources_from_key_access_groups( + source.token, source.team, "access_model_names" + ) + return await _read(replace(facts, assigned_models=tuple(assigned)), source) + case "ReadAssignedScope": + if facts.assigned_models is None: + raise NativeAuthRuntimeError("missing assigned model override grants") + assigned_scope: Final = project_model_facts( + source.model, + source.router, + facts.assigned_models, + source.token.team_model_aliases, + source.token.team_id, + source.key_aliases, + "key", + ) + return await _read(replace(facts, assigned_scope=assigned_scope), source) diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 6c7dc896012..fd3cc7ef39f 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -60,6 +60,10 @@ def request_auth_fallback_candidate(facts: object) -> object: ... def request_auth_personal_budget(facts: object) -> object: ... def request_auth_personal_models(facts: object) -> object: ... def request_auth_owner_teams(facts: object) -> object: ... +def request_auth_team_override(facts: object) -> object: ... +def request_auth_model_helpers(facts: object) -> object: ... +def request_auth_personal_owner(facts: object) -> object: ... +def request_auth_personal_model_request(facts: object) -> object: ... @final class NativeClickHouseSpendConfig: diff --git a/tests/unit/_auth/test_model_helpers.py b/tests/unit/_auth/test_model_helpers.py new file mode 100644 index 00000000000..c038776f218 --- /dev/null +++ b/tests/unit/_auth/test_model_helpers.py @@ -0,0 +1,105 @@ +import re +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.model_access import ModelProvider +from litellm._auth.model_helpers import ModelHelperDecision as Decision +from litellm._auth.model_helpers import ModelHelperFacts as Facts +from litellm._auth.model_helpers import model_helper + + +def reference(facts: Facts) -> Decision: + if facts.kind == "Group" and (facts.grants is None or facts.model in facts.grants): + return Decision("Complete", True) + if facts.router_present and facts.groups is None: + return Decision("ReadGroups") + groups: Final = facts.groups or () + grants: Final = facts.grants or () + if facts.kind == "Group": + return Decision("Complete", any(grant in groups for grant in grants)) + expand: Final = "all-team-models" in grants and facts.team_id is not None and facts.router_present + if expand and facts.names is None: + return Decision("ReadNames") + expanded: Final = ( + ( + tuple( + dict.fromkeys((*tuple(grant for grant in grants if grant != "all-team-models"), *(facts.names or ()))) + ) + or grants + ) + if expand + else grants + ) + if any(grant in groups for grant in expanded) or facts.model in facts.team_aliases: + return Decision("Complete", True) + filtered: Final = tuple(grant for grant in expanded if grant not in groups) + patterns: Final = tuple(grant.replace("*", ".*") + "$" for grant in filtered if "*" in grant) + if any(re.match(pattern, facts.model) is not None for pattern in patterns): + return Decision("Complete", True) + if patterns: + if not facts.provider_loaded: + return Decision("ReadProvider") + provider: Final = facts.provider + if ( + provider is not None + and (provider.stripped_model != facts.model or "/" not in facts.model) + and any( + re.match(pattern, f"{provider.provider}/{provider.stripped_model}") is not None for pattern in patterns + ) + ): + return Decision("Complete", True) + return Decision( + "Complete", + not filtered and not expanded or "*" in filtered or "all-proxy-models" in filtered or facts.model in filtered, + ) + + +_GROUP: Final = Facts("Group", "model", ("model",), router_present=True) +_HELPER: Final = replace(_GROUP, kind="Helper") +_READ: Final = replace(_HELPER, groups=()) + + +@pytest.mark.parametrize("policy", (reference, model_helper)) +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (_GROUP, Decision("Complete", True)), + (replace(_GROUP, grants=None), Decision("Complete", True)), + (replace(_GROUP, grants=()), Decision("ReadGroups")), + (replace(_GROUP, grants=(), groups=()), Decision("Complete")), + (replace(_GROUP, grants=("group",), groups=("group",)), Decision("Complete", True)), + (_HELPER, Decision("ReadGroups")), + (_READ, Decision("Complete", True)), + (replace(_READ, grants=()), Decision("Complete", True)), + (replace(_READ, grants=("other",)), Decision("Complete")), + (replace(_READ, grants=("other",), team_aliases=("model",)), Decision("Complete", True)), + (replace(_READ, grants=("group",), groups=("group",)), Decision("Complete", True)), + (replace(_READ, grants=("all-team-models",), team_id="team"), Decision("ReadNames")), + (replace(_READ, grants=("all-team-models",), team_id="team", names=()), Decision("Complete")), + (replace(_READ, grants=("all-team-models",), team_id="team", names=("model",)), Decision("Complete", True)), + (replace(_READ, grants=("mod*",)), Decision("Complete", True)), + (replace(_READ, grants=("provider/*",)), Decision("ReadProvider")), + (replace(_READ, grants=("provider/*",), provider_loaded=True), Decision("Complete")), + ( + replace(_READ, grants=("provider/*",), provider_loaded=True, provider=ModelProvider("model", "provider")), + Decision("Complete", True), + ), + ( + replace( + _READ, + model="unknown/model", + grants=("provider/*",), + provider_loaded=True, + provider=ModelProvider("unknown/model", "provider"), + ), + Decision("Complete"), + ), + ), +) +def test_model_helper_read_order_and_grants( + policy: Callable[[Facts], Decision], facts: Facts, expected: Decision +) -> None: + assert policy(facts) == expected diff --git a/tests/unit/_auth/test_personal_model_request.py b/tests/unit/_auth/test_personal_model_request.py new file mode 100644 index 00000000000..40b2edf752f --- /dev/null +++ b/tests/unit/_auth/test_personal_model_request.py @@ -0,0 +1,57 @@ +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.owner_teams import OwnerTeamFacts, OwnerTeamLoad +from litellm._auth.personal_model_request import PersonalModelRequestDecision as Decision +from litellm._auth.personal_model_request import PersonalModelRequestFacts as Facts +from litellm._auth.personal_model_request import personal_model_request +from litellm._auth.personal_models import LoadedTeamModelFacts +from tests.unit._auth.test_owner_teams import reference as teams_reference +from tests.unit._auth.test_personal_models import ALLOW, INITIAL, reference as scope_reference + + +def reference(facts: Facts) -> Decision: + if facts.owner_present and facts.teams is None: + return Decision("ReadTeams") + teams: Final = teams_reference(facts.teams).selected if facts.teams is not None else () + for index, model in enumerate(facts.requested): + if not facts.owner_present: + return Decision("Deny", model=model) + if index >= len(facts.models): + return Decision("ReadModel", index=index, teams=teams) + if not any(scope_reference(scope).allowed for scope in facts.models[index].scopes): + return Decision("Deny", model=model) + return Decision("Allow") + + +_BASE: Final = Facts(("first", "second"), True) +_TEAMS: Final = replace( + _BASE, teams=OwnerTeamFacts("Models", ("team",), "ui", loaded=True, loads=(OwnerTeamLoad(True),)) +) +_ONE: Final = replace(_TEAMS, models=(LoadedTeamModelFacts((True,), (ALLOW,)),)) + + +@pytest.mark.parametrize("policy", (reference, personal_model_request)) +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (_BASE, Decision("ReadTeams")), + (replace(_BASE, requested=()), Decision("ReadTeams")), + (replace(_BASE, owner_present=False), Decision("Deny", model="first")), + (replace(_BASE, owner_present=False, requested=()), Decision("Allow")), + (_TEAMS, Decision("ReadModel", teams=(0,))), + (_ONE, Decision("ReadModel", index=1, teams=(0,))), + (replace(_ONE, models=(*_ONE.models, LoadedTeamModelFacts((True,), (ALLOW,)))), Decision("Allow")), + ( + replace(_TEAMS, models=(LoadedTeamModelFacts((True,), (replace(INITIAL, io_failed=True),)),)), + Decision("Deny", model="first"), + ), + ), +) +def test_personal_requests_read_all_teams_but_stop_at_first_model_denial( + policy: Callable[[Facts], Decision], facts: Facts, expected: Decision +) -> None: + assert policy(facts) == expected diff --git a/tests/unit/_auth/test_personal_owner.py b/tests/unit/_auth/test_personal_owner.py new file mode 100644 index 00000000000..02e75764fcd --- /dev/null +++ b/tests/unit/_auth/test_personal_owner.py @@ -0,0 +1,61 @@ +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.owner_teams import OwnerTeamFacts, OwnerTeamLoad +from litellm._auth.personal_owner import PersonalOwnerDecision as Decision +from litellm._auth.personal_owner import PersonalOwnerFacts as Facts +from litellm._auth.personal_owner import personal_owner +from tests.unit._auth.test_owner_teams import reference as select_teams + + +def reference(facts: Facts) -> Decision: + if not all((facts.user_present, facts.database_present, facts.cache_present, facts.logging_present)): + return Decision("Complete") + if not facts.owner_loaded: + return Decision("ReadOwner") + if facts.owner_failed or not facts.owner_present: + return Decision("Complete") + if facts.teams is None: + return Decision("ReadTeams") + teams: Final = select_teams(facts.teams) + return Decision("Complete", teams.selected, teams.withheld, teams.unreadable) + + +_BASE: Final = Facts(True, True, True, True) +_OWNER: Final = replace(_BASE, owner_loaded=True, owner_present=True) + + +@pytest.mark.parametrize("policy", (reference, personal_owner)) +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (_BASE, Decision("ReadOwner")), + (replace(_BASE, user_present=False), Decision("Complete")), + (replace(_BASE, database_present=False), Decision("Complete")), + (replace(_BASE, cache_present=False), Decision("Complete")), + (replace(_BASE, logging_present=False), Decision("Complete")), + (replace(_BASE, owner_loaded=True), Decision("Complete")), + (replace(_OWNER, owner_failed=True), Decision("Complete")), + (_OWNER, Decision("ReadTeams")), + ( + replace( + _OWNER, + teams=OwnerTeamFacts( + "Grants", + (), + "ui", + loaded=True, + loads=(OwnerTeamLoad(True, blocked=True), OwnerTeamLoad(False, failed=True)), + ), + ), + Decision("Complete", withheld=(0,), unreadable=True), + ), + ), +) +def test_owner_lookup_and_team_status_order( + policy: Callable[[Facts], Decision], facts: Facts, expected: Decision +) -> None: + assert policy(facts) == expected diff --git a/tests/unit/_auth/test_team_override.py b/tests/unit/_auth/test_team_override.py new file mode 100644 index 00000000000..8e956b54e3e --- /dev/null +++ b/tests/unit/_auth/test_team_override.py @@ -0,0 +1,74 @@ +from collections.abc import Callable +from dataclasses import replace +from typing import Final + +import pytest + +from litellm._auth.failure import NativeAuthFailure +from litellm._auth.fallback_access import GroupedModelScope +from litellm._auth.model_access import ModelAccessFacts +from litellm._auth.team_override import TeamOverrideDecision as Decision +from litellm._auth.team_override import TeamOverrideFacts as Facts +from litellm._auth.team_override import team_override +from tests.unit._auth.support.fallback_reference import grouped +from tests.unit._auth.support.model_reference import model_access + + +def reference(facts: Facts) -> Decision: + if facts.team is None: + return Decision("ReadTeam") + denied: Final = grouped(facts.team, merge=True) + if denied.kind == "Allow": + return Decision("Allow") + if denied.kind == "ReadGroups": + return Decision("ReadTeamGroups") + if denied.kind != "Deny" or denied.object_type != "team": + return Decision("Deny", denied) + if facts.assigned_models is None: + return Decision("ReadAssignedGroups") + if not facts.assigned_models: + return Decision("Deny", denied) + if facts.assigned_scope is None: + return Decision("ReadAssignedScope") + assigned: Final = model_access(facts.assigned_scope) + if assigned.kind == "Allow": + return Decision("Allow") + return Decision("Deny", denied if assigned.kind == "Deny" else assigned) + + +_MODEL: Final = ModelAccessFacts(("model",), ("other",), object_type="team") +_TEAM: Final = Facts(GroupedModelScope(_MODEL)) +_ASSIGNED: Final = replace(_TEAM, assigned_models=("model",)) + + +@pytest.mark.parametrize("policy", (reference, team_override)) +@pytest.mark.parametrize( + "facts", + ( + Facts(), + _TEAM, + Facts(GroupedModelScope(replace(_MODEL, grants=()))), + Facts(GroupedModelScope(_MODEL, ("team-group",))), + Facts(GroupedModelScope(_MODEL, ("team-group",), ("model",))), + Facts(GroupedModelScope(_MODEL, ("team-group",), ("no-access",))), + Facts(GroupedModelScope(replace(_MODEL, object_type="key"))), + Facts(GroupedModelScope(replace(_MODEL, fallback_depth=10))), + replace(_TEAM, assigned_models=()), + _ASSIGNED, + replace(_ASSIGNED, assigned_scope=replace(_MODEL, object_type="key", grants=("model",))), + replace(_ASSIGNED, assigned_scope=replace(_MODEL, object_type="key")), + replace(_ASSIGNED, assigned_scope=replace(_MODEL, object_type="key", fallback_depth=10)), + ), +) +def test_team_override_reads_groups_only_after_team_denial(policy: Callable[[Facts], Decision], facts: Facts) -> None: + assert policy(facts) == reference(facts) + + +def test_assigned_denial_keeps_team_error_envelope() -> None: + facts: Final = replace(_ASSIGNED, assigned_scope=replace(_MODEL, object_type="key")) + assert team_override(facts) == Decision("Deny", model_access(_MODEL)) + + +def test_bad_assigned_scope_is_terminal_not_team_denial() -> None: + with pytest.raises(NativeAuthFailure): + team_override(replace(_ASSIGNED, assigned_scope=replace(_MODEL, fallback_depth=-1)))