tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-05 14:23:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #21022 from BerriAI/litellm_unified_ag

Browse Source 
[Feature] Access Groups
This commit is contained in:
yuneng-jiang
2026-02-12 15:34:38 -08:00
committed by GitHub
parent 5f40f93846 a37623945d
commit 2864ce73da
14 changed files with 1002 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm-proxy-extras/dist/litellm_proxy_extras-0.4.36-py3-none-any.whl
Vendored
BIN
View File
Binary file not shown.
litellm-proxy-extras/dist/litellm_proxy_extras-0.4.36.tar.gz
Vendored
BIN
View File
Binary file not shown.
litellm-proxy-extras/litellm_proxy_extras/migrations/20260212143306_add_access_group_table/migration.sql
+33
View File
@@ -0,0 +1,33 @@
-- AlterTable
ALTER TABLE "LiteLLM_DeletedTeamTable" ADD COLUMN "access_group_ids" TEXT[] DEFAULT ARRAY[]::TEXT[];
-- AlterTable
ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN "access_group_ids" TEXT[] DEFAULT ARRAY[]::TEXT[];
-- AlterTable
ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN "access_group_ids" TEXT[] DEFAULT ARRAY[]::TEXT[];
-- AlterTable
ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN "access_group_ids" TEXT[] DEFAULT ARRAY[]::TEXT[];
-- CreateTable
CREATE TABLE "LiteLLM_AccessGroupTable" (
"access_group_id" TEXT NOT NULL,
"access_group_name" TEXT NOT NULL,
"description" TEXT,
"access_model_ids" TEXT[] DEFAULT ARRAY[]::TEXT[],
"access_mcp_server_ids" TEXT[] DEFAULT ARRAY[]::TEXT[],
"access_agent_ids" TEXT[] DEFAULT ARRAY[]::TEXT[],
"assigned_team_ids" TEXT[] DEFAULT ARRAY[]::TEXT[],
"assigned_key_ids" TEXT[] DEFAULT ARRAY[]::TEXT[],
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"created_by" TEXT,
"updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"updated_by" TEXT,
CONSTRAINT "LiteLLM_AccessGroupTable_pkey" PRIMARY KEY ("access_group_id")
);
-- CreateIndex
CREATE UNIQUE INDEX "LiteLLM_AccessGroupTable_access_group_name_key" ON "LiteLLM_AccessGroupTable"("access_group_name");
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
+24
View File
@@ -128,6 +128,7 @@ model LiteLLM_TeamTable {
model_max_budget Json @default("{}")
router_settings Json? @default("{}")
team_member_permissions String[] @default([])
access_group_ids String[] @default([])
policies String[] @default([])
model_id Int? @unique // id for LiteLLM_ModelTable -> stores team-level model aliases
allow_team_guardrail_config Boolean @default(false) // if true, team admin can configure guardrails for this team
@@ -161,6 +162,7 @@ model LiteLLM_DeletedTeamTable {
model_max_budget Json @default("{}")
router_settings Json? @default("{}")
team_member_permissions String[] @default([])
access_group_ids String[] @default([])
policies String[] @default([])
model_id Int? // id for LiteLLM_ModelTable -> stores team-level model aliases
allow_team_guardrail_config Boolean @default(false)
@@ -293,6 +295,7 @@ model LiteLLM_VerificationToken {
allowed_cache_controls String[] @default([])
allowed_routes String[] @default([])
policies String[] @default([])
access_group_ids String[] @default([])
model_spend Json @default("{}")
model_max_budget Json @default("{}")
budget_id String?
@@ -348,6 +351,7 @@ model LiteLLM_DeletedVerificationToken {
allowed_cache_controls String[] @default([])
allowed_routes String[] @default([])
policies String[] @default([])
access_group_ids String[] @default([])
model_spend Json @default("{}")
model_max_budget Json @default("{}")
router_settings Json? @default("{}")
@@ -920,3 +924,23 @@ model LiteLLM_PolicyAttachmentTable {
updated_at DateTime @default(now()) @updatedAt
updated_by String?
}
//Unified Access Groups table for storing unified access groups
model LiteLLM_AccessGroupTable {
access_group_id String @id @default(uuid())
access_group_name String @unique
description String?
// Resource memberships - explicit arrays per type
access_model_ids String[] @default([])
access_mcp_server_ids String[] @default([])
access_agent_ids String[] @default([])
assigned_team_ids String[] @default([])
assigned_key_ids String[] @default([])
created_at DateTime @default(now())
created_by String?
updated_at DateTime @default(now()) @updatedAt
updated_by String?
}
litellm-proxy-extras/pyproject.toml
+2 -2
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "litellm-proxy-extras"
version = "0.4.35"
version = "0.4.36"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
authors = ["BerriAI"]
readme = "README.md"
@@ -22,7 +22,7 @@ requires = ["poetry-core"]
build-backend = "poetry.core.masonry.api"
[tool.commitizen]
version = "0.4.35"
version = "0.4.36"
version_files = [
"pyproject.toml:version",
"../requirements.txt:litellm-proxy-extras==",
litellm/proxy/_types.py
+4
View File
@@ -893,6 +893,7 @@ class KeyRequestBase(GenerateRequestBase):
Literal["guaranteed_throughput", "best_effort_throughput", "dynamic"]
] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating tpm
router_settings: Optional[UpdateRouterConfig] = None
access_group_ids: Optional[List[str]] = None
class LiteLLMKeyType(str, enum.Enum):
@@ -1502,6 +1503,7 @@ class TeamBase(LiteLLMPydanticObjectBase):
models: list = []
blocked: bool = False
router_settings: Optional[dict] = None
access_group_ids: Optional[List[str]] = None
class NewTeamRequest(TeamBase):
@@ -1589,6 +1591,7 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase):
model_tpm_limit: Optional[Dict[str, int]] = None
allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None
router_settings: Optional[dict] = None
access_group_ids: Optional[List[str]] = None
class ResetTeamBudgetRequest(LiteLLMPydanticObjectBase):
@@ -2177,6 +2180,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase):
updated_by: Optional[str] = None
object_permission_id: Optional[str] = None
object_permission: Optional[LiteLLM_ObjectPermissionTable] = None
access_group_ids: Optional[List[str]] = None
rotation_count: Optional[int] = 0 # Number of times key has been rotated
auto_rotate: Optional[bool] = False # Whether this key should be auto-rotated
rotation_interval: Optional[str] = None # How often to rotate (e.g., "30d", "90d")
litellm/proxy/management_endpoints/access_group_endpoints.py
+264
View File
@@ -0,0 +1,264 @@
from typing import List
from fastapi import APIRouter, Depends, HTTPException, status
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.proxy.utils import get_prisma_client_or_throw
from litellm.types.access_group import (
AccessGroupCreateRequest,
AccessGroupResponse,
AccessGroupUpdateRequest,
)
router = APIRouter(
tags=["access group management"],
)
def _require_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> None:
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail={"error": CommonProxyErrors.not_allowed_access.value},
)
def _record_to_response(record) -> AccessGroupResponse:
return AccessGroupResponse(
access_group_id=record.access_group_id,
access_group_name=record.access_group_name,
description=record.description,
access_model_ids=record.access_model_ids,
access_mcp_server_ids=record.access_mcp_server_ids,
access_agent_ids=record.access_agent_ids,
assigned_team_ids=record.assigned_team_ids,
assigned_key_ids=record.assigned_key_ids,
created_at=record.created_at,
created_by=record.created_by,
updated_at=record.updated_at,
updated_by=record.updated_by,
)
@router.post(
"/v1/access_group",
response_model=AccessGroupResponse,
status_code=status.HTTP_201_CREATED,
)
async def create_access_group(
data: AccessGroupCreateRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
) -> AccessGroupResponse:
_require_proxy_admin(user_api_key_dict)
prisma_client = get_prisma_client_or_throw(CommonProxyErrors.db_not_connected_error.value)
existing = await prisma_client.db.litellm_accessgrouptable.find_unique(
where={"access_group_name": data.access_group_name}
)
if existing is not None:
raise HTTPException(
status_code=status.HTTP_409_CONFLICT,
detail=f"Access group '{data.access_group_name}' already exists",
)
try:
record = await prisma_client.db.litellm_accessgrouptable.create(
data={
"access_group_name": data.access_group_name,
"description": data.description,
"access_model_ids": data.access_model_ids or [],
"access_mcp_server_ids": data.access_mcp_server_ids or [],
"access_agent_ids": data.access_agent_ids or [],
"assigned_team_ids": data.assigned_team_ids or [],
"assigned_key_ids": data.assigned_key_ids or [],
"created_by": user_api_key_dict.user_id,
"updated_by": user_api_key_dict.user_id,
}
)
except Exception as e:
# Race condition: another request created the same name between find_unique and create.
# Prisma raises UniqueViolationError (P2002) or similar for unique constraint.
if "unique constraint" in str(e).lower() or "P2002" in str(e):
raise HTTPException(
status_code=status.HTTP_409_CONFLICT,
detail=f"Access group '{data.access_group_name}' already exists",
)
raise
return _record_to_response(record)
@router.get(
"/v1/access_group",
response_model=List[AccessGroupResponse],
)
async def list_access_groups(
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
) -> List[AccessGroupResponse]:
_require_proxy_admin(user_api_key_dict)
prisma_client = get_prisma_client_or_throw(CommonProxyErrors.db_not_connected_error.value)
records = await prisma_client.db.litellm_accessgrouptable.find_many(
order={"created_at": "desc"}
)
return [_record_to_response(r) for r in records]
@router.get(
"/v1/access_group/{access_group_id}",
response_model=AccessGroupResponse,
)
async def get_access_group(
access_group_id: str,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
) -> AccessGroupResponse:
_require_proxy_admin(user_api_key_dict)
prisma_client = get_prisma_client_or_throw(CommonProxyErrors.db_not_connected_error.value)
record = await prisma_client.db.litellm_accessgrouptable.find_unique(
where={"access_group_id": access_group_id}
)
if record is None:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail=f"Access group '{access_group_id}' not found",
)
return _record_to_response(record)
@router.put(
"/v1/access_group/{access_group_id}",
response_model=AccessGroupResponse,
)
async def update_access_group(
access_group_id: str,
data: AccessGroupUpdateRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
) -> AccessGroupResponse:
_require_proxy_admin(user_api_key_dict)
prisma_client = get_prisma_client_or_throw(CommonProxyErrors.db_not_connected_error.value)
existing = await prisma_client.db.litellm_accessgrouptable.find_unique(
where={"access_group_id": access_group_id}
)
if existing is None:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail=f"Access group '{access_group_id}' not found",
)
update_data: dict = {"updated_by": user_api_key_dict.user_id}
for field, value in data.model_dump(exclude_unset=True).items():
update_data[field] = value
record = await prisma_client.db.litellm_accessgrouptable.update(
where={"access_group_id": access_group_id},
data=update_data,
)
return _record_to_response(record)
@router.delete(
"/v1/access_group/{access_group_id}",
status_code=status.HTTP_204_NO_CONTENT,
)
async def delete_access_group(
access_group_id: str,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
) -> None:
_require_proxy_admin(user_api_key_dict)
prisma_client = get_prisma_client_or_throw(CommonProxyErrors.db_not_connected_error.value)
try:
async with prisma_client.db.tx() as tx:
existing = await tx.litellm_accessgrouptable.find_unique(
where={"access_group_id": access_group_id}
)
if existing is None:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail=f"Access group '{access_group_id}' not found",
)
# Remove access_group_id from teams and keys that reference it
teams_with_group = await tx.litellm_teamtable.find_many(
where={"access_group_ids": {"hasSome": [access_group_id]}}
)
for team in teams_with_group:
updated_ids = [tid for tid in (team.access_group_ids or []) if tid != access_group_id]
await tx.litellm_teamtable.update(
where={"team_id": team.team_id},
data={"access_group_ids": updated_ids},
)
keys_with_group = await tx.litellm_verificationtoken.find_many(
where={"access_group_ids": {"hasSome": [access_group_id]}}
)
for key in keys_with_group:
updated_ids = [kid for kid in (key.access_group_ids or []) if kid != access_group_id]
await tx.litellm_verificationtoken.update(
where={"token": key.token},
data={"access_group_ids": updated_ids},
)
await tx.litellm_accessgrouptable.delete(
where={"access_group_id": access_group_id}
)
except HTTPException:
raise
except Exception as e:
verbose_proxy_logger.exception(
"delete_access_group failed: access_group_id=%s error=%s",
access_group_id,
e,
)
if PrismaDBExceptionHandler.is_database_connection_error(e):
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
detail=CommonProxyErrors.db_not_connected_error.value,
)
if "P2025" in str(e) or ("record" in str(e).lower() and "not found" in str(e).lower()):
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail=f"Access group '{access_group_id}' not found",
)
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail="Failed to delete access group. Please try again.",
)
# Alias routes for /v1/unified_access_group
router.add_api_route(
"/v1/unified_access_group",
create_access_group,
methods=["POST"],
response_model=AccessGroupResponse,
status_code=status.HTTP_201_CREATED,
)
router.add_api_route(
"/v1/unified_access_group",
list_access_groups,
methods=["GET"],
response_model=List[AccessGroupResponse],
)
router.add_api_route(
"/v1/unified_access_group/{access_group_id}",
get_access_group,
methods=["GET"],
response_model=AccessGroupResponse,
)
router.add_api_route(
"/v1/unified_access_group/{access_group_id}",
update_access_group,
methods=["PUT"],
response_model=AccessGroupResponse,
)
router.add_api_route(
"/v1/unified_access_group/{access_group_id}",
delete_access_group,
methods=["DELETE"],
status_code=status.HTTP_204_NO_CONTENT,
)
litellm/proxy/proxy_server.py
+4
View File
@@ -393,6 +393,9 @@ from litellm.proxy.management_endpoints.tag_management_endpoints import (
from litellm.proxy.management_endpoints.team_callback_endpoints import (
router as team_callback_router,
)
from litellm.proxy.management_endpoints.access_group_endpoints import (
router as access_group_router,
)
from litellm.proxy.management_endpoints.team_endpoints import router as team_router
from litellm.proxy.management_endpoints.team_endpoints import (
update_team,
@@ -11998,6 +12001,7 @@ app.include_router(enterprise_router)
app.include_router(ui_discovery_endpoints_router)
app.include_router(agent_endpoints_router)
app.include_router(a2a_router)
app.include_router(access_group_router)
########################################################
# MCP Server
########################################################
litellm/proxy/schema.prisma
+24
View File
@@ -128,6 +128,7 @@ model LiteLLM_TeamTable {
model_max_budget Json @default("{}")
router_settings Json? @default("{}")
team_member_permissions String[] @default([])
access_group_ids String[] @default([])
policies String[] @default([])
model_id Int? @unique // id for LiteLLM_ModelTable -> stores team-level model aliases
litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id])
@@ -160,6 +161,7 @@ model LiteLLM_DeletedTeamTable {
model_max_budget Json @default("{}")
router_settings Json? @default("{}")
team_member_permissions String[] @default([])
access_group_ids String[] @default([])
policies String[] @default([])
model_id Int? // id for LiteLLM_ModelTable -> stores team-level model aliases
@@ -291,6 +293,7 @@ model LiteLLM_VerificationToken {
allowed_cache_controls String[] @default([])
allowed_routes String[] @default([])
policies String[] @default([])
access_group_ids String[] @default([])
model_spend Json @default("{}")
model_max_budget Json @default("{}")
budget_id String?
@@ -346,6 +349,7 @@ model LiteLLM_DeletedVerificationToken {
allowed_cache_controls String[] @default([])
allowed_routes String[] @default([])
policies String[] @default([])
access_group_ids String[] @default([])
model_spend Json @default("{}")
model_max_budget Json @default("{}")
router_settings Json? @default("{}")
@@ -917,3 +921,23 @@ model LiteLLM_PolicyAttachmentTable {
updated_at DateTime @default(now()) @updatedAt
updated_by String?
}
//Unified Access Groups table for storing unified access groups
model LiteLLM_AccessGroupTable {
access_group_id String @id @default(uuid())
access_group_name String @unique
description String?
// Resource memberships - explicit arrays per type
access_model_ids String[] @default([])
access_mcp_server_ids String[] @default([])
access_agent_ids String[] @default([])
assigned_team_ids String[] @default([])
assigned_key_ids String[] @default([])
created_at DateTime @default(now())
created_by String?
updated_at DateTime @default(now()) @updatedAt
updated_by String?
}
litellm/types/access_group.py
+38
View File
@@ -0,0 +1,38 @@
from datetime import datetime
from typing import List, Optional
from pydantic import BaseModel
class AccessGroupCreateRequest(BaseModel):
access_group_name: str
description: Optional[str] = None
access_model_ids: Optional[List[str]] = None
access_mcp_server_ids: Optional[List[str]] = None
access_agent_ids: Optional[List[str]] = None
assigned_team_ids: Optional[List[str]] = None
assigned_key_ids: Optional[List[str]] = None
class AccessGroupUpdateRequest(BaseModel):
description: Optional[str] = None
access_model_ids: Optional[List[str]] = None
access_mcp_server_ids: Optional[List[str]] = None
access_agent_ids: Optional[List[str]] = None
assigned_team_ids: Optional[List[str]] = None
assigned_key_ids: Optional[List[str]] = None
class AccessGroupResponse(BaseModel):
access_group_id: str
access_group_name: str
description: Optional[str] = None
access_model_ids: List[str]
access_mcp_server_ids: List[str]
access_agent_ids: List[str]
assigned_team_ids: List[str]
assigned_key_ids: List[str]
created_at: datetime
created_by: Optional[str] = None
updated_at: datetime
updated_by: Optional[str] = None
pyproject.toml
+1 -1
View File
@@ -61,7 +61,7 @@ boto3 = { version = "1.40.76", optional = true }
redisvl = {version = "^0.4.1", optional = true, markers = "python_version >= '3.9' and python_version < '3.14'"}
mcp = {version = ">=1.25.0,<2.0.0", optional = true, python = ">=3.10"}
a2a-sdk = {version = "^0.3.22", optional = true, python = ">=3.10"}
litellm-proxy-extras = {version = "0.4.35", optional = true}
litellm-proxy-extras = {version = "0.4.36", optional = true}
rich = {version = "13.7.1", optional = true}
litellm-enterprise = {version = "0.1.31", optional = true}
diskcache = {version = "^5.6.1", optional = true}
requirements.txt
+1 -1
View File
@@ -55,7 +55,7 @@ sentry_sdk==2.21.0 # for sentry error handling
detect-secrets==1.5.0 # Enterprise - secret detection / masking in LLM requests
cryptography==44.0.1
tzdata==2025.1 # IANA time zone database
litellm-proxy-extras==0.4.35 # for proxy extras - e.g. prisma migrations
litellm-proxy-extras==0.4.36 # for proxy extras - e.g. prisma migrations
llm-sandbox==0.3.31 # for skill execution in sandbox
### LITELLM PACKAGE DEPENDENCIES
python-dotenv==1.0.1 # for env
schema.prisma
+24 -1
View File
@@ -128,6 +128,7 @@ model LiteLLM_TeamTable {
model_max_budget Json @default("{}")
router_settings Json? @default("{}")
team_member_permissions String[] @default([])
access_group_ids String[] @default([])
policies String[] @default([])
model_id Int? @unique // id for LiteLLM_ModelTable -> stores team-level model aliases
allow_team_guardrail_config Boolean @default(false) // if true, team admin can configure guardrails for this team
@@ -161,6 +162,7 @@ model LiteLLM_DeletedTeamTable {
model_max_budget Json @default("{}")
router_settings Json? @default("{}")
team_member_permissions String[] @default([])
access_group_ids String[] @default([])
policies String[] @default([])
model_id Int? // id for LiteLLM_ModelTable -> stores team-level model aliases
allow_team_guardrail_config Boolean @default(false)
@@ -293,6 +295,7 @@ model LiteLLM_VerificationToken {
allowed_cache_controls String[] @default([])
allowed_routes String[] @default([])
policies String[] @default([])
access_group_ids String[] @default([])
model_spend Json @default("{}")
model_max_budget Json @default("{}")
budget_id String?
@@ -348,6 +351,7 @@ model LiteLLM_DeletedVerificationToken {
allowed_cache_controls String[] @default([])
allowed_routes String[] @default([])
policies String[] @default([])
access_group_ids String[] @default([])
model_spend Json @default("{}")
model_max_budget Json @default("{}")
router_settings Json? @default("{}")
@@ -363,7 +367,6 @@ model LiteLLM_DeletedVerificationToken {
rotation_interval String?
last_rotation_at DateTime?
key_rotation_at DateTime?
// Deletion metadata
deleted_at DateTime @default(now()) @map("deleted_at")
deleted_by String? @map("deleted_by") // User who deleted the key
@@ -919,3 +922,23 @@ model LiteLLM_PolicyAttachmentTable {
updated_at DateTime @default(now()) @updatedAt
updated_by String?
}
//Unified Access Groups table for storing unified access groups
model LiteLLM_AccessGroupTable {
access_group_id String @id @default(uuid())
access_group_name String @unique
description String?
// Resource memberships - explicit arrays per type
access_model_ids String[] @default([])
access_mcp_server_ids String[] @default([])
access_agent_ids String[] @default([])
assigned_team_ids String[] @default([])
assigned_key_ids String[] @default([])
created_at DateTime @default(now())
created_by String?
updated_at DateTime @default(now()) @updatedAt
updated_by String?
}
tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py
+583
View File
@@ -0,0 +1,583 @@
"""
Tests for access group management endpoints.
"""
import os
import sys
import types
from contextlib import asynccontextmanager
from datetime import datetime
from unittest.mock import AsyncMock, MagicMock
import pytest
from fastapi.testclient import TestClient
from prisma.errors import PrismaError
import litellm.proxy.proxy_server as ps
from litellm.proxy.proxy_server import app
from litellm.proxy._types import (
CommonProxyErrors,
LitellmUserRoles,
UserAPIKeyAuth,
)
sys.path.insert(0, os.path.abspath("../../../"))
def _make_access_group_record(
access_group_id: str = "ag-123",
access_group_name: str = "test-group",
description: str | None = "Test description",
access_model_ids: list | None = None,
access_mcp_server_ids: list | None = None,
access_agent_ids: list | None = None,
assigned_team_ids: list | None = None,
assigned_key_ids: list | None = None,
created_by: str | None = "admin-user",
updated_by: str | None = "admin-user",
created_at: datetime | None = None,
):
record = MagicMock()
record.access_group_id = access_group_id
record.access_group_name = access_group_name
record.description = description
record.access_model_ids = access_model_ids or []
record.access_mcp_server_ids = access_mcp_server_ids or []
record.access_agent_ids = access_agent_ids or []
record.assigned_team_ids = assigned_team_ids or []
record.assigned_key_ids = assigned_key_ids or []
record.created_at = created_at or datetime.now()
record.created_by = created_by
record.updated_at = datetime.now()
record.updated_by = updated_by
return record
@pytest.fixture
def client_and_mocks(monkeypatch):
"""Setup mock prisma and admin auth for access group endpoints."""
mock_access_group_table = MagicMock()
mock_prisma = MagicMock()
def _create_side_effect(*, data):
return _make_access_group_record(
access_group_id="ag-new",
access_group_name=data.get("access_group_name", "new"),
description=data.get("description"),
access_model_ids=data.get("access_model_ids", []),
access_mcp_server_ids=data.get("access_mcp_server_ids", []),
access_agent_ids=data.get("access_agent_ids", []),
assigned_team_ids=data.get("assigned_team_ids", []),
assigned_key_ids=data.get("assigned_key_ids", []),
created_by=data.get("created_by"),
updated_by=data.get("updated_by"),
)
mock_access_group_table.create = AsyncMock(side_effect=_create_side_effect)
mock_access_group_table.find_unique = AsyncMock(return_value=None)
mock_access_group_table.find_many = AsyncMock(return_value=[])
mock_access_group_table.update = AsyncMock(side_effect=lambda *, where, data: _make_access_group_record(
access_group_id=where.get("access_group_id", "ag-123"),
access_group_name=data.get("access_group_name", "updated"),
description=data.get("description"),
access_model_ids=data.get("access_model_ids", []),
access_mcp_server_ids=data.get("access_mcp_server_ids", []),
access_agent_ids=data.get("access_agent_ids", []),
assigned_team_ids=data.get("assigned_team_ids", []),
assigned_key_ids=data.get("assigned_key_ids", []),
updated_by=data.get("updated_by"),
))
mock_access_group_table.delete = AsyncMock(return_value=None)
mock_team_table = MagicMock()
mock_team_table.find_many = AsyncMock(return_value=[])
mock_team_table.update = AsyncMock(return_value=None)
mock_key_table = MagicMock()
mock_key_table.find_many = AsyncMock(return_value=[])
mock_key_table.update = AsyncMock(return_value=None)
@asynccontextmanager
async def mock_tx():
tx = types.SimpleNamespace(
litellm_accessgrouptable=mock_access_group_table,
litellm_teamtable=mock_team_table,
litellm_verificationtoken=mock_key_table,
)
yield tx
mock_db = types.SimpleNamespace(
litellm_accessgrouptable=mock_access_group_table,
litellm_teamtable=mock_team_table,
litellm_verificationtoken=mock_key_table,
tx=mock_tx,
)
mock_prisma.db = mock_db
monkeypatch.setattr(ps, "prisma_client", mock_prisma)
admin_user = UserAPIKeyAuth(
user_id="admin_user",
user_role=LitellmUserRoles.PROXY_ADMIN,
)
app.dependency_overrides[ps.user_api_key_auth] = lambda: admin_user
client = TestClient(app)
yield client, mock_prisma, mock_access_group_table
app.dependency_overrides.clear()
monkeypatch.setattr(ps, "prisma_client", ps.prisma_client)
# Paths for primary and alias endpoints (alias: /v1/unified_access_group)
ACCESS_GROUP_PATHS = ["/v1/access_group", "/v1/unified_access_group"]
# ---------------------------------------------------------------------------
# CREATE
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("base_path", ACCESS_GROUP_PATHS)
@pytest.mark.parametrize(
"payload",
[
{"access_group_name": "group-a"},
{
"access_group_name": "group-b",
"description": "Group B description",
"access_model_ids": ["model-1"],
"access_mcp_server_ids": ["mcp-1"],
"assigned_team_ids": ["team-1"],
},
],
)
def test_create_access_group_success(client_and_mocks, base_path, payload):
"""Create access group with various payloads returns 201."""
client, _, mock_table = client_and_mocks
resp = client.post(base_path, json=payload)
assert resp.status_code == 201
body = resp.json()
assert body["access_group_name"] == payload["access_group_name"]
assert body.get("access_group_id") is not None
mock_table.create.assert_awaited_once()
def test_create_access_group_duplicate_name_conflict(client_and_mocks):
"""Create with duplicate name returns 409."""
client, _, mock_table = client_and_mocks
existing = _make_access_group_record(access_group_name="existing-group")
mock_table.find_unique = AsyncMock(return_value=existing)
resp = client.post("/v1/access_group", json={"access_group_name": "existing-group"})
assert resp.status_code == 409
assert "already exists" in resp.json()["detail"]
@pytest.mark.parametrize(
"error_message",
[
"Unique constraint failed on the fields: (`access_group_name`)",
"P2002: Unique constraint failed",
"unique constraint violation",
],
)
def test_create_access_group_race_condition_returns_409(client_and_mocks, error_message):
"""Create race condition: Prisma unique constraint surfaces as 409, not 500."""
client, _, mock_table = client_and_mocks
mock_table.find_unique = AsyncMock(return_value=None)
mock_table.create = AsyncMock(side_effect=Exception(error_message))
resp = client.post("/v1/access_group", json={"access_group_name": "race-group"})
assert resp.status_code == 409
assert "already exists" in resp.json()["detail"]
@pytest.mark.parametrize("user_role", [LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY])
def test_create_access_group_forbidden_non_admin(client_and_mocks, user_role):
"""Non-admin users cannot create access groups."""
client, _, _ = client_and_mocks
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_id="regular_user",
user_role=user_role,
)
resp = client.post("/v1/access_group", json={"access_group_name": "forbidden"})
assert resp.status_code == 403
assert resp.json()["detail"]["error"] == CommonProxyErrors.not_allowed_access.value
def test_create_access_group_validation_missing_name(client_and_mocks):
"""Create with missing access_group_name returns 422."""
client, _, _ = client_and_mocks
resp = client.post("/v1/access_group", json={})
assert resp.status_code == 422
def test_create_access_group_500_on_non_constraint_prisma_error(client_and_mocks):
"""Create with non-unique-constraint Prisma error returns 500."""
client, _, mock_table = client_and_mocks
mock_table.find_unique = AsyncMock(return_value=None)
mock_table.create = AsyncMock(side_effect=Exception("Some other database error"))
# Use raise_server_exceptions=False so unhandled exceptions become 500 responses
test_client = TestClient(app, raise_server_exceptions=False)
resp = test_client.post("/v1/access_group", json={"access_group_name": "test-group"})
assert resp.status_code == 500
# ---------------------------------------------------------------------------
# LIST
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("base_path", ACCESS_GROUP_PATHS)
def test_list_access_groups_success_empty(client_and_mocks, base_path):
"""List access groups returns empty list when none exist."""
client, _, mock_table = client_and_mocks
resp = client.get(base_path)
assert resp.status_code == 200
assert resp.json() == []
mock_table.find_many.assert_awaited_once()
@pytest.mark.parametrize("base_path", ACCESS_GROUP_PATHS)
def test_list_access_groups_success_with_items(client_and_mocks, base_path):
"""List access groups returns items when they exist."""
client, _, mock_table = client_and_mocks
records = [
_make_access_group_record(access_group_id="ag-1", access_group_name="group-1"),
_make_access_group_record(access_group_id="ag-2", access_group_name="group-2"),
]
mock_table.find_many = AsyncMock(return_value=records)
resp = client.get(base_path)
assert resp.status_code == 200
body = resp.json()
assert len(body) == 2
assert body[0]["access_group_name"] == "group-1"
assert body[1]["access_group_name"] == "group-2"
@pytest.mark.parametrize("base_path", ACCESS_GROUP_PATHS)
def test_list_access_groups_ordered_by_created_at_desc(client_and_mocks, base_path):
"""List access groups calls find_many with created_at desc order."""
client, _, mock_table = client_and_mocks
older = datetime(2025, 1, 1, 12, 0, 0)
newer = datetime(2025, 1, 2, 12, 0, 0)
records = [
_make_access_group_record(
access_group_id="ag-newer",
access_group_name="newer-group",
created_at=newer,
),
_make_access_group_record(
access_group_id="ag-older",
access_group_name="older-group",
created_at=older,
),
]
mock_table.find_many = AsyncMock(return_value=records)
resp = client.get(base_path)
assert resp.status_code == 200
body = resp.json()
assert len(body) == 2
# Mock returns newest first (simulating Prisma order desc)
assert body[0]["access_group_name"] == "newer-group"
assert body[1]["access_group_name"] == "older-group"
mock_table.find_many.assert_awaited_once_with(order={"created_at": "desc"})
@pytest.mark.parametrize("user_role", [LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY])
def test_list_access_groups_forbidden_non_admin(client_and_mocks, user_role):
"""Non-admin users cannot list access groups."""
client, _, _ = client_and_mocks
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_id="regular_user",
user_role=user_role,
)
resp = client.get("/v1/access_group")
assert resp.status_code == 403
assert resp.json()["detail"]["error"] == CommonProxyErrors.not_allowed_access.value
# ---------------------------------------------------------------------------
# GET
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("base_path", ACCESS_GROUP_PATHS)
@pytest.mark.parametrize("access_group_id", ["ag-123", "ag-other-id"])
def test_get_access_group_success(client_and_mocks, base_path, access_group_id):
"""Get access group by id returns record when found."""
client, _, mock_table = client_and_mocks
record = _make_access_group_record(access_group_id=access_group_id)
mock_table.find_unique = AsyncMock(return_value=record)
resp = client.get(f"{base_path}/{access_group_id}")
assert resp.status_code == 200
assert resp.json()["access_group_id"] == access_group_id
def test_get_access_group_not_found(client_and_mocks):
"""Get access group returns 404 when not found."""
client, _, mock_table = client_and_mocks
mock_table.find_unique = AsyncMock(return_value=None)
resp = client.get("/v1/access_group/nonexistent-id")
assert resp.status_code == 404
assert "not found" in resp.json()["detail"]
@pytest.mark.parametrize("user_role", [LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY])
def test_get_access_group_forbidden_non_admin(client_and_mocks, user_role):
"""Non-admin users cannot get access group."""
client, _, _ = client_and_mocks
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_id="regular_user",
user_role=user_role,
)
resp = client.get("/v1/access_group/ag-123")
assert resp.status_code == 403
assert resp.json()["detail"]["error"] == CommonProxyErrors.not_allowed_access.value
# ---------------------------------------------------------------------------
# UPDATE
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("base_path", ACCESS_GROUP_PATHS)
@pytest.mark.parametrize(
"update_payload",
[
{"description": "Updated description"},
{"access_model_ids": ["model-1", "model-2"]},
{"assigned_team_ids": [], "assigned_key_ids": ["key-1"]},
],
)
def test_update_access_group_success(client_and_mocks, base_path, update_payload):
"""Update access group with various payloads returns 200."""
client, _, mock_table = client_and_mocks
existing = _make_access_group_record(access_group_id="ag-update")
mock_table.find_unique = AsyncMock(return_value=existing)
resp = client.put(f"{base_path}/ag-update", json=update_payload)
assert resp.status_code == 200
mock_table.update.assert_awaited_once()
def test_update_access_group_not_found(client_and_mocks):
"""Update access group returns 404 when not found."""
client, _, mock_table = client_and_mocks
mock_table.find_unique = AsyncMock(return_value=None)
resp = client.put(
"/v1/access_group/nonexistent-id",
json={"description": "Updated"},
)
assert resp.status_code == 404
assert "not found" in resp.json()["detail"]
mock_table.update.assert_not_awaited()
@pytest.mark.parametrize("user_role", [LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY])
def test_update_access_group_forbidden_non_admin(client_and_mocks, user_role):
"""Non-admin users cannot update access groups."""
client, _, _ = client_and_mocks
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_id="regular_user",
user_role=user_role,
)
resp = client.put("/v1/access_group/ag-123", json={"description": "Updated"})
assert resp.status_code == 403
assert resp.json()["detail"]["error"] == CommonProxyErrors.not_allowed_access.value
def test_update_access_group_empty_body(client_and_mocks):
"""Update with empty body succeeds; only updated_by is set."""
client, _, mock_table = client_and_mocks
existing = _make_access_group_record(access_group_id="ag-update", access_group_name="unchanged")
mock_table.find_unique = AsyncMock(return_value=existing)
resp = client.put("/v1/access_group/ag-update", json={})
assert resp.status_code == 200
mock_table.update.assert_awaited_once()
call_kwargs = mock_table.update.call_args.kwargs
assert call_kwargs["where"] == {"access_group_id": "ag-update"}
assert "updated_by" in call_kwargs["data"]
assert call_kwargs["data"]["updated_by"] == "admin_user"
# ---------------------------------------------------------------------------
# DELETE
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("base_path", ACCESS_GROUP_PATHS)
@pytest.mark.parametrize("access_group_id", ["ag-123", "ag-delete-me"])
def test_delete_access_group_success(client_and_mocks, base_path, access_group_id):
"""Delete access group returns 204 when found."""
client, _, mock_table = client_and_mocks
existing = _make_access_group_record(access_group_id=access_group_id)
mock_table.find_unique = AsyncMock(return_value=existing)
resp = client.delete(f"{base_path}/{access_group_id}")
assert resp.status_code == 204
mock_table.delete.assert_awaited_once()
def test_delete_access_group_not_found(client_and_mocks):
"""Delete access group returns 404 when not found."""
client, _, mock_table = client_and_mocks
mock_table.find_unique = AsyncMock(return_value=None)
resp = client.delete("/v1/access_group/nonexistent-id")
assert resp.status_code == 404
assert "not found" in resp.json()["detail"]
mock_table.delete.assert_not_awaited()
@pytest.mark.parametrize("user_role", [LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY])
def test_delete_access_group_forbidden_non_admin(client_and_mocks, user_role):
"""Non-admin users cannot delete access groups."""
client, _, _ = client_and_mocks
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
user_id="regular_user",
user_role=user_role,
)
resp = client.delete("/v1/access_group/ag-123")
assert resp.status_code == 403
assert resp.json()["detail"]["error"] == CommonProxyErrors.not_allowed_access.value
def test_delete_access_group_cleans_up_teams_and_keys(client_and_mocks):
"""Delete removes access_group_id from teams and keys before deleting the group."""
client, mock_prisma, mock_access_group_table = client_and_mocks
mock_team_table = mock_prisma.db.litellm_teamtable
mock_key_table = mock_prisma.db.litellm_verificationtoken
existing = _make_access_group_record(access_group_id="ag-to-delete")
mock_access_group_table.find_unique = AsyncMock(return_value=existing)
team_with_group = MagicMock()
team_with_group.team_id = "team-1"
team_with_group.access_group_ids = ["ag-to-delete", "ag-other"]
mock_team_table.find_many = AsyncMock(return_value=[team_with_group])
key_with_group = MagicMock()
key_with_group.token = "key-token-1"
key_with_group.access_group_ids = ["ag-to-delete"]
mock_key_table.find_many = AsyncMock(return_value=[key_with_group])
resp = client.delete("/v1/access_group/ag-to-delete")
assert resp.status_code == 204
mock_team_table.update.assert_awaited_once_with(
where={"team_id": "team-1"},
data={"access_group_ids": ["ag-other"]},
)
mock_key_table.update.assert_awaited_once_with(
where={"token": "key-token-1"},
data={"access_group_ids": []},
)
mock_access_group_table.delete.assert_awaited_once_with(
where={"access_group_id": "ag-to-delete"}
)
def test_delete_access_group_503_on_db_connection_error(client_and_mocks):
"""Delete returns 503 when DB connection error occurs during transaction."""
client, _, mock_table = client_and_mocks
existing = _make_access_group_record(access_group_id="ag-to-delete")
mock_table.find_unique = AsyncMock(return_value=existing)
mock_table.delete = AsyncMock(side_effect=PrismaError())
resp = client.delete("/v1/access_group/ag-to-delete")
assert resp.status_code == 503
assert resp.json()["detail"] == CommonProxyErrors.db_not_connected_error.value
def test_delete_access_group_404_on_p2025_or_record_not_found(client_and_mocks):
"""Delete returns 404 when Prisma raises P2025 or record-not-found error."""
client, _, mock_table = client_and_mocks
existing = _make_access_group_record(access_group_id="ag-to-delete")
mock_table.find_unique = AsyncMock(return_value=existing)
mock_table.delete = AsyncMock(side_effect=Exception("P2025: Record to delete does not exist"))
resp = client.delete("/v1/access_group/ag-to-delete")
assert resp.status_code == 404
assert "not found" in resp.json()["detail"]
def test_delete_access_group_500_on_generic_exception(client_and_mocks):
"""Delete returns 500 when generic exception occurs during transaction."""
client, _, mock_table = client_and_mocks
existing = _make_access_group_record(access_group_id="ag-to-delete")
mock_table.find_unique = AsyncMock(return_value=existing)
mock_table.delete = AsyncMock(side_effect=RuntimeError("Unexpected error"))
resp = client.delete("/v1/access_group/ag-to-delete")
assert resp.status_code == 500
assert "Failed to delete access group" in resp.json()["detail"]
# ---------------------------------------------------------------------------
# DB NOT CONNECTED
# ---------------------------------------------------------------------------
@pytest.mark.parametrize(
"method,url,factory",
[
("post", "/v1/access_group", lambda: {"json": {"access_group_name": "test"}}),
("get", "/v1/access_group", lambda: {}),
("get", "/v1/access_group/ag-123", lambda: {}),
("put", "/v1/access_group/ag-123", lambda: {"json": {"description": "x"}}),
("delete", "/v1/access_group/ag-123", lambda: {}),
# Alias: /v1/unified_access_group
("post", "/v1/unified_access_group", lambda: {"json": {"access_group_name": "test"}}),
("get", "/v1/unified_access_group", lambda: {}),
("get", "/v1/unified_access_group/ag-123", lambda: {}),
("put", "/v1/unified_access_group/ag-123", lambda: {"json": {"description": "x"}}),
("delete", "/v1/unified_access_group/ag-123", lambda: {}),
],
)
def test_access_group_endpoints_db_not_connected(client_and_mocks, monkeypatch, method, url, factory):
"""All endpoints return 500 when DB is not connected."""
client, _, _ = client_and_mocks
monkeypatch.setattr(ps, "prisma_client", None)
resp = getattr(client, method)(url, **factory())
assert resp.status_code == 500
assert resp.json()["detail"]["error"] == CommonProxyErrors.db_not_connected_error.value