tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-13 15:07:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

Revert "[Feature] Deleted Keys and Deleted Teams Table"

Browse Source
This commit is contained in:
YutaSaito
2026-01-17 06:46:41 +09:00
committed by GitHub
parent 3a6c6a1faf
commit 034e3a6d44
16 changed files with 89 additions and 1410 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm-proxy-extras/dist/litellm_proxy_extras-0.4.15-py3-none-any.whl
Vendored
BIN
View File
Binary file not shown.
litellm-proxy-extras/dist/litellm_proxy_extras-0.4.15.tar.gz
Vendored
BIN
View File
Binary file not shown.
litellm-proxy-extras/dist/litellm_proxy_extras-0.4.22-py3-none-any.whl
Vendored
BIN
View File
Binary file not shown.
litellm-proxy-extras/dist/litellm_proxy_extras-0.4.22.tar.gz
Vendored
BIN
View File
Binary file not shown.
litellm-proxy-extras/litellm_proxy_extras/migrations/20251219110931_add_deleted_keys_and_deleted_teams_tables/migration.sql
-117
View File
@@ -1,117 +0,0 @@
-- CreateTable
CREATE TABLE "LiteLLM_DeletedTeamTable" (
"id" TEXT NOT NULL,
"team_id" TEXT NOT NULL,
"team_alias" TEXT,
"organization_id" TEXT,
"object_permission_id" TEXT,
"admins" TEXT[],
"members" TEXT[],
"members_with_roles" JSONB NOT NULL DEFAULT '{}',
"metadata" JSONB NOT NULL DEFAULT '{}',
"max_budget" DOUBLE PRECISION,
"spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0,
"models" TEXT[],
"max_parallel_requests" INTEGER,
"tpm_limit" BIGINT,
"rpm_limit" BIGINT,
"budget_duration" TEXT,
"budget_reset_at" TIMESTAMP(3),
"blocked" BOOLEAN NOT NULL DEFAULT false,
"model_spend" JSONB NOT NULL DEFAULT '{}',
"model_max_budget" JSONB NOT NULL DEFAULT '{}',
"team_member_permissions" TEXT[] DEFAULT ARRAY[]::TEXT[],
"model_id" INTEGER,
"created_at" TIMESTAMP(3),
"updated_at" TIMESTAMP(3),
"deleted_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"deleted_by" TEXT,
"deleted_by_api_key" TEXT,
"litellm_changed_by" TEXT,
CONSTRAINT "LiteLLM_DeletedTeamTable_pkey" PRIMARY KEY ("id")
);
-- CreateTable
CREATE TABLE "LiteLLM_DeletedVerificationToken" (
"id" TEXT NOT NULL,
"token" TEXT NOT NULL,
"key_name" TEXT,
"key_alias" TEXT,
"soft_budget_cooldown" BOOLEAN NOT NULL DEFAULT false,
"spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0,
"expires" TIMESTAMP(3),
"models" TEXT[],
"aliases" JSONB NOT NULL DEFAULT '{}',
"config" JSONB NOT NULL DEFAULT '{}',
"user_id" TEXT,
"team_id" TEXT,
"permissions" JSONB NOT NULL DEFAULT '{}',
"max_parallel_requests" INTEGER,
"metadata" JSONB NOT NULL DEFAULT '{}',
"blocked" BOOLEAN,
"tpm_limit" BIGINT,
"rpm_limit" BIGINT,
"max_budget" DOUBLE PRECISION,
"budget_duration" TEXT,
"budget_reset_at" TIMESTAMP(3),
"allowed_cache_controls" TEXT[] DEFAULT ARRAY[]::TEXT[],
"allowed_routes" TEXT[] DEFAULT ARRAY[]::TEXT[],
"model_spend" JSONB NOT NULL DEFAULT '{}',
"model_max_budget" JSONB NOT NULL DEFAULT '{}',
"budget_id" TEXT,
"organization_id" TEXT,
"object_permission_id" TEXT,
"created_at" TIMESTAMP(3),
"created_by" TEXT,
"updated_at" TIMESTAMP(3),
"updated_by" TEXT,
"rotation_count" INTEGER DEFAULT 0,
"auto_rotate" BOOLEAN DEFAULT false,
"rotation_interval" TEXT,
"last_rotation_at" TIMESTAMP(3),
"key_rotation_at" TIMESTAMP(3),
"deleted_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"deleted_by" TEXT,
"deleted_by_api_key" TEXT,
"litellm_changed_by" TEXT,
CONSTRAINT "LiteLLM_DeletedVerificationToken_pkey" PRIMARY KEY ("id")
);
-- CreateIndex
CREATE INDEX "LiteLLM_DeletedTeamTable_team_id_idx" ON "LiteLLM_DeletedTeamTable"("team_id");
-- CreateIndex
CREATE INDEX "LiteLLM_DeletedTeamTable_deleted_at_idx" ON "LiteLLM_DeletedTeamTable"("deleted_at");
-- CreateIndex
CREATE INDEX "LiteLLM_DeletedTeamTable_organization_id_idx" ON "LiteLLM_DeletedTeamTable"("organization_id");
-- CreateIndex
CREATE INDEX "LiteLLM_DeletedTeamTable_team_alias_idx" ON "LiteLLM_DeletedTeamTable"("team_alias");
-- CreateIndex
CREATE INDEX "LiteLLM_DeletedTeamTable_created_at_idx" ON "LiteLLM_DeletedTeamTable"("created_at");
-- CreateIndex
CREATE INDEX "LiteLLM_DeletedVerificationToken_token_idx" ON "LiteLLM_DeletedVerificationToken"("token");
-- CreateIndex
CREATE INDEX "LiteLLM_DeletedVerificationToken_deleted_at_idx" ON "LiteLLM_DeletedVerificationToken"("deleted_at");
-- CreateIndex
CREATE INDEX "LiteLLM_DeletedVerificationToken_user_id_idx" ON "LiteLLM_DeletedVerificationToken"("user_id");
-- CreateIndex
CREATE INDEX "LiteLLM_DeletedVerificationToken_team_id_idx" ON "LiteLLM_DeletedVerificationToken"("team_id");
-- CreateIndex
CREATE INDEX "LiteLLM_DeletedVerificationToken_organization_id_idx" ON "LiteLLM_DeletedVerificationToken"("organization_id");
-- CreateIndex
CREATE INDEX "LiteLLM_DeletedVerificationToken_key_alias_idx" ON "LiteLLM_DeletedVerificationToken"("key_alias");
-- CreateIndex
CREATE INDEX "LiteLLM_DeletedVerificationToken_created_at_idx" ON "LiteLLM_DeletedVerificationToken"("created_at");
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
-99
View File
@@ -132,49 +132,6 @@ model LiteLLM_TeamTable {
object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
}
// Audit table for deleted teams - preserves spend and team information for historical tracking
model LiteLLM_DeletedTeamTable {
id String @id @default(uuid())
team_id String // Original team_id
team_alias String?
organization_id String?
object_permission_id String?
admins String[]
members String[]
members_with_roles Json @default("{}")
metadata Json @default("{}")
max_budget Float?
spend Float @default(0.0)
models String[]
max_parallel_requests Int?
tpm_limit BigInt?
rpm_limit BigInt?
budget_duration String?
budget_reset_at DateTime?
blocked Boolean @default(false)
model_spend Json @default("{}")
model_max_budget Json @default("{}")
router_settings Json? @default("{}")
team_member_permissions String[] @default([])
model_id Int? // id for LiteLLM_ModelTable -> stores team-level model aliases
// Original timestamps from team creation/updates
created_at DateTime? @map("created_at")
updated_at DateTime? @map("updated_at")
// Deletion metadata
deleted_at DateTime @default(now()) @map("deleted_at")
deleted_by String? @map("deleted_by") // User who deleted the team
deleted_by_api_key String? @map("deleted_by_api_key") // API key hash that performed the deletion
litellm_changed_by String? @map("litellm_changed_by") // From litellm-changed-by header if provided
@@index([team_id])
@@index([deleted_at])
@@index([organization_id])
@@index([team_alias])
@@index([created_at])
}
// Track spend, rate limit, budget Users
model LiteLLM_UserTable {
user_id String @id
@@ -302,62 +259,6 @@ model LiteLLM_VerificationToken {
object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
}
// Audit table for deleted keys - preserves spend and key information for historical tracking
model LiteLLM_DeletedVerificationToken {
id String @id @default(uuid())
token String // Original token (hashed)
key_name String?
key_alias String?
soft_budget_cooldown Boolean @default(false)
spend Float @default(0.0)
expires DateTime?
models String[]
aliases Json @default("{}")
config Json @default("{}")
user_id String?
team_id String?
permissions Json @default("{}")
max_parallel_requests Int?
metadata Json @default("{}")
blocked Boolean?
tpm_limit BigInt?
rpm_limit BigInt?
max_budget Float?
budget_duration String?
budget_reset_at DateTime?
allowed_cache_controls String[] @default([])
allowed_routes String[] @default([])
model_spend Json @default("{}")
model_max_budget Json @default("{}")
router_settings Json? @default("{}")
budget_id String?
organization_id String?
object_permission_id String?
created_at DateTime? // Original creation timestamp
created_by String? // Original creator
updated_at DateTime? // Last update timestamp before deletion
updated_by String? // Last user who updated before deletion
rotation_count Int? @default(0)
auto_rotate Boolean? @default(false)
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
deleted_by_api_key String? @map("deleted_by_api_key") // API key hash that performed the deletion
litellm_changed_by String? @map("litellm_changed_by") // From litellm-changed-by header if provided
@@index([token])
@@index([deleted_at])
@@index([user_id])
@@index([team_id])
@@index([organization_id])
@@index([key_alias])
@@index([created_at])
}
model LiteLLM_EndUserTable {
user_id String @id
alias String? // admin-facing alias
litellm/proxy/_types.py
-30
View File
@@ -1723,21 +1723,6 @@ class LiteLLM_TeamTableCachedObj(LiteLLM_TeamTable):
last_refreshed_at: Optional[float] = None
class LiteLLM_DeletedTeamTable(LiteLLM_TeamTable):
"""
Recording of deleted teams for audit purposes. Mirrors LiteLLM_TeamTable
plus metadata captured at deletion time.
"""
id: Optional[str] = None
deleted_at: Optional[datetime] = None
deleted_by: Optional[str] = None
deleted_by_api_key: Optional[str] = None
litellm_changed_by: Optional[str] = None
model_config = ConfigDict(protected_namespaces=())
class TeamRequest(LiteLLMPydanticObjectBase):
teams: List[str]
@@ -2132,21 +2117,6 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase):
model_config = ConfigDict(protected_namespaces=())
class LiteLLM_DeletedVerificationToken(LiteLLM_VerificationToken):
"""
Recording of deleted keys for audit purposes. Mirrors LiteLLM_VerificationToken
plus metadata captured at deletion time.
"""
id: Optional[str] = None
deleted_at: Optional[datetime] = None
deleted_by: Optional[str] = None
deleted_by_api_key: Optional[str] = None
litellm_changed_by: Optional[str] = None
model_config = ConfigDict(protected_namespaces=())
class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken):
"""
Combined view of litellm verification token + litellm team table (select values)
litellm/proxy/management_endpoints/internal_user_endpoints.py
-13
View File
@@ -412,19 +412,6 @@ async def new_user(
status_code=403,
detail="License is over limit. Please contact support@berri.ai to upgrade your license.",
)
# Only proxy admins can create administrative users
# Check if user_api_key_dict is actually a UserAPIKeyAuth instance (not a Depends object)
# This can happen when the function is called directly in tests
if (
data.user_role in [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY]
and isinstance(user_api_key_dict, UserAPIKeyAuth)
and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN
):
raise HTTPException(
status_code=403,
detail=f"Only proxy admins can create administrative users (proxy_admin, proxy_admin_viewer). Attempted to create user with role: {data.user_role}. Your role: {user_api_key_dict.user_role}"
)
data_json = data.json() # type: ignore
data_json = _update_internal_new_user_params(data_json, data)
litellm/proxy/management_endpoints/key_management_endpoints.py
+29 -112
View File
@@ -16,7 +16,7 @@ import secrets
import traceback
import yaml
from datetime import datetime, timedelta, timezone
from typing import Any, Dict, List, Literal, Optional, Tuple, cast
from typing import List, Literal, Optional, Tuple, cast
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
import fastapi
from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, status
@@ -1791,10 +1791,6 @@ async def delete_key_fn(
if prisma_client is None:
raise Exception("Not connected to DB!")
# Normalize litellm_changed_by: if it's a Header object or not a string, convert to None
if litellm_changed_by is not None and not isinstance(litellm_changed_by, str):
litellm_changed_by = None
## only allow user to delete keys they own
verbose_proxy_logger.debug(
f"user_api_key_dict.user_role: {user_api_key_dict.user_role}"
@@ -1807,7 +1803,6 @@ async def delete_key_fn(
tokens=data.keys,
user_api_key_cache=user_api_key_cache,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
num_keys_to_be_deleted = len(data.keys)
deleted_keys = data.keys
@@ -1817,7 +1812,6 @@ async def delete_key_fn(
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
num_keys_to_be_deleted = len(data.key_aliases)
deleted_keys = data.key_aliases
@@ -2439,7 +2433,6 @@ async def delete_verification_tokens(
tokens: List,
user_api_key_cache: DualCache,
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: Optional[str] = None,
) -> Tuple[Optional[Dict], List[LiteLLM_VerificationToken]]:
"""
Helper that deletes the list of tokens from the database
@@ -2476,43 +2469,38 @@ async def delete_verification_tokens(
detail={"error": "No keys found"},
)
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
authorized_keys = _keys_being_deleted
else:
authorized_keys = []
for key in _keys_being_deleted:
if await can_modify_verification_token(
key_info=key,
user_api_key_cache=user_api_key_cache,
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
):
authorized_keys.append(key)
else:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail={
"error": "You are not authorized to delete this key"
},
)
await _persist_deleted_verification_tokens(
keys=authorized_keys,
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
# Assuming 'db' is your Prisma Client instance
# check if admin making request - don't filter by user-id
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
deleted_tokens = await prisma_client.delete_data(tokens=tokens)
# else
else:
deletion_tasks = [
prisma_client.delete_data(tokens=[key.token])
for key in authorized_keys
]
await asyncio.gather(*deletion_tasks)
tasks = []
deleted_tokens = []
for key in _keys_being_deleted:
deleted_tokens = [key.token for key in authorized_keys]
if len(deleted_tokens) != len(tokens):
async def _delete_key(key: LiteLLM_VerificationToken):
if await can_modify_verification_token(
key_info=key,
user_api_key_cache=user_api_key_cache,
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
):
await prisma_client.delete_data(tokens=[key.token])
deleted_tokens.append(key.token)
else:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail={
"error": "You are not authorized to delete this key"
},
)
tasks.append(_delete_key(key))
await asyncio.gather(*tasks)
_num_deleted_tokens = len(deleted_tokens)
if _num_deleted_tokens != len(tokens):
failed_tokens = [
token for token in tokens if token not in deleted_tokens
]
@@ -2540,81 +2528,11 @@ async def delete_verification_tokens(
return {"deleted_keys": deleted_tokens}, _keys_being_deleted
def _transform_verification_tokens_to_deleted_records(
keys: List[LiteLLM_VerificationToken],
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: Optional[str] = None,
) -> List[Dict[str, Any]]:
"""Transform verification tokens into deleted token records ready for persistence."""
if not keys:
return []
deleted_at = datetime.now(timezone.utc)
records = []
for key in keys:
key_payload = key.model_dump()
deleted_record = LiteLLM_DeletedVerificationToken(
**key_payload,
deleted_at=deleted_at,
deleted_by=user_api_key_dict.user_id,
deleted_by_api_key=user_api_key_dict.api_key,
litellm_changed_by=litellm_changed_by,
)
record = deleted_record.model_dump()
# Map org_id to organization_id (model uses org_id, but schema expects organization_id)
org_id_value = record.pop("org_id", None)
if org_id_value is not None:
record["organization_id"] = org_id_value
for json_field in ["aliases", "config", "permissions", "metadata", "model_spend", "model_max_budget", "router_settings"]:
if json_field in record and record[json_field] is not None:
record[json_field] = json.dumps(record[json_field])
for rel_key in ("litellm_budget_table", "litellm_organization_table", "object_permission", "id"):
record.pop(rel_key, None)
records.append(record)
return records
async def _save_deleted_verification_token_records(
records: List[Dict[str, Any]],
prisma_client: PrismaClient,
) -> None:
"""Save deleted verification token records to the database."""
if not records:
return
await prisma_client.db.litellm_deletedverificationtoken.create_many(
data=records
)
async def _persist_deleted_verification_tokens(
keys: List[LiteLLM_VerificationToken],
prisma_client: PrismaClient,
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: Optional[str] = None,
) -> None:
"""Persist deleted verification token records by transforming and saving them."""
records = _transform_verification_tokens_to_deleted_records(
keys=keys,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
await _save_deleted_verification_token_records(
records=records,
prisma_client=prisma_client,
)
async def delete_key_aliases(
key_aliases: List[str],
user_api_key_cache: DualCache,
prisma_client: PrismaClient,
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: Optional[str] = None,
) -> Tuple[Optional[Dict], List[LiteLLM_VerificationToken]]:
_keys_being_deleted = await prisma_client.db.litellm_verificationtoken.find_many(
where={"key_alias": {"in": key_aliases}}
@@ -2625,7 +2543,6 @@ async def delete_key_aliases(
tokens=tokens,
user_api_key_cache=user_api_key_cache,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
litellm/proxy/management_endpoints/team_endpoints.py
-114
View File
@@ -34,10 +34,8 @@ from litellm.proxy._types import (
LiteLLM_OrganizationTableWithMembers,
LiteLLM_TeamMembership,
LiteLLM_TeamTable,
LiteLLM_DeletedTeamTable,
LiteLLM_TeamTableCachedObj,
LiteLLM_UserTable,
LiteLLM_VerificationToken,
LitellmTableNames,
LitellmUserRoles,
Member,
@@ -2020,28 +2018,6 @@ async def team_member_delete(
## DELETE KEYS CREATED BY USER FOR THIS TEAM
if user_ids_to_delete:
from litellm.proxy.management_endpoints.key_management_endpoints import (
_persist_deleted_verification_tokens,
)
# Fetch keys before deletion to persist them
keys_to_delete: List[LiteLLM_VerificationToken] = (
await prisma_client.db.litellm_verificationtoken.find_many(
where={
"user_id": {"in": list(user_ids_to_delete)},
"team_id": data.team_id,
}
)
)
if keys_to_delete:
await _persist_deleted_verification_tokens(
keys=keys_to_delete,
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=None,
)
await prisma_client.db.litellm_verificationtoken.delete_many(
where={
"user_id": {"in": list(user_ids_to_delete)},
@@ -2427,13 +2403,6 @@ async def delete_team(
team_row_pydantic = LiteLLM_TeamTable(**team_row_base.model_dump())
team_rows.append(team_row_pydantic)
await _persist_deleted_team_records(
teams=team_rows,
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
# Enterprise Feature - Audit Logging. Enable with litellm.store_audit_logs = True
# we do this after the first for loop, since first for loop is for validation. we only want this inserted after validation passes
if litellm.store_audit_logs is True:
@@ -2469,25 +2438,6 @@ async def delete_team(
# End of Audit logging
## DELETE ASSOCIATED KEYS
# Fetch keys before deletion to persist them
from litellm.proxy.management_endpoints.key_management_endpoints import (
_persist_deleted_verification_tokens,
)
keys_to_delete: List[LiteLLM_VerificationToken] = (
await prisma_client.db.litellm_verificationtoken.find_many(
where={"team_id": {"in": data.team_ids}}
)
)
if keys_to_delete:
await _persist_deleted_verification_tokens(
keys=keys_to_delete,
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
await prisma_client.delete_data(team_id_list=data.team_ids, table_name="key")
# ## DELETE TEAM MEMBERSHIPS
@@ -2516,70 +2466,6 @@ async def delete_team(
return deleted_teams
def _transform_teams_to_deleted_records(
teams: List[LiteLLM_TeamTable],
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: Optional[str] = None,
) -> List[Dict[str, Any]]:
"""Transform teams into deleted team records ready for persistence."""
if not teams:
return []
deleted_at = datetime.now(timezone.utc)
records = []
for team in teams:
team_payload = team.model_dump()
deleted_record = LiteLLM_DeletedTeamTable(
**team_payload,
deleted_at=deleted_at,
deleted_by=user_api_key_dict.user_id,
deleted_by_api_key=user_api_key_dict.api_key,
litellm_changed_by=litellm_changed_by,
)
record = deleted_record.model_dump()
for json_field in ["members_with_roles", "metadata", "model_spend", "model_max_budget", "router_settings"]:
if json_field in record and record[json_field] is not None:
record[json_field] = json.dumps(record[json_field])
for rel_key in ("litellm_model_table", "object_permission", "id"):
record.pop(rel_key, None)
records.append(record)
return records
async def _save_deleted_team_records(
records: List[Dict[str, Any]],
prisma_client: PrismaClient,
) -> None:
"""Save deleted team records to the database."""
if not records:
return
await prisma_client.db.litellm_deletedteamtable.create_many(
data=records
)
async def _persist_deleted_team_records(
teams: List[LiteLLM_TeamTable],
prisma_client: PrismaClient,
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: Optional[str] = None,
) -> None:
"""Persist deleted team records by transforming and saving them."""
records = _transform_teams_to_deleted_records(
teams=teams,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
await _save_deleted_team_records(
records=records,
prisma_client=prisma_client,
)
def validate_membership(
user_api_key_dict: UserAPIKeyAuth, team_table: LiteLLM_TeamTable
):
litellm/proxy/schema.prisma
-99
View File
@@ -132,49 +132,6 @@ model LiteLLM_TeamTable {
object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
}
// Audit table for deleted teams - preserves spend and team information for historical tracking
model LiteLLM_DeletedTeamTable {
id String @id @default(uuid())
team_id String // Original team_id
team_alias String?
organization_id String?
object_permission_id String?
admins String[]
members String[]
members_with_roles Json @default("{}")
metadata Json @default("{}")
max_budget Float?
spend Float @default(0.0)
models String[]
max_parallel_requests Int?
tpm_limit BigInt?
rpm_limit BigInt?
budget_duration String?
budget_reset_at DateTime?
blocked Boolean @default(false)
model_spend Json @default("{}")
model_max_budget Json @default("{}")
router_settings Json? @default("{}")
team_member_permissions String[] @default([])
model_id Int? // id for LiteLLM_ModelTable -> stores team-level model aliases
// Original timestamps from team creation/updates
created_at DateTime? @map("created_at")
updated_at DateTime? @map("updated_at")
// Deletion metadata
deleted_at DateTime @default(now()) @map("deleted_at")
deleted_by String? @map("deleted_by") // User who deleted the team
deleted_by_api_key String? @map("deleted_by_api_key") // API key hash that performed the deletion
litellm_changed_by String? @map("litellm_changed_by") // From litellm-changed-by header if provided
@@index([team_id])
@@index([deleted_at])
@@index([organization_id])
@@index([team_alias])
@@index([created_at])
}
// Track spend, rate limit, budget Users
model LiteLLM_UserTable {
user_id String @id
@@ -302,62 +259,6 @@ model LiteLLM_VerificationToken {
object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
}
// Audit table for deleted keys - preserves spend and key information for historical tracking
model LiteLLM_DeletedVerificationToken {
id String @id @default(uuid())
token String // Original token (hashed)
key_name String?
key_alias String?
soft_budget_cooldown Boolean @default(false)
spend Float @default(0.0)
expires DateTime?
models String[]
aliases Json @default("{}")
config Json @default("{}")
user_id String?
team_id String?
permissions Json @default("{}")
max_parallel_requests Int?
metadata Json @default("{}")
blocked Boolean?
tpm_limit BigInt?
rpm_limit BigInt?
max_budget Float?
budget_duration String?
budget_reset_at DateTime?
allowed_cache_controls String[] @default([])
allowed_routes String[] @default([])
model_spend Json @default("{}")
model_max_budget Json @default("{}")
router_settings Json? @default("{}")
budget_id String?
organization_id String?
object_permission_id String?
created_at DateTime? // Original creation timestamp
created_by String? // Original creator
updated_at DateTime? // Last update timestamp before deletion
updated_by String? // Last user who updated before deletion
rotation_count Int? @default(0)
auto_rotate Boolean? @default(false)
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
deleted_by_api_key String? @map("deleted_by_api_key") // API key hash that performed the deletion
litellm_changed_by String? @map("litellm_changed_by") // From litellm-changed-by header if provided
@@index([token])
@@index([deleted_at])
@@index([user_id])
@@index([team_id])
@@index([organization_id])
@@index([key_alias])
@@index([created_at])
}
model LiteLLM_EndUserTable {
user_id String @id
alias String? // admin-facing alias
schema.prisma
-99
View File
@@ -132,49 +132,6 @@ model LiteLLM_TeamTable {
object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
}
// Audit table for deleted teams - preserves spend and team information for historical tracking
model LiteLLM_DeletedTeamTable {
id String @id @default(uuid())
team_id String // Original team_id
team_alias String?
organization_id String?
object_permission_id String?
admins String[]
members String[]
members_with_roles Json @default("{}")
metadata Json @default("{}")
max_budget Float?
spend Float @default(0.0)
models String[]
max_parallel_requests Int?
tpm_limit BigInt?
rpm_limit BigInt?
budget_duration String?
budget_reset_at DateTime?
blocked Boolean @default(false)
model_spend Json @default("{}")
model_max_budget Json @default("{}")
router_settings Json? @default("{}")
team_member_permissions String[] @default([])
model_id Int? // id for LiteLLM_ModelTable -> stores team-level model aliases
// Original timestamps from team creation/updates
created_at DateTime? @map("created_at")
updated_at DateTime? @map("updated_at")
// Deletion metadata
deleted_at DateTime @default(now()) @map("deleted_at")
deleted_by String? @map("deleted_by") // User who deleted the team
deleted_by_api_key String? @map("deleted_by_api_key") // API key hash that performed the deletion
litellm_changed_by String? @map("litellm_changed_by") // From litellm-changed-by header if provided
@@index([team_id])
@@index([deleted_at])
@@index([organization_id])
@@index([team_alias])
@@index([created_at])
}
// Track spend, rate limit, budget Users
model LiteLLM_UserTable {
user_id String @id
@@ -302,62 +259,6 @@ model LiteLLM_VerificationToken {
object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
}
// Audit table for deleted keys - preserves spend and key information for historical tracking
model LiteLLM_DeletedVerificationToken {
id String @id @default(uuid())
token String // Original token (hashed)
key_name String?
key_alias String?
soft_budget_cooldown Boolean @default(false)
spend Float @default(0.0)
expires DateTime?
models String[]
aliases Json @default("{}")
config Json @default("{}")
user_id String?
team_id String?
permissions Json @default("{}")
max_parallel_requests Int?
metadata Json @default("{}")
blocked Boolean?
tpm_limit BigInt?
rpm_limit BigInt?
max_budget Float?
budget_duration String?
budget_reset_at DateTime?
allowed_cache_controls String[] @default([])
allowed_routes String[] @default([])
model_spend Json @default("{}")
model_max_budget Json @default("{}")
router_settings Json? @default("{}")
budget_id String?
organization_id String?
object_permission_id String?
created_at DateTime? // Original creation timestamp
created_by String? // Original creator
updated_at DateTime? // Last update timestamp before deletion
updated_by String? // Last user who updated before deletion
rotation_count Int? @default(0)
auto_rotate Boolean? @default(false)
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
deleted_by_api_key String? @map("deleted_by_api_key") // API key hash that performed the deletion
litellm_changed_by String? @map("litellm_changed_by") // From litellm-changed-by header if provided
@@index([token])
@@index([deleted_at])
@@index([user_id])
@@index([team_id])
@@index([organization_id])
@@index([key_alias])
@@index([created_at])
}
model LiteLLM_EndUserTable {
user_id String @id
alias String? // admin-facing alias
tests/proxy_admin_ui_tests/test_key_management.py
-2
View File
@@ -1061,7 +1061,6 @@ async def test_list_key_helper(prisma_client):
api_key="sk-1234",
user_id="admin",
),
litellm_changed_by=None,
)
@@ -1182,7 +1181,6 @@ async def test_list_key_helper_team_filtering(prisma_client):
api_key="sk-1234",
user_id="admin",
),
litellm_changed_by=None,
)
tests/proxy_unit_tests/test_key_generate_prisma.py
+4 -16
View File
@@ -1166,10 +1166,8 @@ def test_delete_key_auth(prisma_client):
asyncio.run(test())
except Exception as e:
print("Got Exception", e)
# Handle different exception types - ProxyException has .message, others might have .detail or str(e)
error_message = getattr(e, "message", None) or getattr(e, "detail", None) or str(e)
print(f"Error message: {error_message}")
assert "Authentication Error" in error_message or "Invalid proxy server token" in error_message or "not found in db" in error_message
print(e.message)
assert "Authentication Error" in e.message
pass
@@ -2710,12 +2708,7 @@ async def test_reset_spend_authentication(prisma_client):
_response = await new_user(
data=NewUserRequest(
tpm_limit=20,
),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key=master_key,
user_id="1234",
),
)
)
generate_key = "Bearer " + _response.key
@@ -2735,12 +2728,7 @@ async def test_reset_spend_authentication(prisma_client):
data=NewUserRequest(
user_role=LitellmUserRoles.PROXY_ADMIN,
tpm_limit=20,
),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key=master_key,
user_id="1234",
),
)
)
generate_key = "Bearer " + _response.key
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
+56 -360
View File
@@ -31,13 +31,9 @@ from litellm.proxy.management_endpoints.key_management_endpoints import (
_check_team_key_limits,
_common_key_generation_helper,
_list_key_helper,
_persist_deleted_verification_tokens,
_save_deleted_verification_token_records,
_transform_verification_tokens_to_deleted_records,
can_modify_verification_token,
check_org_key_model_specific_limits,
check_team_key_model_specific_limits,
delete_verification_tokens,
generate_key_helper_fn,
prepare_key_update_data,
validate_key_team_change,
@@ -2732,364 +2728,64 @@ def test_check_org_key_model_specific_limits_org_model_tpm_overallocation():
)
def test_transform_verification_tokens_to_deleted_records():
from datetime import datetime, timezone
user_api_key_dict = UserAPIKeyAuth(
user_id="user-123",
api_key="sk-test",
user_role=LitellmUserRoles.PROXY_ADMIN.value,
)
key1 = LiteLLM_VerificationToken(
token="hashed-token-1",
user_id="user-123",
team_id="team-456",
key_alias="test-key-1",
spend=100.0,
max_budget=1000.0,
models=["gpt-4"],
aliases={},
config={},
permissions={},
metadata={"test": "value"},
model_max_budget={},
model_spend={},
soft_budget_cooldown=False,
allowed_routes=[],
)
key2 = LiteLLM_VerificationToken(
token="hashed-token-2",
user_id="user-789",
team_id=None,
key_alias="test-key-2",
spend=50.0,
max_budget=500.0,
models=["gpt-3.5-turbo"],
aliases={"alias": "model"},
config={"config": "value"},
permissions={"permission": True},
metadata={},
model_max_budget={"gpt-4": {"budget_limit": 100.0}},
model_spend={},
soft_budget_cooldown=False,
allowed_routes=[],
)
records = _transform_verification_tokens_to_deleted_records(
keys=[key1, key2],
user_api_key_dict=user_api_key_dict,
litellm_changed_by="admin-user",
)
assert len(records) == 2
assert all("deleted_at" in record for record in records)
assert all("deleted_by" in record for record in records)
assert all("deleted_by_api_key" in record for record in records)
assert all("litellm_changed_by" in record for record in records)
assert all(record["deleted_by"] == "user-123" for record in records)
assert all(record["deleted_by_api_key"] == user_api_key_dict.api_key for record in records)
assert all(record["litellm_changed_by"] == "admin-user" for record in records)
record1 = records[0]
assert record1["token"] == "hashed-token-1"
assert record1["user_id"] == "user-123"
assert record1["team_id"] == "team-456"
assert isinstance(record1["aliases"], str)
assert isinstance(record1["config"], str)
assert isinstance(record1["permissions"], str)
assert isinstance(record1["metadata"], str)
assert "litellm_budget_table" not in record1
assert "litellm_organization_table" not in record1
assert "object_permission" not in record1
assert "id" not in record1
record2 = records[1]
assert record2["token"] == "hashed-token-2"
assert isinstance(record2["model_max_budget"], str)
def test_transform_verification_tokens_to_deleted_records_empty_list():
user_api_key_dict = UserAPIKeyAuth(
user_id="user-123",
api_key="sk-test",
user_role=LitellmUserRoles.PROXY_ADMIN.value,
)
records = _transform_verification_tokens_to_deleted_records(
keys=[],
user_api_key_dict=user_api_key_dict,
)
assert records == []
@pytest.mark.asyncio
async def test_save_deleted_verification_token_records():
mock_prisma_client = AsyncMock()
mock_create_many = AsyncMock()
mock_prisma_client.db.litellm_deletedverificationtoken.create_many = (
mock_create_many
)
records = [
{
"token": "hashed-token-1",
"user_id": "user-123",
"deleted_at": "2024-01-01T00:00:00Z",
"deleted_by": "admin",
},
{
"token": "hashed-token-2",
"user_id": "user-456",
"deleted_at": "2024-01-01T00:00:00Z",
"deleted_by": "admin",
},
]
await _save_deleted_verification_token_records(
records=records, prisma_client=mock_prisma_client
)
mock_create_many.assert_called_once_with(data=records)
@pytest.mark.asyncio
async def test_save_deleted_verification_token_records_empty_list():
mock_prisma_client = AsyncMock()
mock_create_many = AsyncMock()
mock_prisma_client.db.litellm_deletedverificationtoken.create_many = (
mock_create_many
)
await _save_deleted_verification_token_records(
records=[], prisma_client=mock_prisma_client
)
mock_create_many.assert_not_called()
@pytest.mark.asyncio
async def test_persist_deleted_verification_tokens():
mock_prisma_client = AsyncMock()
mock_create_many = AsyncMock()
mock_prisma_client.db.litellm_deletedverificationtoken.create_many = (
mock_create_many
)
user_api_key_dict = UserAPIKeyAuth(
user_id="user-123",
api_key="sk-test",
user_role=LitellmUserRoles.PROXY_ADMIN.value,
)
key = LiteLLM_VerificationToken(
token="hashed-token-1",
user_id="user-123",
team_id="team-456",
key_alias="test-key",
spend=100.0,
max_budget=1000.0,
models=["gpt-4"],
aliases={},
config={},
permissions={},
metadata={},
model_max_budget={},
model_spend={},
soft_budget_cooldown=False,
allowed_routes=[],
)
await _persist_deleted_verification_tokens(
keys=[key],
prisma_client=mock_prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_changed_by="admin-user",
)
mock_create_many.assert_called_once()
call_args = mock_create_many.call_args
assert "data" in call_args.kwargs
records = call_args.kwargs["data"]
assert len(records) == 1
assert records[0]["token"] == "hashed-token-1"
assert records[0]["deleted_by"] == "user-123"
assert records[0]["litellm_changed_by"] == "admin-user"
@pytest.mark.asyncio
async def test_delete_verification_tokens_persists_deleted_keys(monkeypatch):
mock_prisma_client = AsyncMock()
mock_user_api_key_cache = MagicMock()
user_api_key_dict = UserAPIKeyAuth(
user_id="admin-user",
api_key="sk-admin",
user_role=LitellmUserRoles.PROXY_ADMIN.value,
)
key1 = LiteLLM_VerificationToken(
token="hashed-token-1",
user_id="user-123",
team_id="team-456",
key_alias="test-key-1",
spend=100.0,
max_budget=1000.0,
models=["gpt-4"],
aliases={},
config={},
permissions={},
metadata={},
model_max_budget={},
model_spend={},
soft_budget_cooldown=False,
allowed_routes=[],
)
key2 = LiteLLM_VerificationToken(
token="hashed-token-2",
user_id="user-789",
team_id=None,
key_alias="test-key-2",
spend=50.0,
max_budget=500.0,
models=["gpt-3.5-turbo"],
aliases={},
config={},
permissions={},
metadata={},
model_max_budget={},
model_spend={},
soft_budget_cooldown=False,
allowed_routes=[],
)
mock_find_many = AsyncMock(return_value=[key1, key2])
mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many
# delete_data returns {"deleted_keys": ...} from utils.py line 3049
# The function at line 2410 assigns it to deleted_tokens
# Then at line 2444 returns {"deleted_keys": deleted_tokens}
# So if delete_data returns {"deleted_keys": list}, then result would be nested
# But looking at the error, it seems like delete_data might return just the list
# Or the code extracts it. Let's return the list directly since that's what the test expects
mock_delete_data = AsyncMock(return_value=["hashed-token-1", "hashed-token-2"])
mock_prisma_client.delete_data = mock_delete_data
# Mock cache delete_cache method
mock_user_api_key_cache.delete_cache = MagicMock()
mock_create_many = AsyncMock()
mock_prisma_client.db.litellm_deletedverificationtoken.create_many = (
mock_create_many
)
def mock_hash_token(token):
return token if not token.startswith("sk-") else f"hashed-{token}"
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints._hash_token_if_needed",
mock_hash_token,
)
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.hash_token",
mock_hash_token,
)
monkeypatch.setattr(
"litellm.proxy.proxy_server.prisma_client",
mock_prisma_client,
)
result, deleted_keys = await delete_verification_tokens(
tokens=["sk-token-1", "sk-token-2"],
user_api_key_cache=mock_user_api_key_cache,
user_api_key_dict=user_api_key_dict,
litellm_changed_by="admin-user",
)
mock_create_many.assert_called_once()
call_args = mock_create_many.call_args
assert "data" in call_args.kwargs
records = call_args.kwargs["data"]
assert len(records) == 2
assert all(record["deleted_by"] == "admin-user" for record in records)
assert all(record["litellm_changed_by"] == "admin-user" for record in records)
# delete_data returns the list directly, which gets wrapped in {"deleted_keys": ...}
assert isinstance(result["deleted_keys"], list)
assert set(result["deleted_keys"]) == {"hashed-token-1", "hashed-token-2"}
assert len(deleted_keys) == 2
@pytest.mark.asyncio
async def test_delete_key_fn_persists_deleted_keys(monkeypatch):
from litellm.proxy._types import KeyRequest
from litellm.proxy.management_endpoints.key_management_endpoints import (
delete_key_fn,
delete_verification_tokens,
)
mock_prisma_client = AsyncMock()
mock_user_api_key_cache = MagicMock()
user_api_key_dict = UserAPIKeyAuth(
user_id="admin-user",
api_key="sk-admin",
user_role=LitellmUserRoles.PROXY_ADMIN.value,
)
key1 = LiteLLM_VerificationToken(
token="hashed-token-1",
user_id="user-123",
team_id="team-456",
key_alias="test-key-1",
spend=100.0,
max_budget=1000.0,
models=["gpt-4"],
aliases={},
config={},
permissions={},
metadata={},
model_max_budget={},
model_spend={},
soft_budget_cooldown=False,
allowed_routes=[],
)
async def mock_delete_verification_tokens(*args, **kwargs):
return ({"deleted_keys": ["sk-token-1"]}, [key1])
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.delete_verification_tokens",
mock_delete_verification_tokens,
)
monkeypatch.setattr(
"litellm.proxy.proxy_server.prisma_client",
mock_prisma_client,
)
monkeypatch.setattr(
"litellm.proxy.proxy_server.user_api_key_cache",
mock_user_api_key_cache,
)
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_deleted_hook",
AsyncMock(),
)
data = KeyRequest(keys=["sk-token-1"])
result = await delete_key_fn(
data=data,
user_api_key_dict=user_api_key_dict,
litellm_changed_by="admin-user",
)
assert result["deleted_keys"] == ["sk-token-1"]
@pytest.mark.asyncio
async def test_can_delete_verification_token_proxy_admin_team_key(monkeypatch):
"""Test that proxy admin can delete any team key."""
key_info = LiteLLM_VerificationToken(
token="test-token",
user_id="other-user",
team_id="test-team-123",
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
user_id="admin-user",
api_key="sk-admin",
)
mock_prisma_client = AsyncMock()
mock_user_api_key_cache = MagicMock()
result = await can_modify_verification_token(
key_info=key_info,
user_api_key_cache=mock_user_api_key_cache,
user_api_key_dict=user_api_key_dict,
prisma_client=mock_prisma_client,
)
assert result is True
@pytest.mark.asyncio
async def test_can_delete_verification_token_proxy_admin_personal_key(monkeypatch):
"""Test that proxy admin can delete any personal key."""
key_info = LiteLLM_VerificationToken(
token="test-token",
user_id="other-user",
team_id=None,
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
user_id="admin-user",
api_key="sk-admin",
)
mock_prisma_client = AsyncMock()
mock_user_api_key_cache = MagicMock()
result = await can_modify_verification_token(
key_info=key_info,
user_api_key_cache=mock_user_api_key_cache,
user_api_key_dict=user_api_key_dict,
prisma_client=mock_prisma_client,
)
assert result is True
@pytest.mark.asyncio
async def test_can_delete_verification_token_team_admin_own_team(monkeypatch):
"""Test that team admin can delete team keys from their own team."""
key_info = LiteLLM_VerificationToken(
token="test-token",
tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
-349
View File
@@ -33,13 +33,8 @@ from litellm.proxy.management_endpoints.team_endpoints import (
from litellm.proxy.management_endpoints.team_endpoints import (
GetTeamMemberPermissionsResponse,
UpdateTeamMemberPermissionsRequest,
_persist_deleted_team_records,
_save_deleted_team_records,
_transform_teams_to_deleted_records,
delete_team,
router,
team_member_add_duplication_check,
team_member_delete,
validate_team_org_change,
)
from litellm.proxy.management_helpers.team_member_permission_checks import (
@@ -2265,7 +2260,6 @@ async def test_team_member_delete_cleans_membership(mock_db_client, mock_admin_a
# Verification token deletion should be called
mock_db_client.db.litellm_verificationtoken = MagicMock()
mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
mock_db_client.db.litellm_verificationtoken.delete_many = AsyncMock(return_value=MagicMock())
# Execute
@@ -2313,7 +2307,6 @@ async def test_team_member_delete_cleans_verification_tokens(mock_db_client, moc
mock_db_client.db.litellm_teammembership.delete_many = AsyncMock(return_value=MagicMock())
mock_db_client.db.litellm_verificationtoken = MagicMock()
mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
mock_db_client.db.litellm_verificationtoken.delete_many = AsyncMock(return_value=MagicMock())
await team_member_delete(
@@ -4332,348 +4325,6 @@ async def test_update_team_guardrails_with_org_id():
assert first_call_kwargs["include"]["teams"] is True
def test_transform_teams_to_deleted_records():
from datetime import datetime, timezone
user_api_key_dict = UserAPIKeyAuth(
user_id="user-123",
api_key="sk-test",
user_role=LitellmUserRoles.PROXY_ADMIN.value,
)
team1 = LiteLLM_TeamTable(
team_id="team-1",
team_alias="test-team-1",
members_with_roles=[
Member(user_id="user-1", role="admin"),
Member(user_id="user-2", role="user"),
],
metadata={"test": "value"},
model_max_budget={},
model_spend={},
)
team2 = LiteLLM_TeamTable(
team_id="team-2",
team_alias="test-team-2",
members_with_roles=[],
metadata=None,
model_max_budget={"gpt-4": {"budget_limit": 100.0}},
model_spend={},
)
records = _transform_teams_to_deleted_records(
teams=[team1, team2],
user_api_key_dict=user_api_key_dict,
litellm_changed_by="admin-user",
)
assert len(records) == 2
assert all("deleted_at" in record for record in records)
assert all("deleted_by" in record for record in records)
assert all("deleted_by_api_key" in record for record in records)
assert all("litellm_changed_by" in record for record in records)
assert all(record["deleted_by"] == "user-123" for record in records)
# UserAPIKeyAuth hashes the api_key, so we check against the hashed value
assert all(record["deleted_by_api_key"] == user_api_key_dict.api_key for record in records)
assert all(record["litellm_changed_by"] == "admin-user" for record in records)
record1 = records[0]
assert record1["team_id"] == "team-1"
assert isinstance(record1["members_with_roles"], str)
assert isinstance(record1["metadata"], str)
assert "litellm_model_table" not in record1
assert "object_permission" not in record1
assert "id" not in record1
record2 = records[1]
assert record2["team_id"] == "team-2"
# model_max_budget should be converted to JSON string if it exists
if "model_max_budget" in record2:
assert isinstance(record2["model_max_budget"], str)
def test_transform_teams_to_deleted_records_empty_list():
user_api_key_dict = UserAPIKeyAuth(
user_id="user-123",
api_key="sk-test",
user_role=LitellmUserRoles.PROXY_ADMIN.value,
)
records = _transform_teams_to_deleted_records(
teams=[],
user_api_key_dict=user_api_key_dict,
)
assert records == []
@pytest.mark.asyncio
async def test_save_deleted_team_records():
mock_prisma_client = AsyncMock()
mock_create_many = AsyncMock()
mock_prisma_client.db.litellm_deletedteamtable.create_many = mock_create_many
records = [
{
"team_id": "team-1",
"team_alias": "test-team-1",
"deleted_at": "2024-01-01T00:00:00Z",
"deleted_by": "admin",
},
{
"team_id": "team-2",
"team_alias": "test-team-2",
"deleted_at": "2024-01-01T00:00:00Z",
"deleted_by": "admin",
},
]
await _save_deleted_team_records(records=records, prisma_client=mock_prisma_client)
mock_create_many.assert_called_once_with(data=records)
@pytest.mark.asyncio
async def test_save_deleted_team_records_empty_list():
mock_prisma_client = AsyncMock()
mock_create_many = AsyncMock()
mock_prisma_client.db.litellm_deletedteamtable.create_many = mock_create_many
await _save_deleted_team_records(records=[], prisma_client=mock_prisma_client)
mock_create_many.assert_not_called()
@pytest.mark.asyncio
async def test_persist_deleted_team_records():
mock_prisma_client = AsyncMock()
mock_create_many = AsyncMock()
mock_prisma_client.db.litellm_deletedteamtable.create_many = mock_create_many
user_api_key_dict = UserAPIKeyAuth(
user_id="user-123",
api_key="sk-test",
user_role=LitellmUserRoles.PROXY_ADMIN.value,
)
team = LiteLLM_TeamTable(
team_id="team-1",
team_alias="test-team",
members_with_roles=[
Member(user_id="user-1", role="admin"),
],
metadata={},
model_max_budget={},
model_spend={},
)
await _persist_deleted_team_records(
teams=[team],
prisma_client=mock_prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_changed_by="admin-user",
)
mock_create_many.assert_called_once()
call_args = mock_create_many.call_args
assert "data" in call_args.kwargs
records = call_args.kwargs["data"]
assert len(records) == 1
assert records[0]["team_id"] == "team-1"
assert records[0]["deleted_by"] == "user-123"
assert records[0]["litellm_changed_by"] == "admin-user"
@pytest.mark.asyncio
async def test_delete_team_persists_deleted_teams(monkeypatch):
from litellm.proxy._types import DeleteTeamRequest
mock_prisma_client = AsyncMock()
mock_user_api_key_dict = UserAPIKeyAuth(
user_id="admin-user",
api_key="sk-admin",
user_role=LitellmUserRoles.PROXY_ADMIN.value,
)
team1 = LiteLLM_TeamTable(
team_id="team-1",
team_alias="test-team-1",
members_with_roles=[
Member(user_id="user-1", role="admin"),
],
metadata={},
model_max_budget={},
model_spend={},
)
mock_find_unique = AsyncMock(return_value=team1)
mock_prisma_client.db.litellm_teamtable.find_unique = mock_find_unique
mock_delete_data = AsyncMock(return_value={"deleted_teams": ["team-1"]})
mock_prisma_client.delete_data = mock_delete_data
mock_create_many_teams = AsyncMock()
mock_prisma_client.db.litellm_deletedteamtable.create_many = mock_create_many_teams
mock_create_many_keys = AsyncMock()
mock_prisma_client.db.litellm_deletedverificationtoken.create_many = (
mock_create_many_keys
)
mock_find_many_keys = AsyncMock(return_value=[])
mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many_keys
monkeypatch.setattr(
"litellm.proxy.proxy_server.prisma_client",
mock_prisma_client,
)
monkeypatch.setattr(
"litellm.proxy.proxy_server.create_audit_log_for_update",
AsyncMock(),
)
monkeypatch.setattr(
"litellm.proxy.proxy_server.litellm_proxy_admin_name",
"admin",
)
monkeypatch.setattr(
"litellm.proxy.management_endpoints.team_endpoints.team_member_delete",
AsyncMock(return_value=team1),
)
data = DeleteTeamRequest(team_ids=["team-1"])
result = await delete_team(
data=data,
http_request=MagicMock(),
user_api_key_dict=mock_user_api_key_dict,
litellm_changed_by="admin-user",
)
mock_create_many_teams.assert_called_once()
call_args = mock_create_many_teams.call_args
assert "data" in call_args.kwargs
records = call_args.kwargs["data"]
assert len(records) == 1
assert records[0]["team_id"] == "team-1"
assert records[0]["deleted_by"] == "admin-user"
assert records[0]["litellm_changed_by"] == "admin-user"
@pytest.mark.asyncio
async def test_team_member_delete_persists_deleted_keys(monkeypatch):
from litellm.proxy._types import TeamMemberDeleteRequest
from litellm.proxy.management_endpoints.key_management_endpoints import (
LiteLLM_VerificationToken,
)
mock_prisma_client = AsyncMock()
mock_user_api_key_dict = UserAPIKeyAuth(
user_id="admin-user",
api_key="sk-admin",
user_role=LitellmUserRoles.PROXY_ADMIN.value,
)
team = LiteLLM_TeamTable(
team_id="team-1",
team_alias="test-team",
members_with_roles=[
Member(user_id="user-123", role="admin"),
],
metadata={},
model_max_budget={},
model_spend={},
)
key1 = LiteLLM_VerificationToken(
token="hashed-token-1",
user_id="user-123",
team_id="team-1",
key_alias="test-key-1",
spend=100.0,
max_budget=1000.0,
models=["gpt-4"],
aliases={},
config={},
permissions={},
metadata={},
model_max_budget={},
)
key2 = LiteLLM_VerificationToken(
token="hashed-token-2",
user_id="user-123",
team_id="team-1",
key_alias="test-key-2",
spend=50.0,
max_budget=500.0,
models=["gpt-3.5-turbo"],
aliases={},
config={},
permissions={},
metadata={},
model_max_budget={},
)
mock_find_unique_team = AsyncMock(return_value=team)
mock_prisma_client.db.litellm_teamtable.find_unique = mock_find_unique_team
mock_find_many_user = AsyncMock(
return_value=[
MagicMock(
user_id="user-123",
teams=["team-1"],
model_dump=lambda: {"user_id": "user-123", "teams": ["team-1"]},
)
]
)
mock_prisma_client.db.litellm_usertable.find_many = mock_find_many_user
mock_update_team = AsyncMock()
mock_prisma_client.db.litellm_teamtable.update = mock_update_team
mock_update_user = AsyncMock()
mock_prisma_client.db.litellm_usertable.update = mock_update_user
mock_delete_membership = AsyncMock()
mock_prisma_client.db.litellm_teammembership.delete_many = mock_delete_membership
mock_find_many_keys = AsyncMock(return_value=[key1, key2])
mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many_keys
mock_delete_keys = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.delete_many = mock_delete_keys
mock_create_many_keys = AsyncMock()
mock_prisma_client.db.litellm_deletedverificationtoken.create_many = (
mock_create_many_keys
)
monkeypatch.setattr(
"litellm.proxy.proxy_server.prisma_client",
mock_prisma_client,
)
monkeypatch.setattr(
"litellm.proxy.management_endpoints.team_endpoints._is_user_team_admin",
lambda **kwargs: True,
)
data = TeamMemberDeleteRequest(team_id="team-1", user_id="user-123")
result = await team_member_delete(
data=data,
user_api_key_dict=mock_user_api_key_dict,
)
mock_create_many_keys.assert_called_once()
call_args = mock_create_many_keys.call_args
assert "data" in call_args.kwargs
records = call_args.kwargs["data"]
assert len(records) == 2
assert all(record["deleted_by"] == "admin-user" for record in records)
assert all(record["team_id"] == "team-1" for record in records)
assert all(record["user_id"] == "user-123" for record in records)
mock_delete_keys.assert_called_once()
@pytest.mark.asyncio
async def test_new_team_negative_max_budget():
"""