From 02ff899fa6f8cbc2c6164bbc2ab4ea312da6db36 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 8 Nov 2026 23:54:47 +0000 Subject: [PATCH] refactor: stage personal model listings and agent ceilings in native policy --- litellm-rust/crates/python-bridge/src/lib.rs | 3 + .../crates/python-bridge/src/request_auth.rs | 20 +++ .../request-auth-core/src/agent_ceilings.rs | 151 +++++++++++++++++ .../crates/request-auth-core/src/lib.rs | 2 + .../src/personal_model_listing.rs | 81 +++++++++ litellm/_auth/agent_ceilings.py | 60 +++++++ litellm/_auth/personal_model_listing.py | 32 ++++ .../auth/agent_access_groups.py | 62 +------ litellm/proxy/auth/agent_ceiling_adapter.py | 156 ++++++++++++++++++ litellm/proxy/auth/model_access_adapter.py | 4 +- .../auth/personal_model_listing_adapter.py | 111 +++++++++++++ litellm/proxy/auth/personal_owner_adapter.py | 8 +- litellm/rust_bridge/_native.pyi | 3 + tests/unit/_auth/test_agent_ceilings.py | 91 ++++++++++ .../unit/_auth/test_personal_model_listing.py | 66 ++++++++ 15 files changed, 789 insertions(+), 61 deletions(-) create mode 100644 litellm-rust/crates/request-auth-core/src/agent_ceilings.rs create mode 100644 litellm-rust/crates/request-auth-core/src/personal_model_listing.rs create mode 100644 litellm/_auth/agent_ceilings.py create mode 100644 litellm/_auth/personal_model_listing.py create mode 100644 litellm/proxy/auth/agent_ceiling_adapter.py create mode 100644 litellm/proxy/auth/personal_model_listing_adapter.py create mode 100644 tests/unit/_auth/test_agent_ceilings.py create mode 100644 tests/unit/_auth/test_personal_model_listing.py diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index cab06f702e4..a1888e0e632 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -56,6 +56,9 @@ mod _native { request_auth_personal_model_request, request_auth_member_models, request_auth_resolved_models, + request_auth_personal_model_listing, + request_auth_agent_ceiling, + request_auth_agent_group_load, }; #[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 0b63c012f4e..6a349aa5db1 100644 --- a/litellm-rust/crates/python-bridge/src/request_auth.rs +++ b/litellm-rust/crates/python-bridge/src/request_auth.rs @@ -287,3 +287,23 @@ pub(crate) fn request_auth_resolved_models(py: Python<'_>, facts: &Bound<'_, PyA let decision = py.detach(move || core::resolved_models::decide(&facts)).map_err(map_error)?; to_py(py, &decision) } + +#[pyfunction] +pub(crate) fn request_auth_personal_model_listing(py: Python<'_>, facts: &Bound<'_, PyAny>) -> PyResult> { + let facts: core::personal_model_listing::Facts = from_py(facts)?; + let decision = py.detach(move || core::personal_model_listing::decide(&facts)).map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_agent_ceiling(py: Python<'_>, facts: &Bound<'_, PyAny>) -> PyResult> { + let facts: core::agent_ceilings::Facts = from_py(facts)?; + let decision = py.detach(move || core::agent_ceilings::decide(&facts)).map_err(map_error)?; + to_py(py, &decision) +} + +#[pyfunction] +pub(crate) fn request_auth_agent_group_load(py: Python<'_>, facts: &Bound<'_, PyAny>) -> PyResult> { + let facts: core::agent_ceilings::LoadFacts = from_py(facts)?; + to_py(py, &py.detach(move || core::agent_ceilings::load(&facts))) +} diff --git a/litellm-rust/crates/request-auth-core/src/agent_ceilings.rs b/litellm-rust/crates/request-auth-core/src/agent_ceilings.rs new file mode 100644 index 00000000000..e270e291078 --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/agent_ceilings.rs @@ -0,0 +1,151 @@ +use crate::Error; +use serde::{Deserialize, Serialize}; +use std::collections::BTreeSet; + +#[derive(Clone, Copy, Deserialize, PartialEq, Eq, strum::VariantArray)] +pub enum Status { + Found, + Missing, + Unavailable, +} + +#[derive(Clone, Deserialize)] +pub struct Group { + pub status: Status, + pub models: Vec, + pub mcp_server_ids: Vec, + pub agent_ids: Vec, +} + +#[derive(Clone, Deserialize)] +pub struct Facts { + pub ids: Option>, + pub groups: Option>, + pub check_db_only: bool, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + ReadIds, + ReadGroups { + authoritative: bool, + }, + Complete { + present: bool, + ids: Vec, + models: BTreeSet, + mcp_server_ids: BTreeSet, + agent_ids: BTreeSet, + indices: Vec, + }, +} + +pub fn decide(facts: &Facts) -> Result { + let Some(ids) = &facts.ids else { + return Ok(Decision::ReadIds); + }; + let empty = |present| Decision::Complete { + present, + ids: vec![], + models: BTreeSet::new(), + mcp_server_ids: BTreeSet::new(), + agent_ids: BTreeSet::new(), + indices: vec![], + }; + if ids.is_empty() { + return Ok(empty(false)); + } + let Some(groups) = &facts.groups else { + return Ok(Decision::ReadGroups { + authoritative: facts.check_db_only, + }); + }; + if groups.len() != ids.len() { + return Err(Error::InvalidAnswer); + } + let indices: Vec<_> = groups + .iter() + .enumerate() + .filter_map(|(index, group)| (group.status == Status::Found).then_some(index)) + .collect(); + Ok(Decision::Complete { + present: true, + ids: ids + .iter() + .zip(groups) + .filter(|(_, group)| group.status != Status::Missing) + .map(|(id, _)| id.clone()) + .collect(), + models: indices + .iter() + .flat_map(|index| groups[*index].models.iter().cloned()) + .collect(), + mcp_server_ids: indices + .iter() + .flat_map(|index| groups[*index].mcp_server_ids.iter().cloned()) + .collect(), + agent_ids: indices + .iter() + .flat_map(|index| groups[*index].agent_ids.iter().cloned()) + .collect(), + indices, + }) +} + +#[derive(Clone, Copy, Deserialize, strum::VariantArray)] +pub enum ReadStatus { + Success, + NotFound, + HttpError, +} + +#[derive(Clone, Deserialize)] +pub struct LoadFacts { + pub database_present: bool, + pub check_db_only: bool, + pub status: Option, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum LoadDecision { + Read, + Return, + DatabaseUnavailable, + HttpUnavailable, + Raise, +} + +pub fn load(facts: &LoadFacts) -> LoadDecision { + if !facts.database_present { + return LoadDecision::DatabaseUnavailable; + } + match facts.status { + None => LoadDecision::Read, + Some(ReadStatus::Success) => LoadDecision::Return, + Some(ReadStatus::NotFound) => LoadDecision::Raise, + Some(ReadStatus::HttpError) if facts.check_db_only => LoadDecision::Raise, + Some(ReadStatus::HttpError) => LoadDecision::HttpUnavailable, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::soft(false, LoadDecision::HttpUnavailable)] + #[case::authoritative(true, LoadDecision::Raise)] + fn http_lookup_policy(#[case] check_db_only: bool, #[case] expected: LoadDecision) { + assert_eq!( + load(&LoadFacts { + database_present: true, + check_db_only, + status: Some(ReadStatus::HttpError) + }), + expected + ); + } +} diff --git a/litellm-rust/crates/request-auth-core/src/lib.rs b/litellm-rust/crates/request-auth-core/src/lib.rs index a1e5dcc4aee..b8ee07066f7 100644 --- a/litellm-rust/crates/request-auth-core/src/lib.rs +++ b/litellm-rust/crates/request-auth-core/src/lib.rs @@ -23,6 +23,8 @@ pub mod personal_owner; pub mod personal_model_request; pub mod member_models; pub mod resolved_models; +pub mod personal_model_listing; +pub mod agent_ceilings; pub mod owner_teams; pub mod protocol; pub mod route_access; diff --git a/litellm-rust/crates/request-auth-core/src/personal_model_listing.rs b/litellm-rust/crates/request-auth-core/src/personal_model_listing.rs new file mode 100644 index 00000000000..98c66742bbe --- /dev/null +++ b/litellm-rust/crates/request-auth-core/src/personal_model_listing.rs @@ -0,0 +1,81 @@ +use crate::{Error, personal_model_request::Loaded, personal_models, personal_owner}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Deserialize)] +pub struct Facts { + pub owner: personal_owner::Facts, + pub names: Option>, + pub models: Vec, +} + +#[derive(Debug, PartialEq, Eq, Serialize)] +#[serde(tag = "kind")] +pub enum Decision { + ReadOwner, + ReadNames, + ReadModel { index: usize, teams: Vec }, + Complete { selected: Vec }, +} + +pub fn decide(facts: &Facts) -> Result { + let teams = match personal_owner::decide(&facts.owner)? { + personal_owner::Decision::Complete { selected, .. } => selected, + _ => return Ok(Decision::ReadOwner), + }; + if facts + .owner + .teams + .as_ref() + .is_none_or(|teams| teams.loads.is_empty()) + { + return Ok(Decision::Complete { selected: vec![] }); + } + let Some(names) = &facts.names else { + return Ok(Decision::ReadNames); + }; + let mut selected = vec![]; + for index in 0..names.len() { + 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 } => { + if allowed { + selected.push(index); + } + } + _ => return Err(Error::InvalidAnswer), + } + } + Ok(Decision::Complete { selected }) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::missing_database(false, Decision::Complete { selected: vec![] })] + #[case::load(true, Decision::ReadOwner)] + fn owner_before_names(#[case] database_present: bool, #[case] expected: Decision) { + let facts = Facts { + owner: personal_owner::Facts { + user_present: true, + database_present, + cache_present: true, + logging_present: true, + owner_loaded: false, + owner_present: false, + owner_failed: false, + teams: None, + }, + names: None, + models: vec![], + }; + assert_eq!(decide(&facts).unwrap(), expected); + } +} diff --git a/litellm/_auth/agent_ceilings.py b/litellm/_auth/agent_ceilings.py new file mode 100644 index 00000000000..544bf05d556 --- /dev/null +++ b/litellm/_auth/agent_ceilings.py @@ -0,0 +1,60 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class AgentCeilingGroup: + status: Literal["Found", "Missing", "Unavailable"] + models: tuple[str, ...] = () + mcp_server_ids: tuple[str, ...] = () + agent_ids: tuple[str, ...] = () + + +@dataclass(frozen=True, slots=True) +class AgentCeilingFacts: + ids: tuple[str, ...] | None = None + groups: tuple[AgentCeilingGroup, ...] | None = None + check_db_only: bool = False + + +@dataclass(frozen=True, slots=True) +class AgentCeilingDecision: + kind: Literal["ReadIds", "ReadGroups", "Complete"] + authoritative: bool = False + present: bool = False + ids: tuple[str, ...] = () + models: frozenset[str] = frozenset() + mcp_server_ids: frozenset[str] = frozenset() + agent_ids: frozenset[str] = frozenset() + indices: tuple[int, ...] = () + + +@dataclass(frozen=True, slots=True) +class AgentGroupLoadFacts: + database_present: bool + check_db_only: bool = False + status: Literal["Success", "NotFound", "HttpError"] | None = None + + +@dataclass(frozen=True, slots=True) +class AgentGroupLoadDecision: + kind: Literal["Read", "Return", "DatabaseUnavailable", "HttpUnavailable", "Raise"] + + +_DECISION: Final = TypeAdapter(AgentCeilingDecision) +_LOAD: Final = TypeAdapter(AgentGroupLoadDecision) + + +@native_boundary +def agent_ceiling(facts: AgentCeilingFacts) -> AgentCeilingDecision: + return _DECISION.validate_python(_native.request_auth_agent_ceiling(asdict(facts))) + + +@native_boundary +def agent_group_load(facts: AgentGroupLoadFacts) -> AgentGroupLoadDecision: + return _LOAD.validate_python(_native.request_auth_agent_group_load(asdict(facts))) diff --git a/litellm/_auth/personal_model_listing.py b/litellm/_auth/personal_model_listing.py new file mode 100644 index 00000000000..e90cc43c388 --- /dev/null +++ b/litellm/_auth/personal_model_listing.py @@ -0,0 +1,32 @@ +from dataclasses import asdict, dataclass +from typing import Final, Literal + +from pydantic import TypeAdapter + +from litellm._auth.failure import native_boundary +from litellm._auth.personal_models import LoadedTeamModelFacts +from litellm._auth.personal_owner import PersonalOwnerFacts +from litellm.rust_bridge import _native + + +@dataclass(frozen=True, slots=True) +class PersonalModelListingFacts: + owner: PersonalOwnerFacts + names: tuple[str, ...] | None = None + models: tuple[LoadedTeamModelFacts, ...] = () + + +@dataclass(frozen=True, slots=True) +class PersonalModelListingDecision: + kind: Literal["ReadOwner", "ReadNames", "ReadModel", "Complete"] + index: int = 0 + teams: tuple[int, ...] = () + selected: tuple[int, ...] = () + + +_DECISION: Final = TypeAdapter(PersonalModelListingDecision) + + +@native_boundary +def personal_model_listing(facts: PersonalModelListingFacts) -> PersonalModelListingDecision: + return _DECISION.validate_python(_native.request_auth_personal_model_listing(asdict(facts))) diff --git a/litellm/proxy/agent_endpoints/auth/agent_access_groups.py b/litellm/proxy/agent_endpoints/auth/agent_access_groups.py index 557d1c7bc0a..34bb6781df4 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_access_groups.py +++ b/litellm/proxy/agent_endpoints/auth/agent_access_groups.py @@ -1,11 +1,7 @@ -import asyncio from collections.abc import Awaitable, Callable from dataclasses import dataclass from typing import TYPE_CHECKING, Final, TypeAlias -from fastapi import HTTPException - -from litellm._logging import verbose_proxy_logger from litellm.proxy._types import LiteLLM_AccessGroupTable if TYPE_CHECKING: @@ -31,13 +27,6 @@ class AgentAccessGroupCeiling: CeilingResolver: TypeAlias = Callable[[str], Awaitable[AgentAccessGroupCeiling | None]] -class _MissingGroup: - __slots__ = () - - -_MISSING: Final = _MissingGroup() - - async def _registry_access_group_ids(agent_id: str) -> AccessGroupIds: from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through @@ -46,29 +35,9 @@ async def _registry_access_group_ids(agent_id: str) -> AccessGroupIds: async def _load_access_group(access_group_id: str, *, check_db_only: bool = False) -> LoadedAccessGroup: - from litellm.proxy.auth.auth_checks import AccessGroupNotFoundError, get_access_object - from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + from litellm.proxy.auth.agent_ceiling_adapter import load_agent_access_group - if prisma_client is None: - verbose_proxy_logger.warning("Agent access group %s cannot be loaded without a DB", access_group_id) - return None - try: - return await get_access_object( - access_group_id=access_group_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - check_db_only=check_db_only, - ) - except AccessGroupNotFoundError: - raise - except HTTPException as e: - if check_db_only: - raise - verbose_proxy_logger.warning( - "Agent access group %s could not be loaded, treating it as empty: %s", access_group_id, e.detail - ) - return None + return await load_agent_access_group(access_group_id, check_db_only=check_db_only) async def resolve_agent_access_group_ceiling( @@ -82,31 +51,10 @@ async def resolve_agent_access_group_ceiling( provably does not exist is dropped from the ceiling, it grants and bounds nothing, while a group that exists but could not be read stays named and loads as ``None``. The ceiling itself stays once any group is named, so an agent whose every group is missing is still capped to nothing.""" - named_ids: Final = await load_access_group_ids(agent_id) - if not named_ids: - return None + from litellm.proxy.auth.agent_ceiling_adapter import project_agent_ceiling, read_agent_ceiling_facts - from litellm.proxy.auth.auth_checks import AccessGroupNotFoundError - - async def load(group_id: str) -> LoadedAccessGroup | _MissingGroup: - try: - return await ( - _load_access_group(group_id, check_db_only=True) if check_db_only else load_access_group(group_id) - ) - except AccessGroupNotFoundError: - return _MISSING - - results: Final = await asyncio.gather(*(load(group_id) for group_id in named_ids)) - access_group_ids: Final = tuple( - group_id for group_id, result in zip(named_ids, results, strict=True) if not isinstance(result, _MissingGroup) - ) - groups: Final = tuple(result for result in results if isinstance(result, LiteLLM_AccessGroupTable)) - return AgentAccessGroupCeiling( - access_group_ids=access_group_ids, - models=frozenset(model for group in groups for model in group.access_model_names), - mcp_server_ids=frozenset(server_id for group in groups for server_id in group.access_mcp_server_ids), - agent_ids=frozenset(target_id for group in groups for target_id in group.access_agent_ids), - access_groups=groups, + return project_agent_ceiling( + await read_agent_ceiling_facts(agent_id, load_access_group_ids, load_access_group, check_db_only=check_db_only) ) diff --git a/litellm/proxy/auth/agent_ceiling_adapter.py b/litellm/proxy/auth/agent_ceiling_adapter.py new file mode 100644 index 00000000000..aee14fc9c6b --- /dev/null +++ b/litellm/proxy/auth/agent_ceiling_adapter.py @@ -0,0 +1,156 @@ +from __future__ import annotations + +import asyncio +from dataclasses import dataclass, replace +from typing import TYPE_CHECKING, Final + +from fastapi import HTTPException + +from litellm._auth.agent_ceilings import ( + AgentCeilingFacts, + AgentCeilingGroup, + AgentGroupLoadFacts, + agent_ceiling, + agent_group_load, +) +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import LiteLLM_AccessGroupTable + +if TYPE_CHECKING: + from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( + AccessGroupIdsLoader, + AccessGroupLoader, + AgentAccessGroupCeiling, + ) + + +@dataclass(frozen=True, slots=True) +class AgentCeilingReading: + facts: AgentCeilingFacts + rows: tuple[LiteLLM_AccessGroupTable | None, ...] = () + + +@dataclass(frozen=True, slots=True) +class _Source: + agent_id: str + ids: AccessGroupIdsLoader + group: AccessGroupLoader + + +async def read_agent_ceiling_facts( + agent_id: str, + load_access_group_ids: AccessGroupIdsLoader, + load_access_group: AccessGroupLoader, + *, + check_db_only: bool = False, +) -> AgentCeilingReading: + return await _read( + AgentCeilingReading(AgentCeilingFacts(check_db_only=check_db_only)), + _Source(agent_id, load_access_group_ids, load_access_group), + ) + + +def project_agent_ceiling(reading: AgentCeilingReading) -> AgentAccessGroupCeiling | None: + from litellm.proxy.agent_endpoints.auth.agent_access_groups import AgentAccessGroupCeiling + + decision: Final = agent_ceiling(reading.facts) + if decision.kind != "Complete": + raise NativeAuthRuntimeError("incomplete agent ceiling facts") + if not decision.present: + return None + return AgentAccessGroupCeiling( + decision.ids, + decision.models, + decision.mcp_server_ids, + decision.agent_ids, + tuple(_row(reading, index) for index in decision.indices), + ) + + +def _row(reading: AgentCeilingReading, index: int) -> LiteLLM_AccessGroupTable: + row: Final = reading.rows[index] + if row is None: + raise NativeAuthRuntimeError("agent ceiling selected a missing row") + return row + + +async def _group(group_id: str, loader: AccessGroupLoader) -> tuple[AgentCeilingGroup, LiteLLM_AccessGroupTable | None]: + from litellm.proxy.auth.auth_checks import AccessGroupNotFoundError + + try: + row: Final = await loader(group_id) + except NativeAuthFailure: + raise + except AccessGroupNotFoundError: + return AgentCeilingGroup("Missing"), None + if not isinstance(row, LiteLLM_AccessGroupTable): + return AgentCeilingGroup("Unavailable"), None + return AgentCeilingGroup( + "Found", tuple(row.access_model_names), tuple(row.access_mcp_server_ids), tuple(row.access_agent_ids) + ), row + + +async def _authoritative(group_id: str) -> LiteLLM_AccessGroupTable | None: + from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( + _load_access_group, # pyright: ignore[reportPrivateUsage] # authoritative raw reader preserves the public injected-loader contract + ) + + return await _load_access_group(group_id, check_db_only=True) + + +async def _read(reading: AgentCeilingReading, source: _Source) -> AgentCeilingReading: + decision: Final = agent_ceiling(reading.facts) + match decision.kind: + case "Complete": + return reading + case "ReadIds": + ids: Final = await source.ids(source.agent_id) + return await _read(replace(reading, facts=replace(reading.facts, ids=ids)), source) + case "ReadGroups": + if reading.facts.ids is None: + raise NativeAuthRuntimeError("agent group reads without named IDs") + loader: Final = _authoritative if decision.authoritative else source.group + rows: Final = tuple(await asyncio.gather(*(_group(group_id, loader) for group_id in reading.facts.ids))) + return AgentCeilingReading( + replace(reading.facts, groups=tuple(row[0] for row in rows)), tuple(row[1] for row in rows) + ) + + +async def load_agent_access_group( + access_group_id: str, *, check_db_only: bool = False +) -> LiteLLM_AccessGroupTable | None: + from litellm.proxy.auth.auth_checks import AccessGroupNotFoundError, get_access_object + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + initial: Final = AgentGroupLoadFacts(prisma_client is not None, check_db_only) + initial_decision: Final = agent_group_load(initial) + if initial_decision.kind == "DatabaseUnavailable": + verbose_proxy_logger.warning("Agent access group %s cannot be loaded without a DB", access_group_id) + return None + if initial_decision.kind != "Read": + raise NativeAuthRuntimeError("invalid initial agent group read decision") + try: + return await get_access_object( + access_group_id=access_group_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + check_db_only=check_db_only, + ) + except NativeAuthFailure: + raise + except AccessGroupNotFoundError: + if agent_group_load(replace(initial, status="NotFound")).kind != "Raise": + raise NativeAuthRuntimeError("invalid missing agent group decision") from None + raise + except HTTPException as error: + decision: Final = agent_group_load(replace(initial, status="HttpError")) + if decision.kind == "Raise": + raise + if decision.kind != "HttpUnavailable": + raise NativeAuthRuntimeError("invalid failed agent group decision") from error + verbose_proxy_logger.warning( + "Agent access group %s could not be loaded, treating it as empty: %s", access_group_id, error.detail + ) + return None diff --git a/litellm/proxy/auth/model_access_adapter.py b/litellm/proxy/auth/model_access_adapter.py index c04b4fcc04e..b8df7e88ef5 100644 --- a/litellm/proxy/auth/model_access_adapter.py +++ b/litellm/proxy/auth/model_access_adapter.py @@ -6,7 +6,7 @@ from fastapi import HTTPException from pydantic import BaseModel import litellm -from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError +from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError, native_boundary from litellm._auth.model_access import ( ModelAccessDecision, ModelAccessFacts, @@ -45,6 +45,7 @@ class ModelLists(BaseModel, frozen=True): models: tuple[str, ...] = () +@native_boundary def object_model_grants(source: object) -> tuple[str, ...]: return ModelLists.model_validate(source, from_attributes=True).models @@ -71,6 +72,7 @@ def _record(model: str, router: RegistryReader | None, team_id: str | None) -> M ) +@native_boundary def project_model_facts( model: str | Sequence[str], llm_router: RegistryReader | None, diff --git a/litellm/proxy/auth/personal_model_listing_adapter.py b/litellm/proxy/auth/personal_model_listing_adapter.py new file mode 100644 index 00000000000..5e79f062535 --- /dev/null +++ b/litellm/proxy/auth/personal_model_listing_adapter.py @@ -0,0 +1,111 @@ +from collections.abc import Sequence +from dataclasses import dataclass, replace +from typing import Final + +from litellm._auth.failure import NativeAuthRuntimeError +from litellm._auth.personal_model_listing import PersonalModelListingFacts, personal_model_listing +from litellm._auth.personal_owner import PersonalOwnerFacts, personal_owner +from litellm.proxy._types import LiteLLM_TeamTableCachedObj, UserAPIKeyAuth +from litellm.proxy.auth.personal_model_adapter import read_loaded_team_model_facts +from litellm.proxy.auth.personal_owner_adapter import read_personal_owner_facts, 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: + names: Sequence[str] + token: UserAPIKeyAuth + router: RegistryReader | None + prisma: PrismaClient | None + cache: UserApiKeyCache | None + logging: ProxyLogging | None + teams: tuple[LiteLLM_TeamTableCachedObj | None, ...] = () + + +async def personal_key_listing_teams( + user_api_key_dict: UserAPIKeyAuth, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache | None, + proxy_logging_obj: ProxyLogging | None, + *, + check_db_only: bool = False, +) -> tuple[LiteLLM_TeamTableCachedObj | None, ...]: + reading: Final = await read_personal_owner_facts( + user_api_key_dict, + prisma_client, + user_api_key_cache, + proxy_logging_obj, + check_db_only=check_db_only, + purpose="Models", + ) + decision: Final = personal_owner(reading.facts) + if decision.kind != "Complete": + raise NativeAuthRuntimeError("incomplete listing team facts") + return tuple(team if index in decision.selected else None for index, team in enumerate(reading.teams)) + + +async def read_personal_model_listing_facts( + models: Sequence[str], + user_api_key_dict: UserAPIKeyAuth, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache | None, + proxy_logging_obj: ProxyLogging | None, +) -> PersonalModelListingFacts: + owner: Final = PersonalOwnerFacts( + user_api_key_dict.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( + PersonalModelListingFacts(owner), + _Source(models, user_api_key_dict, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj), + ) + + +async def personal_key_team_visible_models( + models: Sequence[str], + user_api_key_dict: UserAPIKeyAuth, + llm_router: RegistryReader | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache | None, + proxy_logging_obj: ProxyLogging | None, +) -> tuple[str, ...]: + facts: Final = await read_personal_model_listing_facts( + models, user_api_key_dict, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj + ) + decision: Final = personal_model_listing(facts) + if decision.kind != "Complete": + raise NativeAuthRuntimeError("incomplete personal model listing facts") + return tuple((facts.names or ())[index] for index in decision.selected) + + +async def _read(facts: PersonalModelListingFacts, source: _Source) -> PersonalModelListingFacts: + decision: Final = personal_model_listing(facts) + match decision.kind: + case "Complete": + return facts + case "ReadOwner": + owner: Final = await read_personal_owner_facts( + source.token, source.prisma, source.cache, source.logging, purpose="Models" + ) + return await _read(replace(facts, owner=owner.facts), replace(source, teams=owner.teams)) + case "ReadNames": + return await _read(replace(facts, names=tuple(source.names)), source) + case "ReadModel": + if source.prisma is None or source.cache is None or source.logging is None or facts.names is None: + raise NativeAuthRuntimeError("personal listing scope read without dependencies") + teams: Final = tuple(selected_team(source.teams, index) for index in decision.teams) + model: Final = await read_loaded_team_model_facts( + facts.names[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 index 20add242af3..01c3ea5ddd1 100644 --- a/litellm/proxy/auth/personal_owner_adapter.py +++ b/litellm/proxy/auth/personal_owner_adapter.py @@ -1,7 +1,7 @@ from __future__ import annotations from dataclasses import dataclass, replace -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Final, Literal from litellm._auth.failure import NativeAuthFailure, NativeAuthRuntimeError from litellm._auth.personal_owner import PersonalOwnerFacts, personal_owner @@ -29,6 +29,7 @@ class _Source: logging: ProxyLogging | None check_db_only: bool owner: LiteLLM_UserTable | None = None + purpose: Literal["Models", "Grants"] = "Grants" async def read_personal_owner_facts( @@ -38,6 +39,7 @@ async def read_personal_owner_facts( proxy_logging_obj: ProxyLogging | None, *, check_db_only: bool = False, + purpose: Literal["Models", "Grants"] = "Grants", ) -> PersonalOwnerReading: facts: Final = PersonalOwnerFacts( valid_token.user_id is not None, @@ -47,7 +49,7 @@ async def read_personal_owner_facts( ) return await _read( PersonalOwnerReading(facts), - _Source(valid_token, prisma_client, user_api_key_cache, proxy_logging_obj, check_db_only), + _Source(valid_token, prisma_client, user_api_key_cache, proxy_logging_obj, check_db_only, purpose=purpose), ) @@ -121,7 +123,7 @@ async def _read(reading: PersonalOwnerReading, source: _Source) -> PersonalOwner source.prisma, source.cache, source.logging, - purpose="Grants", + purpose=source.purpose, check_db_only=source.check_db_only, ) return PersonalOwnerReading(replace(reading.facts, teams=teams.facts), teams.teams) diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index d8109c9337c..2c6d76bb449 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -45,6 +45,9 @@ def request_auth_personal_owner(facts: object) -> object: ... def request_auth_personal_model_request(facts: object) -> object: ... def request_auth_member_models(facts: object) -> object: ... def request_auth_resolved_models(facts: object) -> object: ... +def request_auth_personal_model_listing(facts: object) -> object: ... +def request_auth_agent_ceiling(facts: object) -> object: ... +def request_auth_agent_group_load(facts: object) -> object: ... @final class NativeClickHouseSpendConfig: diff --git a/tests/unit/_auth/test_agent_ceilings.py b/tests/unit/_auth/test_agent_ceilings.py new file mode 100644 index 00000000000..a32fb1f6488 --- /dev/null +++ b/tests/unit/_auth/test_agent_ceilings.py @@ -0,0 +1,91 @@ +from collections.abc import Callable +from itertools import chain +from typing import Final + +import pytest + +from litellm._auth.agent_ceilings import AgentCeilingDecision as Decision +from litellm._auth.agent_ceilings import AgentCeilingFacts as Facts +from litellm._auth.agent_ceilings import AgentCeilingGroup as Group +from litellm._auth.agent_ceilings import AgentGroupLoadDecision as LoadDecision +from litellm._auth.agent_ceilings import AgentGroupLoadFacts as LoadFacts +from litellm._auth.agent_ceilings import agent_ceiling, agent_group_load +from litellm._auth.failure import NativeAuthFailure + + +def reference(facts: Facts) -> Decision: + if facts.ids is None: + return Decision("ReadIds") + if not facts.ids: + return Decision("Complete") + if facts.groups is None: + return Decision("ReadGroups", authoritative=facts.check_db_only) + indices: Final = tuple(index for index, group in enumerate(facts.groups) if group.status == "Found") + return Decision( + "Complete", + present=True, + ids=tuple(name for name, group in zip(facts.ids, facts.groups, strict=True) if group.status != "Missing"), + models=frozenset(chain.from_iterable(facts.groups[index].models for index in indices)), + mcp_server_ids=frozenset(chain.from_iterable(facts.groups[index].mcp_server_ids for index in indices)), + agent_ids=frozenset(chain.from_iterable(facts.groups[index].agent_ids for index in indices)), + indices=indices, + ) + + +def load_reference(facts: LoadFacts) -> LoadDecision: + if not facts.database_present: + return LoadDecision("DatabaseUnavailable") + if facts.status is None: + return LoadDecision("Read") + if facts.status == "Success": + return LoadDecision("Return") + return LoadDecision("Raise" if facts.status == "NotFound" or facts.check_db_only else "HttpUnavailable") + + +@pytest.mark.parametrize("policy", (reference, agent_ceiling)) +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (Facts(), Decision("ReadIds")), + (Facts(()), Decision("Complete")), + (Facts(("group",)), Decision("ReadGroups")), + (Facts(("group",), check_db_only=True), Decision("ReadGroups", authoritative=True)), + (Facts(("group",), (Group("Missing"),)), Decision("Complete", present=True)), + (Facts(("group",), (Group("Unavailable"),)), Decision("Complete", present=True, ids=("group",))), + ( + Facts( + ("one", "two", "one"), + (Group("Found", ("model",), ("server",), ("agent",)), Group("Missing"), Group("Found", ("model",))), + ), + Decision( + "Complete", + present=True, + ids=("one", "one"), + models=frozenset(("model",)), + mcp_server_ids=frozenset(("server",)), + agent_ids=frozenset(("agent",)), + indices=(0, 2), + ), + ), + ), +) +def test_ceiling_union_and_missing_vs_unavailable( + policy: Callable[[Facts], Decision], facts: Facts, expected: Decision +) -> None: + assert policy(facts) == expected + + +@pytest.mark.parametrize("policy", (load_reference, agent_group_load)) +@pytest.mark.parametrize("database", (False, True)) +@pytest.mark.parametrize("strict", (False, True)) +@pytest.mark.parametrize("status", (None, "Success", "NotFound", "HttpError")) +def test_group_read_error_policy( + policy: Callable[[LoadFacts], LoadDecision], database: bool, strict: bool, status: str | None +) -> None: + facts: Final = LoadFacts(database, strict, status) # pyright: ignore[reportArgumentType] # parameter table covers the complete typed status domain + assert policy(facts) == load_reference(facts) + + +def test_mismatched_group_answers_are_terminal() -> None: + with pytest.raises(NativeAuthFailure): + agent_ceiling(Facts(("group",), ())) diff --git a/tests/unit/_auth/test_personal_model_listing.py b/tests/unit/_auth/test_personal_model_listing.py new file mode 100644 index 00000000000..0d40963c03f --- /dev/null +++ b/tests/unit/_auth/test_personal_model_listing.py @@ -0,0 +1,66 @@ +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_listing import PersonalModelListingDecision as Decision +from litellm._auth.personal_model_listing import PersonalModelListingFacts as Facts +from litellm._auth.personal_model_listing import personal_model_listing +from litellm._auth.personal_models import LoadedTeamModelFacts +from litellm._auth.personal_owner import PersonalOwnerFacts +from tests.unit._auth.test_personal_models import ALLOW, INITIAL +from tests.unit._auth.test_personal_models import reference as scope_reference +from tests.unit._auth.test_personal_owner import reference as owner_reference + + +def reference(facts: Facts) -> Decision: + owner: Final = owner_reference(facts.owner) + if owner.kind != "Complete": + return Decision("ReadOwner") + if facts.owner.teams is None or not facts.owner.teams.loads: + return Decision("Complete") + if facts.names is None: + return Decision("ReadNames") + if len(facts.models) < len(facts.names): + return Decision("ReadModel", index=len(facts.models), teams=owner.selected) + return Decision( + "Complete", + selected=tuple( + index + for index, model in enumerate(facts.models) + if any(tuple(scope_reference(scope).allowed for scope in model.scopes)) + ), + ) + + +_BASE: Final = Facts(PersonalOwnerFacts(True, True, True, True)) +_OWNER: Final = replace( + _BASE.owner, + owner_loaded=True, + owner_present=True, + teams=OwnerTeamFacts("Models", ("team",), "ui", loaded=True, loads=(OwnerTeamLoad(True),)), +) +_LOADED: Final = replace(_BASE, owner=_OWNER, names=("first", "second", "first")) +_ALLOWED: Final = LoadedTeamModelFacts((True,), (ALLOW,)) +_DENIED: Final = LoadedTeamModelFacts((True,), (replace(INITIAL, io_failed=True),)) + + +@pytest.mark.parametrize("policy", (reference, personal_model_listing)) +@pytest.mark.parametrize( + ("facts", "expected"), + ( + (_BASE, Decision("ReadOwner")), + (replace(_BASE, owner=replace(_BASE.owner, database_present=False)), Decision("Complete")), + (replace(_BASE, owner=replace(_BASE.owner, owner_loaded=True, owner_failed=True)), Decision("Complete")), + (replace(_BASE, owner=_OWNER), Decision("ReadNames")), + (_LOADED, Decision("ReadModel", teams=(0,))), + (replace(_LOADED, models=(_DENIED,)), Decision("ReadModel", index=1, teams=(0,))), + (replace(_LOADED, models=(_ALLOWED, _DENIED, _ALLOWED)), Decision("Complete", selected=(0, 2))), + ), +) +def test_listing_preserves_order_duplicates_and_checks_after_denial( + policy: Callable[[Facts], Decision], facts: Facts, expected: Decision +) -> None: + assert policy(facts) == expected