Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/agent_endpoints/auth/agent_access_groups.py: 48%

1import asyncio 

2from collections.abc import Awaitable, Callable 

3from dataclasses import dataclass 

4from typing import Final, TypeAlias 

5 

6from fastapi import HTTPException 

7 

8from litellm._logging import verbose_proxy_logger 

9from litellm.proxy._types import LiteLLM_AccessGroupTable 

10 

11AccessGroupIds: TypeAlias = tuple[str, ...] 

12AccessGroupIdsLoader: TypeAlias = Callable[[str], Awaitable[AccessGroupIds]] # mutable-ok: Callable params 

13LoadedAccessGroup: TypeAlias = LiteLLM_AccessGroupTable | None 

14AccessGroupLoader: TypeAlias = Callable[[str], Awaitable[LoadedAccessGroup]] # mutable-ok: Callable parameter syntax 

15 

16 

17@dataclass(frozen=True, slots=True) 

18class AgentAccessGroupCeiling: 

19 """Everything the agent's attached access groups allow. An empty set denies that resource kind.""" 

20 

21 access_group_ids: AccessGroupIds 

22 models: frozenset[str] 

23 mcp_server_ids: frozenset[str] 

24 agent_ids: frozenset[str] 

25 

26 

27CeilingResolver: TypeAlias = Callable[[str], Awaitable[AgentAccessGroupCeiling | None]] # mutable-ok: Callable params 

28 

29 

30async def _registry_access_group_ids(agent_id: str) -> AccessGroupIds: 

31 from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through 

32 

33 agent: Final = await get_agent_with_read_through(agent_id) 

34 return tuple(agent.access_group_ids or ()) if agent is not None else () 

35 

36 

37async def _load_access_group(access_group_id: str) -> LoadedAccessGroup: 

38 from litellm.proxy.auth.auth_checks import get_access_object 

39 from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache 

40 

41 if prisma_client is None: 

42 verbose_proxy_logger.warning("Agent access group %s cannot be loaded without a DB", access_group_id) 

43 return None 

44 try: 

45 return await get_access_object( 

46 access_group_id=access_group_id, 

47 prisma_client=prisma_client, 

48 user_api_key_cache=user_api_key_cache, 

49 proxy_logging_obj=proxy_logging_obj, 

50 ) 

51 except HTTPException as e: 

52 verbose_proxy_logger.warning( 

53 "Agent access group %s could not be loaded, treating it as empty: %s", access_group_id, e.detail 

54 ) 

55 return None 

56 

57 

58async def resolve_agent_access_group_ceiling( 

59 agent_id: str, 

60 load_access_group_ids: AccessGroupIdsLoader = _registry_access_group_ids, 

61 load_access_group: AccessGroupLoader = _load_access_group, 

62) -> AgentAccessGroupCeiling | None: 

63 """``None`` when the agent has no access groups attached, so nothing is capped.""" 

64 access_group_ids: Final = await load_access_group_ids(agent_id) 

65 if not access_group_ids: 

66 return None 

67 

68 loaded: Final = await asyncio.gather(*(load_access_group(group_id) for group_id in access_group_ids)) 

69 groups: Final = tuple(group for group in loaded if group is not None) 

70 return AgentAccessGroupCeiling( 

71 access_group_ids=access_group_ids, 

72 models=frozenset(model for group in groups for model in group.access_model_names), 

73 mcp_server_ids=frozenset(server_id for group in groups for server_id in group.access_mcp_server_ids), 

74 agent_ids=frozenset(target_id for group in groups for target_id in group.access_agent_ids), 

75 )