chore: initial public snapshot for github upload
This commit is contained in:
Your Name
@@ -0,0 +1,451 @@
|
||||
"""
|
||||
Agent Permission Handler for LiteLLM Proxy.
|
||||
|
||||
Handles agent permission checking for keys and teams using object_permission_id.
|
||||
Follows the same pattern as MCP permission handling.
|
||||
"""
|
||||
|
||||
from typing import List, Optional, Set
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_ObjectPermissionTable,
|
||||
LiteLLM_TeamTable,
|
||||
UI_TEAM_ID,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
|
||||
|
||||
class AgentRequestHandler:
|
||||
"""
|
||||
Class to handle agent permission checking, including:
|
||||
1. Key-level agent permissions
|
||||
2. Team-level agent permissions
|
||||
3. Agent access group resolution
|
||||
|
||||
Follows the same inheritance logic as MCP:
|
||||
- If team has restrictions and key has restrictions: use intersection
|
||||
- If team has restrictions and key has none: inherit from team
|
||||
- If team has no restrictions: use key restrictions
|
||||
- If no restrictions: allow all agents
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
async def get_allowed_agents(
|
||||
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
|
||||
) -> List[str]:
|
||||
"""
|
||||
Get list of allowed agent IDs for the given user/key based on permissions.
|
||||
|
||||
Returns:
|
||||
List[str]: List of allowed agent IDs. Empty list means no restrictions (allow all).
|
||||
"""
|
||||
try:
|
||||
allowed_agents: List[str] = []
|
||||
allowed_agents_for_key = (
|
||||
await AgentRequestHandler._get_allowed_agents_for_key(user_api_key_auth)
|
||||
)
|
||||
allowed_agents_for_team = (
|
||||
await AgentRequestHandler._get_allowed_agents_for_team(
|
||||
user_api_key_auth
|
||||
)
|
||||
)
|
||||
|
||||
# If team has agent restrictions, handle inheritance and intersection logic
|
||||
if len(allowed_agents_for_team) > 0:
|
||||
if len(allowed_agents_for_key) > 0:
|
||||
# Key has its own agent permissions - use intersection with team permissions
|
||||
for agent_id in allowed_agents_for_key:
|
||||
if agent_id in allowed_agents_for_team:
|
||||
allowed_agents.append(agent_id)
|
||||
else:
|
||||
# Key has no agent permissions - inherit from team
|
||||
allowed_agents = allowed_agents_for_team
|
||||
else:
|
||||
allowed_agents = allowed_agents_for_key
|
||||
|
||||
return list(set(allowed_agents))
|
||||
except Exception as e:
|
||||
verbose_logger.warning(f"Failed to get allowed agents: {str(e)}")
|
||||
return []
|
||||
|
||||
@staticmethod
|
||||
async def is_agent_allowed(
|
||||
agent_id: str,
|
||||
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
|
||||
) -> bool:
|
||||
"""
|
||||
Check if a specific agent is allowed for the given user/key.
|
||||
|
||||
Args:
|
||||
agent_id: The agent ID to check
|
||||
user_api_key_auth: User authentication info
|
||||
|
||||
Returns:
|
||||
bool: True if agent is allowed, False otherwise
|
||||
"""
|
||||
allowed_agents = await AgentRequestHandler.get_allowed_agents(user_api_key_auth)
|
||||
|
||||
# Empty list means no restrictions - allow all
|
||||
if len(allowed_agents) == 0:
|
||||
return True
|
||||
|
||||
return agent_id in allowed_agents
|
||||
|
||||
@staticmethod
|
||||
def _get_key_object_permission(
|
||||
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
|
||||
) -> Optional[LiteLLM_ObjectPermissionTable]:
|
||||
"""
|
||||
Get key object_permission - already loaded by get_key_object() in main auth flow.
|
||||
|
||||
Note: object_permission is automatically populated when the key is fetched via
|
||||
get_key_object() in litellm/proxy/auth/auth_checks.py
|
||||
"""
|
||||
if not user_api_key_auth:
|
||||
return None
|
||||
|
||||
return user_api_key_auth.object_permission
|
||||
|
||||
@staticmethod
|
||||
async def _get_team_object_permission(
|
||||
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
|
||||
) -> Optional[LiteLLM_ObjectPermissionTable]:
|
||||
"""
|
||||
Get team object_permission - automatically loaded by get_team_object() in main auth flow.
|
||||
|
||||
Note: object_permission is automatically populated when the team is fetched via
|
||||
get_team_object() in litellm/proxy/auth/auth_checks.py
|
||||
"""
|
||||
from litellm.proxy.auth.auth_checks import get_team_object
|
||||
from litellm.proxy.proxy_server import (
|
||||
prisma_client,
|
||||
proxy_logging_obj,
|
||||
user_api_key_cache,
|
||||
)
|
||||
|
||||
if not user_api_key_auth or not user_api_key_auth.team_id or not prisma_client:
|
||||
return None
|
||||
|
||||
# Get the team object (which has object_permission already loaded)
|
||||
team_obj: Optional[LiteLLM_TeamTable] = await get_team_object(
|
||||
team_id=user_api_key_auth.team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=user_api_key_auth.parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
if not team_obj:
|
||||
return None
|
||||
|
||||
return team_obj.object_permission
|
||||
|
||||
@staticmethod
|
||||
async def _get_allowed_agents_for_key(
|
||||
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
|
||||
) -> List[str]:
|
||||
"""
|
||||
Get allowed agents for a key.
|
||||
|
||||
1. First checks native key-level agent permissions (object_permission)
|
||||
2. Also includes agents from key's access_group_ids (unified access groups)
|
||||
|
||||
Note: object_permission is already loaded by get_key_object() in main auth flow.
|
||||
"""
|
||||
if user_api_key_auth is None:
|
||||
return []
|
||||
|
||||
try:
|
||||
all_agents: List[str] = []
|
||||
|
||||
# 1. Get agents from object_permission (native permissions)
|
||||
key_object_permission = AgentRequestHandler._get_key_object_permission(
|
||||
user_api_key_auth
|
||||
)
|
||||
if key_object_permission is not None:
|
||||
# Get direct agents
|
||||
direct_agents = key_object_permission.agents or []
|
||||
|
||||
# Get agents from access groups
|
||||
access_group_agents = (
|
||||
await AgentRequestHandler._get_agents_from_access_groups(
|
||||
key_object_permission.agent_access_groups or []
|
||||
)
|
||||
)
|
||||
|
||||
all_agents = direct_agents + access_group_agents
|
||||
|
||||
# 2. Fallback: get agent IDs from key's access_group_ids (unified access groups)
|
||||
key_access_group_ids = user_api_key_auth.access_group_ids or []
|
||||
if key_access_group_ids:
|
||||
from litellm.proxy.auth.auth_checks import (
|
||||
_get_agent_ids_from_access_groups,
|
||||
)
|
||||
|
||||
unified_agents = await _get_agent_ids_from_access_groups(
|
||||
access_group_ids=key_access_group_ids,
|
||||
)
|
||||
all_agents.extend(unified_agents)
|
||||
|
||||
return list(set(all_agents))
|
||||
except Exception as e:
|
||||
verbose_logger.warning(f"Failed to get allowed agents for key: {str(e)}")
|
||||
return []
|
||||
|
||||
@staticmethod
|
||||
async def _get_allowed_agents_for_team(
|
||||
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
|
||||
) -> List[str]:
|
||||
"""
|
||||
Get allowed agents for a team.
|
||||
|
||||
1. First checks native team-level agent permissions (object_permission)
|
||||
2. Also includes agents from team's access_group_ids (unified access groups)
|
||||
|
||||
Fetches the team object once and reuses it for both permission sources.
|
||||
"""
|
||||
if user_api_key_auth is None:
|
||||
return []
|
||||
|
||||
if user_api_key_auth.team_id is None:
|
||||
return []
|
||||
|
||||
try:
|
||||
from litellm.proxy.auth.auth_checks import get_team_object
|
||||
from litellm.proxy.proxy_server import (
|
||||
prisma_client,
|
||||
proxy_logging_obj,
|
||||
user_api_key_cache,
|
||||
)
|
||||
|
||||
if not prisma_client:
|
||||
return []
|
||||
|
||||
# Fetch the team object once for both permission sources
|
||||
team_obj = await get_team_object(
|
||||
team_id=user_api_key_auth.team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=user_api_key_auth.parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
if team_obj is None:
|
||||
return []
|
||||
|
||||
all_agents: List[str] = []
|
||||
|
||||
# 1. Get agents from object_permission (native permissions)
|
||||
object_permissions = team_obj.object_permission
|
||||
if object_permissions is not None:
|
||||
# Get direct agents
|
||||
direct_agents = object_permissions.agents or []
|
||||
|
||||
# Get agents from access groups
|
||||
access_group_agents = (
|
||||
await AgentRequestHandler._get_agents_from_access_groups(
|
||||
object_permissions.agent_access_groups or []
|
||||
)
|
||||
)
|
||||
|
||||
all_agents = direct_agents + access_group_agents
|
||||
|
||||
# 2. Also include agents from team's access_group_ids (unified access groups)
|
||||
team_access_group_ids = team_obj.access_group_ids or []
|
||||
if team_access_group_ids:
|
||||
from litellm.proxy.auth.auth_checks import (
|
||||
_get_agent_ids_from_access_groups,
|
||||
)
|
||||
|
||||
unified_agents = await _get_agent_ids_from_access_groups(
|
||||
access_group_ids=team_access_group_ids,
|
||||
)
|
||||
all_agents.extend(unified_agents)
|
||||
|
||||
return list(set(all_agents))
|
||||
except Exception as e:
|
||||
# litellm-dashboard is the default UI team and will never have agents;
|
||||
# skip noisy warnings for it.
|
||||
if user_api_key_auth.team_id != UI_TEAM_ID:
|
||||
verbose_logger.warning(
|
||||
f"Failed to get allowed agents for team: {str(e)}"
|
||||
)
|
||||
return []
|
||||
|
||||
@staticmethod
|
||||
def _get_config_agent_ids_for_access_groups(
|
||||
config_agents: List, access_groups: List[str]
|
||||
) -> Set[str]:
|
||||
"""
|
||||
Helper to get agent_ids from config-loaded agents that match any of the given access groups.
|
||||
"""
|
||||
server_ids: Set[str] = set()
|
||||
for agent in config_agents:
|
||||
agent_access_groups = getattr(agent, "agent_access_groups", None)
|
||||
if agent_access_groups:
|
||||
if any(group in agent_access_groups for group in access_groups):
|
||||
server_ids.add(agent.agent_id)
|
||||
return server_ids
|
||||
|
||||
@staticmethod
|
||||
async def _get_db_agent_ids_for_access_groups(
|
||||
prisma_client, access_groups: List[str]
|
||||
) -> Set[str]:
|
||||
"""
|
||||
Helper to get agent_ids from DB agents that match any of the given access groups.
|
||||
"""
|
||||
agent_ids: Set[str] = set()
|
||||
if access_groups and prisma_client is not None:
|
||||
try:
|
||||
agents = await prisma_client.db.litellm_agentstable.find_many(
|
||||
where={"agent_access_groups": {"hasSome": access_groups}}
|
||||
)
|
||||
for agent in agents:
|
||||
agent_ids.add(agent.agent_id)
|
||||
except Exception as e:
|
||||
verbose_logger.debug(f"Error getting agents from access groups: {e}")
|
||||
return agent_ids
|
||||
|
||||
@staticmethod
|
||||
async def _get_agents_from_access_groups(
|
||||
access_groups: List[str],
|
||||
) -> List[str]:
|
||||
"""
|
||||
Resolve agent access groups to agent IDs by querying BOTH the agent table (DB) AND config-loaded agents.
|
||||
"""
|
||||
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
try:
|
||||
# Use the helper for config-loaded agents
|
||||
agent_ids = AgentRequestHandler._get_config_agent_ids_for_access_groups(
|
||||
global_agent_registry.agent_list, access_groups
|
||||
)
|
||||
|
||||
# Use the helper for DB agents
|
||||
db_agent_ids = (
|
||||
await AgentRequestHandler._get_db_agent_ids_for_access_groups(
|
||||
prisma_client, access_groups
|
||||
)
|
||||
)
|
||||
agent_ids.update(db_agent_ids)
|
||||
|
||||
return list(agent_ids)
|
||||
except Exception as e:
|
||||
verbose_logger.warning(f"Failed to get agents from access groups: {str(e)}")
|
||||
return []
|
||||
|
||||
@staticmethod
|
||||
async def get_agent_access_groups(
|
||||
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
|
||||
) -> List[str]:
|
||||
"""
|
||||
Get list of agent access groups for the given user/key based on permissions.
|
||||
"""
|
||||
access_groups: List[str] = []
|
||||
access_groups_for_key = (
|
||||
await AgentRequestHandler._get_agent_access_groups_for_key(
|
||||
user_api_key_auth
|
||||
)
|
||||
)
|
||||
access_groups_for_team = (
|
||||
await AgentRequestHandler._get_agent_access_groups_for_team(
|
||||
user_api_key_auth
|
||||
)
|
||||
)
|
||||
|
||||
# If team has access groups, then key must have a subset of the team's access groups
|
||||
if len(access_groups_for_team) > 0:
|
||||
for access_group in access_groups_for_key:
|
||||
if access_group in access_groups_for_team:
|
||||
access_groups.append(access_group)
|
||||
else:
|
||||
access_groups = access_groups_for_key
|
||||
|
||||
return list(set(access_groups))
|
||||
|
||||
@staticmethod
|
||||
async def _get_agent_access_groups_for_key(
|
||||
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
|
||||
) -> List[str]:
|
||||
"""Get agent access groups for the key."""
|
||||
from litellm.proxy.auth.auth_checks import get_object_permission
|
||||
from litellm.proxy.proxy_server import (
|
||||
prisma_client,
|
||||
proxy_logging_obj,
|
||||
user_api_key_cache,
|
||||
)
|
||||
|
||||
if user_api_key_auth is None:
|
||||
return []
|
||||
|
||||
if user_api_key_auth.object_permission_id is None:
|
||||
return []
|
||||
|
||||
if prisma_client is None:
|
||||
verbose_logger.debug("prisma_client is None")
|
||||
return []
|
||||
|
||||
try:
|
||||
key_object_permission = await get_object_permission(
|
||||
object_permission_id=user_api_key_auth.object_permission_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=user_api_key_auth.parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
if key_object_permission is None:
|
||||
return []
|
||||
|
||||
return key_object_permission.agent_access_groups or []
|
||||
except Exception as e:
|
||||
verbose_logger.warning(
|
||||
f"Failed to get agent access groups for key: {str(e)}"
|
||||
)
|
||||
return []
|
||||
|
||||
@staticmethod
|
||||
async def _get_agent_access_groups_for_team(
|
||||
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
|
||||
) -> List[str]:
|
||||
"""Get agent access groups for the team."""
|
||||
from litellm.proxy.auth.auth_checks import get_team_object
|
||||
from litellm.proxy.proxy_server import (
|
||||
prisma_client,
|
||||
proxy_logging_obj,
|
||||
user_api_key_cache,
|
||||
)
|
||||
|
||||
if user_api_key_auth is None:
|
||||
return []
|
||||
|
||||
if user_api_key_auth.team_id is None:
|
||||
return []
|
||||
|
||||
if prisma_client is None:
|
||||
verbose_logger.debug("prisma_client is None")
|
||||
return []
|
||||
|
||||
try:
|
||||
team_obj: Optional[LiteLLM_TeamTable] = await get_team_object(
|
||||
team_id=user_api_key_auth.team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=user_api_key_auth.parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
if team_obj is None:
|
||||
verbose_logger.debug("team_obj is None")
|
||||
return []
|
||||
|
||||
object_permissions = team_obj.object_permission
|
||||
if object_permissions is None:
|
||||
return []
|
||||
|
||||
return object_permissions.agent_access_groups or []
|
||||
except Exception as e:
|
||||
verbose_logger.warning(
|
||||
f"Failed to get agent access groups for team: {str(e)}"
|
||||
)
|
||||
return []
|
||||
Reference in New Issue
Block a user