mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
feat(proxy): customer org, team and member hierarchy with budget templates (#51459)
* feat: add customer hierarchy budget management Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * refactor: clarify customer hierarchy management and enforce ownership Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * fix: return 201 for customer hierarchy creates Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * Refine customer budget templates and expiring member overrides Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * fix: guard customer hierarchy migration and register tables for span redaction Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * fix: chunk template budget lookup and lock the organization row while rewriting templates Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * fix: allow management_routes keys on customer routes and block /budget/delete on org default budgets Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * fix: clearing a member override drops its expiry and a null budget clears org and team budgets Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>
This commit is contained in:
ishaan-j-berri
GitHub
parent
0fe6927426
commit
b251f125a7
20 changed files with 3427 additions and 5 deletions
|
|
@ -263,6 +263,9 @@ _PRISMA_MODELS: Final[frozenset[str]] = frozenset(
|
|||
"LiteLLM_Alert",
|
||||
"LiteLLM_PerformanceAlertState",
|
||||
"LiteLLM_BudgetTable",
|
||||
"LiteLLM_CustomerOrganizationTable",
|
||||
"LiteLLM_CustomerTeamTable",
|
||||
"LiteLLM_CustomerMemberTable",
|
||||
"LiteLLM_CredentialsTable",
|
||||
"LiteLLM_ProxyModelTable",
|
||||
"LiteLLM_AgentsTable",
|
||||
|
|
|
|||
|
|
@ -0,0 +1,133 @@
|
|||
SET LOCAL lock_timeout = '10s';
|
||||
|
||||
CREATE TABLE IF NOT EXISTS "LiteLLM_CustomerOrganizationTable" (
|
||||
"spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0,
|
||||
"budget_reset_at" TIMESTAMP(3),
|
||||
"customer_org_id" TEXT NOT NULL,
|
||||
"alias" TEXT,
|
||||
"default_budget_id" TEXT,
|
||||
"budget_ids" TEXT[] NOT NULL DEFAULT ARRAY[]::TEXT[],
|
||||
"budget_names" JSONB NOT NULL DEFAULT '{}',
|
||||
"budget_id" TEXT,
|
||||
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"created_by" TEXT NOT NULL,
|
||||
"updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"updated_by" TEXT NOT NULL,
|
||||
|
||||
CONSTRAINT "LiteLLM_CustomerOrganizationTable_pkey" PRIMARY KEY ("customer_org_id")
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS "LiteLLM_CustomerTeamTable" (
|
||||
"spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0,
|
||||
"budget_reset_at" TIMESTAMP(3),
|
||||
"customer_team_id" TEXT NOT NULL,
|
||||
"alias" TEXT,
|
||||
"parent_org_id" TEXT NOT NULL,
|
||||
"budget_id" TEXT,
|
||||
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"created_by" TEXT NOT NULL,
|
||||
"updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"updated_by" TEXT NOT NULL,
|
||||
|
||||
CONSTRAINT "LiteLLM_CustomerTeamTable_pkey" PRIMARY KEY ("customer_team_id")
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS "LiteLLM_CustomerMemberTable" (
|
||||
"spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0,
|
||||
"budget_reset_at" TIMESTAMP(3),
|
||||
"customer_member_id" TEXT NOT NULL,
|
||||
"alias" TEXT,
|
||||
"parent_team_id" TEXT NOT NULL,
|
||||
"budget_id" TEXT,
|
||||
"budget_template_id" TEXT,
|
||||
"budget_override_expires_at" TIMESTAMP(3),
|
||||
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"created_by" TEXT NOT NULL,
|
||||
"updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"updated_by" TEXT NOT NULL,
|
||||
|
||||
CONSTRAINT "LiteLLM_CustomerMemberTable_pkey" PRIMARY KEY ("customer_member_id")
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_CustomerTeamTable_parent_org_id_idx" ON "LiteLLM_CustomerTeamTable"("parent_org_id");
|
||||
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_CustomerMemberTable_parent_team_id_idx" ON "LiteLLM_CustomerMemberTable"("parent_team_id");
|
||||
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_CustomerMemberTable_budget_template_id_idx" ON "LiteLLM_CustomerMemberTable"("budget_template_id");
|
||||
|
||||
DO $$
|
||||
BEGIN
|
||||
IF NOT EXISTS (
|
||||
SELECT 1 FROM pg_constraint
|
||||
WHERE conname = 'LiteLLM_CustomerOrganizationTable_default_budget_id_fkey'
|
||||
AND conrelid = '"LiteLLM_CustomerOrganizationTable"'::regclass
|
||||
) THEN
|
||||
ALTER TABLE "LiteLLM_CustomerOrganizationTable" ADD CONSTRAINT "LiteLLM_CustomerOrganizationTable_default_budget_id_fkey" FOREIGN KEY ("default_budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE;
|
||||
END IF;
|
||||
END $$;
|
||||
|
||||
DO $$
|
||||
BEGIN
|
||||
IF NOT EXISTS (
|
||||
SELECT 1 FROM pg_constraint
|
||||
WHERE conname = 'LiteLLM_CustomerTeamTable_parent_org_id_fkey'
|
||||
AND conrelid = '"LiteLLM_CustomerTeamTable"'::regclass
|
||||
) THEN
|
||||
ALTER TABLE "LiteLLM_CustomerTeamTable" ADD CONSTRAINT "LiteLLM_CustomerTeamTable_parent_org_id_fkey" FOREIGN KEY ("parent_org_id") REFERENCES "LiteLLM_CustomerOrganizationTable"("customer_org_id") ON DELETE RESTRICT ON UPDATE CASCADE;
|
||||
END IF;
|
||||
END $$;
|
||||
|
||||
DO $$
|
||||
BEGIN
|
||||
IF NOT EXISTS (
|
||||
SELECT 1 FROM pg_constraint
|
||||
WHERE conname = 'LiteLLM_CustomerTeamTable_budget_id_fkey'
|
||||
AND conrelid = '"LiteLLM_CustomerTeamTable"'::regclass
|
||||
) THEN
|
||||
ALTER TABLE "LiteLLM_CustomerTeamTable" ADD CONSTRAINT "LiteLLM_CustomerTeamTable_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE;
|
||||
END IF;
|
||||
END $$;
|
||||
|
||||
DO $$
|
||||
BEGIN
|
||||
IF NOT EXISTS (
|
||||
SELECT 1 FROM pg_constraint
|
||||
WHERE conname = 'LiteLLM_CustomerMemberTable_parent_team_id_fkey'
|
||||
AND conrelid = '"LiteLLM_CustomerMemberTable"'::regclass
|
||||
) THEN
|
||||
ALTER TABLE "LiteLLM_CustomerMemberTable" ADD CONSTRAINT "LiteLLM_CustomerMemberTable_parent_team_id_fkey" FOREIGN KEY ("parent_team_id") REFERENCES "LiteLLM_CustomerTeamTable"("customer_team_id") ON DELETE RESTRICT ON UPDATE CASCADE;
|
||||
END IF;
|
||||
END $$;
|
||||
|
||||
DO $$
|
||||
BEGIN
|
||||
IF NOT EXISTS (
|
||||
SELECT 1 FROM pg_constraint
|
||||
WHERE conname = 'LiteLLM_CustomerMemberTable_budget_id_fkey'
|
||||
AND conrelid = '"LiteLLM_CustomerMemberTable"'::regclass
|
||||
) THEN
|
||||
ALTER TABLE "LiteLLM_CustomerMemberTable" ADD CONSTRAINT "LiteLLM_CustomerMemberTable_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE RESTRICT ON UPDATE CASCADE;
|
||||
END IF;
|
||||
END $$;
|
||||
|
||||
DO $$
|
||||
BEGIN
|
||||
IF NOT EXISTS (
|
||||
SELECT 1 FROM pg_constraint
|
||||
WHERE conname = 'LiteLLM_CustomerOrganizationTable_budget_id_fkey'
|
||||
AND conrelid = '"LiteLLM_CustomerOrganizationTable"'::regclass
|
||||
) THEN
|
||||
ALTER TABLE "LiteLLM_CustomerOrganizationTable" ADD CONSTRAINT "LiteLLM_CustomerOrganizationTable_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE;
|
||||
END IF;
|
||||
END $$;
|
||||
|
||||
DO $$
|
||||
BEGIN
|
||||
IF NOT EXISTS (
|
||||
SELECT 1 FROM pg_constraint
|
||||
WHERE conname = 'LiteLLM_CustomerMemberTable_budget_template_id_fkey'
|
||||
AND conrelid = '"LiteLLM_CustomerMemberTable"'::regclass
|
||||
) THEN
|
||||
ALTER TABLE "LiteLLM_CustomerMemberTable" ADD CONSTRAINT "LiteLLM_CustomerMemberTable_budget_template_id_fkey" FOREIGN KEY ("budget_template_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE RESTRICT ON UPDATE CASCADE;
|
||||
END IF;
|
||||
END $$;
|
||||
|
|
@ -11,6 +11,11 @@ generator client {
|
|||
|
||||
// Budget / Rate Limits for an org
|
||||
model LiteLLM_BudgetTable {
|
||||
customer_organization_defaults LiteLLM_CustomerOrganizationTable[] @relation("CustomerOrganizationDefaultBudget")
|
||||
customer_organizations LiteLLM_CustomerOrganizationTable[] @relation("CustomerOrganizationBudget")
|
||||
customer_member_templates LiteLLM_CustomerMemberTable[] @relation("CustomerMemberTemplate")
|
||||
customer_teams LiteLLM_CustomerTeamTable[] @relation("CustomerTeamBudget")
|
||||
customer_members LiteLLM_CustomerMemberTable[] @relation("CustomerMemberBudget")
|
||||
budget_id String @id @default(uuid())
|
||||
max_budget Float?
|
||||
soft_budget Float?
|
||||
|
|
@ -40,6 +45,61 @@ model LiteLLM_BudgetTable {
|
|||
organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization
|
||||
}
|
||||
|
||||
model LiteLLM_CustomerOrganizationTable {
|
||||
spend Float @default(0.0)
|
||||
budget_reset_at DateTime?
|
||||
budget_ids String[] @default([])
|
||||
budget_names Json @default("{}")
|
||||
budget_id String?
|
||||
budget LiteLLM_BudgetTable? @relation("CustomerOrganizationBudget", fields: [budget_id], references: [budget_id])
|
||||
customer_org_id String @id @default(uuid())
|
||||
alias String?
|
||||
default_budget_id String?
|
||||
default_budget LiteLLM_BudgetTable? @relation("CustomerOrganizationDefaultBudget", fields: [default_budget_id], references: [budget_id])
|
||||
teams LiteLLM_CustomerTeamTable[]
|
||||
created_at DateTime @default(now())
|
||||
created_by String
|
||||
updated_at DateTime @default(now()) @updatedAt
|
||||
updated_by String
|
||||
}
|
||||
|
||||
model LiteLLM_CustomerTeamTable {
|
||||
spend Float @default(0.0)
|
||||
budget_reset_at DateTime?
|
||||
customer_team_id String @id @default(uuid())
|
||||
alias String?
|
||||
parent_org_id String
|
||||
parent_org LiteLLM_CustomerOrganizationTable @relation(fields: [parent_org_id], references: [customer_org_id])
|
||||
budget_id String?
|
||||
budget LiteLLM_BudgetTable? @relation("CustomerTeamBudget", fields: [budget_id], references: [budget_id])
|
||||
members LiteLLM_CustomerMemberTable[]
|
||||
created_at DateTime @default(now())
|
||||
created_by String
|
||||
updated_at DateTime @default(now()) @updatedAt
|
||||
updated_by String
|
||||
@@index([parent_org_id])
|
||||
}
|
||||
|
||||
model LiteLLM_CustomerMemberTable {
|
||||
spend Float @default(0.0)
|
||||
budget_reset_at DateTime?
|
||||
customer_member_id String @id @default(uuid())
|
||||
alias String?
|
||||
parent_team_id String
|
||||
parent_team LiteLLM_CustomerTeamTable @relation(fields: [parent_team_id], references: [customer_team_id])
|
||||
budget_id String?
|
||||
budget LiteLLM_BudgetTable? @relation("CustomerMemberBudget", fields: [budget_id], references: [budget_id], onDelete: Restrict)
|
||||
budget_template_id String?
|
||||
budget_template LiteLLM_BudgetTable? @relation("CustomerMemberTemplate", fields: [budget_template_id], references: [budget_id], onDelete: Restrict)
|
||||
budget_override_expires_at DateTime?
|
||||
created_at DateTime @default(now())
|
||||
created_by String
|
||||
updated_at DateTime @default(now()) @updatedAt
|
||||
updated_by String
|
||||
@@index([parent_team_id])
|
||||
@@index([budget_template_id])
|
||||
}
|
||||
|
||||
// Models on proxy
|
||||
model LiteLLM_CredentialsTable {
|
||||
credential_id String @id @default(uuid())
|
||||
|
|
|
|||
|
|
@ -724,6 +724,14 @@ class LiteLLMRoutes(enum.Enum):
|
|||
"/jwt/key/mapping/delete",
|
||||
"/jwt/key/mapping/list",
|
||||
"/jwt/key/mapping/info",
|
||||
"/management/v1/customer_organizations",
|
||||
"/management/v1/customer_organizations/{customer_org_id}",
|
||||
"/management/v1/customer_organizations/{customer_org_id}/default_budgets",
|
||||
"/management/v1/customer_organizations/{customer_org_id}/default_budgets/{budget_id}",
|
||||
"/management/v1/customer_teams",
|
||||
"/management/v1/customer_teams/{customer_team_id}",
|
||||
"/management/v1/customer_members",
|
||||
"/management/v1/customer_members/{customer_member_id}",
|
||||
]
|
||||
+ key_management_routes
|
||||
+ mcp_management_routes
|
||||
|
|
|
|||
|
|
@ -31,6 +31,7 @@ from litellm.proxy.management_endpoints.common_utils import ( # noqa: F401 # l
|
|||
)
|
||||
from litellm.proxy.utils import jsonify_object
|
||||
from litellm.repositories.budget_repository import BudgetRepository
|
||||
from litellm.repositories.table_repositories import CustomerOrganizationRepository
|
||||
|
||||
router: Final = APIRouter()
|
||||
|
||||
|
|
@ -371,6 +372,18 @@ async def delete_budget(
|
|||
detail={"error": f"{CommonProxyErrors.not_allowed_access.value}, your role={user_api_key_dict.user_role}"},
|
||||
)
|
||||
|
||||
customer_org_count: Final = await CustomerOrganizationRepository(prisma_client).table.count(
|
||||
where={"budget_ids": {"has": data.id}}
|
||||
)
|
||||
if customer_org_count:
|
||||
raise HTTPException(
|
||||
status_code=409,
|
||||
detail={
|
||||
"error": "Budget is a customer organization default budget. Delete it through "
|
||||
"/management/v1/customer_organizations/{customer_org_id}/default_budgets/{budget_id}"
|
||||
},
|
||||
)
|
||||
|
||||
response: Final = await BudgetRepository(prisma_client).table.delete(where={"budget_id": data.id})
|
||||
|
||||
return response
|
||||
|
|
|
|||
|
|
@ -8,6 +8,9 @@ from litellm.proxy.management.users.endpoints import router as current_user_rout
|
|||
from litellm.proxy.management_endpoints.management_v1.budgets import (
|
||||
router as budgets_router,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.management_v1.customers import (
|
||||
router as customers_router,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.management_v1.spend_logs import (
|
||||
router as spend_logs_router,
|
||||
)
|
||||
|
|
@ -20,6 +23,7 @@ from litellm.proxy.management_endpoints.management_v1.users import (
|
|||
|
||||
router: Final = APIRouter()
|
||||
router.include_router(budgets_router)
|
||||
router.include_router(customers_router)
|
||||
router.include_router(spend_logs_router)
|
||||
router.include_router(teams_router)
|
||||
router.include_router(users_router)
|
||||
|
|
|
|||
415
litellm/proxy/management_endpoints/management_v1/customers.py
Normal file
415
litellm/proxy/management_endpoints/management_v1/customers.py
Normal file
|
|
@ -0,0 +1,415 @@
|
|||
from math import ceil
|
||||
from typing import Annotated, Final, NoReturn, TypeVar
|
||||
|
||||
from fastapi import APIRouter, Depends, Query, Request, Response
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.list_api.common import (
|
||||
PROBLEM_TYPE_BASE,
|
||||
ManagementProblem,
|
||||
build_list_links,
|
||||
reject_unknown_query_params,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.management_v1.common import MANAGEMENT_V1_PREFIX
|
||||
from litellm.proxy.management_helpers.customer_hierarchy import CustomerHierarchy, ErrorCode, Failure, Page
|
||||
from litellm.proxy.management_helpers.utils import (
|
||||
management_endpoint_wrapper, # pyright: ignore[reportUnknownVariableType] # legacy decorator is untyped
|
||||
)
|
||||
from litellm.types.proxy.management_endpoints.customer_hierarchy import (
|
||||
BudgetTemplate,
|
||||
CustomerMember,
|
||||
MemberCreate,
|
||||
MemberPatch,
|
||||
Organization,
|
||||
OrganizationCreate,
|
||||
OrganizationPatch,
|
||||
Team,
|
||||
TeamCreate,
|
||||
TeamPatch,
|
||||
TemplatePatch,
|
||||
)
|
||||
from litellm.types.proxy.management_endpoints.management_v1 import ListMeta, ListResponse, ProblemDetail
|
||||
|
||||
T: Final = TypeVar("T")
|
||||
|
||||
|
||||
def error_details(code: ErrorCode) -> tuple[int, str, str]:
|
||||
match code:
|
||||
case ErrorCode.NOT_FOUND:
|
||||
return 404, "Not found", "Customer entity not found."
|
||||
case ErrorCode.PARENT_NOT_FOUND:
|
||||
return 404, "Parent not found", "Customer parent not found."
|
||||
case ErrorCode.ALREADY_EXISTS:
|
||||
return 409, "Already exists", "Customer entity ID already exists."
|
||||
case ErrorCode.INVALID_DEFAULT_BUDGET:
|
||||
return (
|
||||
422,
|
||||
"Invalid default budget",
|
||||
"Exactly one organization budget must be marked default.",
|
||||
)
|
||||
case ErrorCode.TEMPLATE_NOT_IN_ORG:
|
||||
return (
|
||||
422,
|
||||
"Template not in organization",
|
||||
"Budget template must belong to the member's organization.",
|
||||
)
|
||||
case ErrorCode.HAS_CHILDREN:
|
||||
return 409, "Has children", "Delete the customer's children before deleting the parent."
|
||||
case ErrorCode.INVALID_TEMPLATE_NAMES:
|
||||
return 422, "Invalid template names", "Budget template names must be unique within the organization."
|
||||
case ErrorCode.TEMPLATE_IN_USE:
|
||||
return 409, "Template in use", "A default or referenced template cannot be deleted."
|
||||
case ErrorCode.BUDGET_REQUIRED:
|
||||
return 422, "Budget required", "A member requires a template or a budget override."
|
||||
case ErrorCode.FORBIDDEN:
|
||||
return 403, "Forbidden", "Only proxy admins can manage customer hierarchies."
|
||||
case ErrorCode.DATABASE_UNAVAILABLE:
|
||||
return 503, "Database not connected", "Database not connected."
|
||||
case _:
|
||||
assert_never(code)
|
||||
|
||||
|
||||
def raise_public(failure: Failure) -> NoReturn:
|
||||
status, title, detail = error_details(failure.code)
|
||||
raise ManagementProblem(
|
||||
ProblemDetail(
|
||||
type=f"{PROBLEM_TYPE_BASE}customer-hierarchy-{failure.code.value.replace('_', '-')}",
|
||||
title=title,
|
||||
status=status,
|
||||
detail=detail,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def require_customer_admin(auth: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)]) -> UserAPIKeyAuth:
|
||||
if auth.user_role != LitellmUserRoles.PROXY_ADMIN:
|
||||
raise_public(Failure(ErrorCode.FORBIDDEN))
|
||||
return auth
|
||||
|
||||
|
||||
def customer_actor(auth: Annotated[UserAPIKeyAuth, Depends(require_customer_admin)]) -> str:
|
||||
return auth.user_id or "proxy_admin"
|
||||
|
||||
|
||||
def get_customer_hierarchy() -> CustomerHierarchy:
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
if prisma_client is None:
|
||||
raise_public(Failure(ErrorCode.DATABASE_UNAVAILABLE))
|
||||
return CustomerHierarchy(prisma_client)
|
||||
|
||||
|
||||
def resolve(result: T | Failure) -> T:
|
||||
if isinstance(result, Failure):
|
||||
raise_public(result)
|
||||
return result
|
||||
|
||||
|
||||
def list_response(result: Page[T], request: Request, page: int, page_size: int) -> ListResponse[T]:
|
||||
total_pages: Final = ceil(result.total_count / page_size)
|
||||
return ListResponse[T].model_validate(
|
||||
{
|
||||
"data": result.data,
|
||||
"meta": ListMeta(total_count=result.total_count, page=page, page_size=page_size, total_pages=total_pages),
|
||||
"links": build_list_links(request, page, total_pages),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def reject_duplicate_list_params(request: Request) -> None:
|
||||
repeated: Final = tuple(
|
||||
sorted(name for name in request.query_params if len(request.query_params.getlist(name)) > 1)
|
||||
)
|
||||
if repeated:
|
||||
raise ManagementProblem(
|
||||
ProblemDetail(
|
||||
type=f"{PROBLEM_TYPE_BASE}duplicate-query-parameter",
|
||||
title="Duplicate query parameter",
|
||||
status=400,
|
||||
detail=f"Repeated query parameter(s): {', '.join(repeated)}. Each may appear once.",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
router: Final = APIRouter(
|
||||
prefix=MANAGEMENT_V1_PREFIX,
|
||||
tags=["customer hierarchy"],
|
||||
dependencies=(Depends(require_customer_admin), Depends(reject_unknown_query_params)),
|
||||
)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/customer_organizations",
|
||||
status_code=201,
|
||||
response_model=Organization,
|
||||
description="Create a customer organization with resolved budgets. Proxy admins only.",
|
||||
)
|
||||
@management_endpoint_wrapper
|
||||
async def create_customer_organization(
|
||||
data: OrganizationCreate,
|
||||
actor: Annotated[str, Depends(customer_actor)],
|
||||
hierarchy: Annotated[CustomerHierarchy, Depends(get_customer_hierarchy)],
|
||||
) -> Organization:
|
||||
return resolve(await hierarchy.create_organization(data, actor))
|
||||
|
||||
|
||||
@router.get(
|
||||
"/customer_organizations/{customer_org_id}",
|
||||
response_model=Organization,
|
||||
description="Get a customer organization with resolved budgets. Proxy admins only.",
|
||||
)
|
||||
@management_endpoint_wrapper
|
||||
async def get_customer_organization(
|
||||
customer_org_id: str,
|
||||
hierarchy: Annotated[CustomerHierarchy, Depends(get_customer_hierarchy)],
|
||||
) -> Organization:
|
||||
return resolve(await hierarchy.get_organization(customer_org_id))
|
||||
|
||||
|
||||
@router.patch(
|
||||
"/customer_organizations/{customer_org_id}",
|
||||
response_model=Organization,
|
||||
description="Update a customer organization alias or overall budget. Proxy admins only.",
|
||||
)
|
||||
@management_endpoint_wrapper
|
||||
async def update_customer_organization(
|
||||
customer_org_id: str,
|
||||
data: OrganizationPatch,
|
||||
actor: Annotated[str, Depends(customer_actor)],
|
||||
hierarchy: Annotated[CustomerHierarchy, Depends(get_customer_hierarchy)],
|
||||
) -> Organization:
|
||||
return resolve(await hierarchy.update_organization(customer_org_id, data, actor))
|
||||
|
||||
|
||||
@router.delete(
|
||||
"/customer_organizations/{customer_org_id}",
|
||||
status_code=204,
|
||||
description="Delete a customer organization and its owned budgets. Parents with children return 409. Proxy admins only.",
|
||||
)
|
||||
@management_endpoint_wrapper
|
||||
async def delete_customer_organization(
|
||||
customer_org_id: str,
|
||||
hierarchy: Annotated[CustomerHierarchy, Depends(get_customer_hierarchy)],
|
||||
) -> Response:
|
||||
resolve(await hierarchy.delete_organization(customer_org_id))
|
||||
return Response(status_code=204)
|
||||
|
||||
|
||||
@router.get(
|
||||
"/customer_organizations",
|
||||
response_model=ListResponse[Organization],
|
||||
dependencies=(Depends(reject_duplicate_list_params),),
|
||||
description="List customer organizations with page-number pagination and resolved budgets. Proxy admins only.",
|
||||
)
|
||||
@management_endpoint_wrapper
|
||||
async def list_customer_organizations(
|
||||
request: Request,
|
||||
hierarchy: Annotated[CustomerHierarchy, Depends(get_customer_hierarchy)],
|
||||
page: Annotated[int, Query(ge=1)] = 1,
|
||||
page_size: Annotated[int, Query(ge=1)] = 50,
|
||||
) -> ListResponse[Organization]:
|
||||
size: Final = min(page_size, 100)
|
||||
result: Final = resolve(await hierarchy.list_organizations(page, size))
|
||||
return list_response(result, request, page, size)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/customer_teams",
|
||||
status_code=201,
|
||||
response_model=Team,
|
||||
description="Create a customer team with resolved budgets. Proxy admins only.",
|
||||
)
|
||||
@management_endpoint_wrapper
|
||||
async def create_customer_team(
|
||||
data: TeamCreate,
|
||||
actor: Annotated[str, Depends(customer_actor)],
|
||||
hierarchy: Annotated[CustomerHierarchy, Depends(get_customer_hierarchy)],
|
||||
) -> Team:
|
||||
return resolve(await hierarchy.create_team(data, actor))
|
||||
|
||||
|
||||
@router.get(
|
||||
"/customer_teams/{customer_team_id}",
|
||||
response_model=Team,
|
||||
description="Get a customer team with resolved budgets. Proxy admins only.",
|
||||
)
|
||||
@management_endpoint_wrapper
|
||||
async def get_customer_team(
|
||||
customer_team_id: str,
|
||||
hierarchy: Annotated[CustomerHierarchy, Depends(get_customer_hierarchy)],
|
||||
) -> Team:
|
||||
return resolve(await hierarchy.get_team(customer_team_id))
|
||||
|
||||
|
||||
@router.patch(
|
||||
"/customer_teams/{customer_team_id}",
|
||||
response_model=Team,
|
||||
description="Update a customer team and its budgets. Proxy admins only.",
|
||||
)
|
||||
@management_endpoint_wrapper
|
||||
async def update_customer_team(
|
||||
customer_team_id: str,
|
||||
data: TeamPatch,
|
||||
actor: Annotated[str, Depends(customer_actor)],
|
||||
hierarchy: Annotated[CustomerHierarchy, Depends(get_customer_hierarchy)],
|
||||
) -> Team:
|
||||
return resolve(await hierarchy.update_team(customer_team_id, data, actor))
|
||||
|
||||
|
||||
@router.delete(
|
||||
"/customer_teams/{customer_team_id}",
|
||||
status_code=204,
|
||||
description="Delete a customer team and its owned budgets. Parents with children return 409. Proxy admins only.",
|
||||
)
|
||||
@management_endpoint_wrapper
|
||||
async def delete_customer_team(
|
||||
customer_team_id: str,
|
||||
hierarchy: Annotated[CustomerHierarchy, Depends(get_customer_hierarchy)],
|
||||
) -> Response:
|
||||
resolve(await hierarchy.delete_team(customer_team_id))
|
||||
return Response(status_code=204)
|
||||
|
||||
|
||||
@router.get(
|
||||
"/customer_teams",
|
||||
response_model=ListResponse[Team],
|
||||
dependencies=(Depends(reject_duplicate_list_params),),
|
||||
description="List customer teams with page-number pagination and resolved budgets. Proxy admins only.",
|
||||
)
|
||||
@management_endpoint_wrapper
|
||||
async def list_customer_teams(
|
||||
request: Request,
|
||||
hierarchy: Annotated[CustomerHierarchy, Depends(get_customer_hierarchy)],
|
||||
page: Annotated[int, Query(ge=1)] = 1,
|
||||
page_size: Annotated[int, Query(ge=1)] = 50,
|
||||
parent_org_id: str | None = None,
|
||||
) -> ListResponse[Team]:
|
||||
size: Final = min(page_size, 100)
|
||||
result: Final = resolve(await hierarchy.list_teams(parent_org_id, page, size))
|
||||
return list_response(result, request, page, size)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/customer_members",
|
||||
status_code=201,
|
||||
response_model=CustomerMember,
|
||||
description="Create a customer member with resolved budgets. Proxy admins only.",
|
||||
)
|
||||
@management_endpoint_wrapper
|
||||
async def create_customer_member(
|
||||
data: MemberCreate,
|
||||
actor: Annotated[str, Depends(customer_actor)],
|
||||
hierarchy: Annotated[CustomerHierarchy, Depends(get_customer_hierarchy)],
|
||||
) -> CustomerMember:
|
||||
return resolve(await hierarchy.create_member(data, actor))
|
||||
|
||||
|
||||
@router.get(
|
||||
"/customer_members/{customer_member_id}",
|
||||
response_model=CustomerMember,
|
||||
description="Get a customer member with resolved budgets. Proxy admins only.",
|
||||
)
|
||||
@management_endpoint_wrapper
|
||||
async def get_customer_member(
|
||||
customer_member_id: str,
|
||||
hierarchy: Annotated[CustomerHierarchy, Depends(get_customer_hierarchy)],
|
||||
) -> CustomerMember:
|
||||
return resolve(await hierarchy.get_member(customer_member_id))
|
||||
|
||||
|
||||
@router.patch(
|
||||
"/customer_members/{customer_member_id}",
|
||||
response_model=CustomerMember,
|
||||
description="Update a customer member and its budgets. Proxy admins only.",
|
||||
)
|
||||
@management_endpoint_wrapper
|
||||
async def update_customer_member(
|
||||
customer_member_id: str,
|
||||
data: MemberPatch,
|
||||
actor: Annotated[str, Depends(customer_actor)],
|
||||
hierarchy: Annotated[CustomerHierarchy, Depends(get_customer_hierarchy)],
|
||||
) -> CustomerMember:
|
||||
return resolve(await hierarchy.update_member(customer_member_id, data, actor))
|
||||
|
||||
|
||||
@router.delete(
|
||||
"/customer_members/{customer_member_id}",
|
||||
status_code=204,
|
||||
description="Delete a customer member and its owned budgets. Parents with children return 409. Proxy admins only.",
|
||||
)
|
||||
@management_endpoint_wrapper
|
||||
async def delete_customer_member(
|
||||
customer_member_id: str,
|
||||
hierarchy: Annotated[CustomerHierarchy, Depends(get_customer_hierarchy)],
|
||||
) -> Response:
|
||||
resolve(await hierarchy.delete_member(customer_member_id))
|
||||
return Response(status_code=204)
|
||||
|
||||
|
||||
@router.get(
|
||||
"/customer_members",
|
||||
response_model=ListResponse[CustomerMember],
|
||||
dependencies=(Depends(reject_duplicate_list_params),),
|
||||
description="List customer members with page-number pagination and resolved budgets. Proxy admins only.",
|
||||
)
|
||||
@management_endpoint_wrapper
|
||||
async def list_customer_members(
|
||||
request: Request,
|
||||
hierarchy: Annotated[CustomerHierarchy, Depends(get_customer_hierarchy)],
|
||||
page: Annotated[int, Query(ge=1)] = 1,
|
||||
page_size: Annotated[int, Query(ge=1)] = 50,
|
||||
parent_team_id: str | None = None,
|
||||
) -> ListResponse[CustomerMember]:
|
||||
size: Final = min(page_size, 100)
|
||||
result: Final = resolve(await hierarchy.list_members(parent_team_id, page, size))
|
||||
return list_response(result, request, page, size)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/customer_organizations/{customer_org_id}/default_budgets",
|
||||
status_code=201,
|
||||
response_model=Organization,
|
||||
description="Add a named budget template to a customer organization. Proxy admins only.",
|
||||
)
|
||||
@management_endpoint_wrapper
|
||||
async def add_customer_budget_template(
|
||||
customer_org_id: str,
|
||||
data: BudgetTemplate,
|
||||
actor: Annotated[str, Depends(customer_actor)],
|
||||
hierarchy: Annotated[CustomerHierarchy, Depends(get_customer_hierarchy)],
|
||||
) -> Organization:
|
||||
return resolve(await hierarchy.add_template(customer_org_id, data, actor))
|
||||
|
||||
|
||||
@router.patch(
|
||||
"/customer_organizations/{customer_org_id}/default_budgets/{budget_id}",
|
||||
response_model=Organization,
|
||||
description="Edit a budget template in place. Members inherit the updated values. Proxy admins only.",
|
||||
)
|
||||
@management_endpoint_wrapper
|
||||
async def update_customer_budget_template(
|
||||
customer_org_id: str,
|
||||
budget_id: str,
|
||||
data: TemplatePatch,
|
||||
actor: Annotated[str, Depends(customer_actor)],
|
||||
hierarchy: Annotated[CustomerHierarchy, Depends(get_customer_hierarchy)],
|
||||
) -> Organization:
|
||||
return resolve(await hierarchy.update_template(customer_org_id, budget_id, data, actor))
|
||||
|
||||
|
||||
@router.delete(
|
||||
"/customer_organizations/{customer_org_id}/default_budgets/{budget_id}",
|
||||
status_code=204,
|
||||
description="Delete an unreferenced, non-default customer budget template. Proxy admins only.",
|
||||
)
|
||||
@management_endpoint_wrapper
|
||||
async def delete_customer_budget_template(
|
||||
customer_org_id: str,
|
||||
budget_id: str,
|
||||
actor: Annotated[str, Depends(customer_actor)],
|
||||
hierarchy: Annotated[CustomerHierarchy, Depends(get_customer_hierarchy)],
|
||||
) -> Response:
|
||||
resolve(await hierarchy.delete_template(customer_org_id, budget_id, actor))
|
||||
return Response(status_code=204)
|
||||
692
litellm/proxy/management_helpers/customer_hierarchy.py
Normal file
692
litellm/proxy/management_helpers/customer_hierarchy.py
Normal file
|
|
@ -0,0 +1,692 @@
|
|||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||
from contextlib import AbstractAsyncContextManager
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from enum import Enum
|
||||
from itertools import chain
|
||||
from typing import TYPE_CHECKING, Final, Generic, Protocol, TypeVar
|
||||
from uuid import uuid4
|
||||
|
||||
from pydantic import ConfigDict, Field, TypeAdapter, model_validator
|
||||
|
||||
from litellm.models.budget import LiteLLM_BudgetTable
|
||||
from litellm.repositories.budget_repository import create_budget_row, update_budget_row
|
||||
from litellm.repositories.chunked_in import find_many_in
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
from litellm.types.proxy.management_endpoints.customer_hierarchy import (
|
||||
Budget,
|
||||
BudgetTemplate,
|
||||
CustomerMember,
|
||||
MemberCreate,
|
||||
MemberPatch,
|
||||
Organization,
|
||||
OrganizationBudget,
|
||||
OrganizationCreate,
|
||||
OrganizationPatch,
|
||||
Team,
|
||||
TeamCreate,
|
||||
TeamPatch,
|
||||
TemplatePatch,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from prisma import Prisma
|
||||
|
||||
|
||||
class ErrorCode(str, Enum):
|
||||
NOT_FOUND = "not_found"
|
||||
PARENT_NOT_FOUND = "parent_not_found"
|
||||
ALREADY_EXISTS = "already_exists"
|
||||
INVALID_DEFAULT_BUDGET = "invalid_default_budget"
|
||||
TEMPLATE_NOT_IN_ORG = "template_not_in_org"
|
||||
HAS_CHILDREN = "has_children"
|
||||
FORBIDDEN = "forbidden"
|
||||
DATABASE_UNAVAILABLE = "database_unavailable"
|
||||
INVALID_TEMPLATE_NAMES = "invalid_template_names"
|
||||
TEMPLATE_IN_USE = "template_in_use"
|
||||
BUDGET_REQUIRED = "budget_required"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Failure:
|
||||
code: ErrorCode
|
||||
|
||||
|
||||
_LOCK_ORGANIZATION_SQL: Final = (
|
||||
'SELECT 1 AS locked FROM "LiteLLM_CustomerOrganizationTable" WHERE customer_org_id = $1 FOR UPDATE'
|
||||
)
|
||||
|
||||
|
||||
class TransactionOperation(str, Enum):
|
||||
READ_WRITE = "read_write"
|
||||
DELETE = "delete"
|
||||
|
||||
|
||||
class PrismaRecord(Protocol):
|
||||
def model_dump(self) -> Mapping[str, object]: ...
|
||||
|
||||
|
||||
class TransactionSource(Protocol):
|
||||
def tx(self) -> AbstractAsyncContextManager["Prisma"]: ...
|
||||
|
||||
|
||||
T: Final = TypeVar("T")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Page(Generic[T]):
|
||||
data: tuple[T, ...]
|
||||
total_count: int
|
||||
|
||||
|
||||
class BudgetRecord(Budget):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
|
||||
_VALUES: Final = TypeAdapter(Mapping[str, object])
|
||||
_BUDGET: Final[TypeAdapter[BudgetRecord | None]] = TypeAdapter(BudgetRecord | None)
|
||||
_EXPIRY: Final[TypeAdapter[datetime | None]] = TypeAdapter(datetime | None)
|
||||
|
||||
|
||||
def budget_values(budget: Budget) -> Mapping[str, object]:
|
||||
return _VALUES.validate_python(budget.model_dump())
|
||||
|
||||
|
||||
def related_budget(value: object, now: datetime | None = None, inherited_maximum: float | None = None) -> Budget | None:
|
||||
record: Final = _BUDGET.validate_python(value)
|
||||
if record is None:
|
||||
return None
|
||||
if now is None:
|
||||
return Budget.model_validate(record.model_dump())
|
||||
budget: Final = LiteLLM_BudgetTable.model_validate(value)
|
||||
maximum: Final = budget.effective_max_budget(utc_datetime(now))
|
||||
if maximum is not None:
|
||||
return Budget.model_validate({**record.model_dump(), "max_budget": maximum})
|
||||
if inherited_maximum is None:
|
||||
return Budget.model_validate(record.model_dump())
|
||||
return Budget.model_validate(
|
||||
{
|
||||
**record.model_dump(),
|
||||
"max_budget": inherited_maximum + budget.active_temp_budget_increase(utc_datetime(now)),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class OrganizationRecord(Organization):
|
||||
budget_ids: tuple[str, ...] = ()
|
||||
budget_names: Mapping[str, str] = Field(default_factory=dict)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def valid_ownership(self) -> "OrganizationRecord":
|
||||
if set(self.budget_names) != set(self.budget_ids):
|
||||
raise ValueError("Budget names must match the organization's budget IDs")
|
||||
if len(set(self.budget_names.values())) != len(self.budget_ids):
|
||||
raise ValueError("Template names must be unique")
|
||||
if self.default_budget_id is not None and self.default_budget_id not in self.budget_ids:
|
||||
raise ValueError("Default budget must belong to the organization")
|
||||
return self
|
||||
|
||||
|
||||
class StoredBudget(BudgetRecord):
|
||||
budget_id: str
|
||||
|
||||
|
||||
def organization_response(row: OrganizationRecord, budgets: Mapping[str, StoredBudget]) -> Organization:
|
||||
return Organization.model_validate(
|
||||
{
|
||||
**row.model_dump(exclude={"budget_ids", "budget_names"}),
|
||||
"default_budgets": tuple(
|
||||
OrganizationBudget.model_validate(
|
||||
{
|
||||
**budgets[budget_id].model_dump(),
|
||||
"name": row.budget_names[budget_id],
|
||||
"is_default": budget_id == row.default_budget_id,
|
||||
}
|
||||
)
|
||||
for budget_id in row.budget_ids
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def team_response(row: PrismaRecord) -> Team:
|
||||
values: Final = row.model_dump()
|
||||
return Team.model_validate({**values, "budget": related_budget(values.get("budget"))})
|
||||
|
||||
|
||||
def utc_now() -> datetime:
|
||||
return datetime.now(timezone.utc)
|
||||
|
||||
|
||||
def utc_datetime(value: datetime) -> datetime:
|
||||
return value.replace(tzinfo=timezone.utc) if value.tzinfo is None else value.astimezone(timezone.utc)
|
||||
|
||||
|
||||
def effective_member_budget(
|
||||
template: Budget | None, override: Budget | None, expires_at: datetime | None, now: datetime
|
||||
) -> Budget:
|
||||
inherited: Final = template if template is not None else Budget(allowed_models=())
|
||||
if override is None:
|
||||
return inherited
|
||||
if expires_at is not None and utc_datetime(expires_at) <= utc_datetime(now):
|
||||
return inherited
|
||||
values: Final = _VALUES.validate_python(override.model_dump(exclude_none=True, exclude={"allowed_models"}))
|
||||
models: Final[Mapping[str, object]] = {"allowed_models": override.allowed_models} if override.allowed_models else {}
|
||||
return Budget.model_validate({**budget_values(inherited), **values, **models})
|
||||
|
||||
|
||||
def member_response(row: PrismaRecord, now: datetime) -> CustomerMember:
|
||||
values: Final = row.model_dump()
|
||||
override: Final = related_budget(values.get("budget"))
|
||||
template: Final = related_budget(values.get("budget_template"), now)
|
||||
effective_override: Final = related_budget(
|
||||
values.get("budget"), now, template.max_budget if template is not None else None
|
||||
)
|
||||
expiry: Final = _EXPIRY.validate_python(values.get("budget_override_expires_at"))
|
||||
return CustomerMember.model_validate(
|
||||
{
|
||||
**values,
|
||||
"budget_override": override,
|
||||
"effective_budget": effective_member_budget(template, effective_override, expiry, now),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def default_template_index(templates: tuple[BudgetTemplate, ...]) -> int | Failure | None:
|
||||
if not templates:
|
||||
return None
|
||||
flagged: Final = tuple(index for index, template in enumerate(templates) if template.is_default)
|
||||
if len(flagged) > 1:
|
||||
return Failure(ErrorCode.INVALID_DEFAULT_BUDGET)
|
||||
if flagged:
|
||||
return flagged[0]
|
||||
if len(templates) == 1:
|
||||
return 0
|
||||
return Failure(ErrorCode.INVALID_DEFAULT_BUDGET)
|
||||
|
||||
|
||||
def alias_patch(payload: OrganizationPatch | TeamPatch | MemberPatch) -> Mapping[str, object]:
|
||||
return {"alias": payload.alias} if "alias" in payload.model_fields_set else {}
|
||||
|
||||
|
||||
_BUDGET_INCLUDE: Final = {"budget": True}
|
||||
_MEMBER_INCLUDE: Final = {"budget": True, "budget_template": True}
|
||||
|
||||
|
||||
class CustomerHierarchyStore:
|
||||
def __init__(self, db: "Prisma", now: datetime) -> None:
|
||||
self.db = db
|
||||
self.now = now
|
||||
self.organizations: TableActions[PrismaRecord] = db.litellm_customerorganizationtable # pyright: ignore[reportAttributeAccessIssue] # protocol accepts immutable query mappings
|
||||
self.teams: TableActions[PrismaRecord] = db.litellm_customerteamtable # pyright: ignore[reportAttributeAccessIssue] # protocol accepts immutable query mappings
|
||||
self.members: TableActions[PrismaRecord] = db.litellm_customermembertable # pyright: ignore[reportAttributeAccessIssue] # protocol accepts immutable query mappings
|
||||
self.budgets: TableActions[PrismaRecord] = db.litellm_budgettable # pyright: ignore[reportAttributeAccessIssue] # protocol accepts immutable query mappings
|
||||
|
||||
async def organization_responses(self, rows: Sequence[PrismaRecord]) -> tuple[Organization, ...]:
|
||||
records: Final = tuple(
|
||||
OrganizationRecord.model_validate(
|
||||
{**row.model_dump(), "budget": related_budget(row.model_dump().get("budget"))}
|
||||
)
|
||||
for row in rows
|
||||
)
|
||||
budget_ids: Final = tuple(chain.from_iterable(record.budget_ids for record in records))
|
||||
budget_rows: Final = await find_many_in(self.budgets, "budget_id", budget_ids)
|
||||
resolved: Final = tuple(StoredBudget.model_validate(row.model_dump()) for row in budget_rows)
|
||||
budgets: Final = {budget.budget_id: budget for budget in resolved}
|
||||
return tuple(organization_response(record, budgets) for record in records)
|
||||
|
||||
async def get_organization(self, organization_id: str) -> Organization | Failure:
|
||||
row: Final = await self.organizations.find_unique(
|
||||
where={"customer_org_id": organization_id}, include=_BUDGET_INCLUDE
|
||||
)
|
||||
if row is None:
|
||||
return Failure(ErrorCode.NOT_FOUND)
|
||||
organizations: Final = await self.organization_responses((row,))
|
||||
return organizations[0]
|
||||
|
||||
async def get_team(self, team_id: str) -> Team | Failure:
|
||||
row: Final = await self.teams.find_unique(where={"customer_team_id": team_id}, include=_BUDGET_INCLUDE)
|
||||
return team_response(row) if row is not None else Failure(ErrorCode.NOT_FOUND)
|
||||
|
||||
async def get_member(self, member_id: str) -> CustomerMember | Failure:
|
||||
row: Final = await self.members.find_unique(where={"customer_member_id": member_id}, include=_MEMBER_INCLUDE)
|
||||
return member_response(row, self.now) if row is not None else Failure(ErrorCode.NOT_FOUND)
|
||||
|
||||
async def create_budget(self, budget: Budget, actor: str) -> str:
|
||||
budget_id: Final = str(uuid4())
|
||||
await create_budget_row(
|
||||
self.budgets,
|
||||
{
|
||||
**budget_values(budget),
|
||||
"allowed_models": budget.allowed_models or (),
|
||||
"budget_id": budget_id,
|
||||
},
|
||||
actor,
|
||||
)
|
||||
return budget_id
|
||||
|
||||
async def update_budget(self, budget_id: str | None, budget: Budget, actor: str) -> str:
|
||||
if budget_id is None:
|
||||
return await self.create_budget(budget, actor)
|
||||
patch: Final = _VALUES.validate_python(budget.model_dump(exclude_unset=True))
|
||||
await update_budget_row(
|
||||
self.budgets,
|
||||
budget_id,
|
||||
{
|
||||
**patch,
|
||||
**({"allowed_models": budget.allowed_models or ()} if "allowed_models" in patch else {}),
|
||||
},
|
||||
actor,
|
||||
)
|
||||
return budget_id
|
||||
|
||||
async def write_template_metadata(
|
||||
self, organization_id: str, names: Mapping[str, str], default_id: str | None, actor: str
|
||||
) -> None:
|
||||
from prisma import Json
|
||||
|
||||
await self.organizations.update(
|
||||
where={"customer_org_id": organization_id},
|
||||
data={
|
||||
"budget_ids": tuple(names),
|
||||
"budget_names": Json(dict(names)),
|
||||
"default_budget_id": default_id,
|
||||
"updated_by": actor,
|
||||
},
|
||||
)
|
||||
|
||||
async def lock_organization(self, organization_id: str) -> None:
|
||||
await self.db.query_raw(_LOCK_ORGANIZATION_SQL, organization_id)
|
||||
|
||||
async def create_organization(self, payload: OrganizationCreate, actor: str) -> Organization | Failure:
|
||||
default_index: Final = default_template_index(payload.default_budgets)
|
||||
if isinstance(default_index, Failure):
|
||||
return default_index
|
||||
names: Final = tuple(template.name for template in payload.default_budgets)
|
||||
if len(set(names)) != len(names):
|
||||
return Failure(ErrorCode.INVALID_TEMPLATE_NAMES)
|
||||
organization_id: Final = payload.customer_org_id or str(uuid4())
|
||||
overall_id: Final = await self.create_budget(payload.budget, actor) if payload.budget is not None else None
|
||||
await self.organizations.create(
|
||||
data={
|
||||
"customer_org_id": organization_id,
|
||||
"alias": payload.alias,
|
||||
"budget_id": overall_id,
|
||||
"created_by": actor,
|
||||
"updated_by": actor,
|
||||
}
|
||||
)
|
||||
created: Final = tuple(
|
||||
[
|
||||
await self.create_budget(
|
||||
Budget.model_validate(template.model_dump(exclude={"name", "is_default"})), actor
|
||||
)
|
||||
for template in payload.default_budgets
|
||||
]
|
||||
)
|
||||
await self.write_template_metadata(
|
||||
organization_id,
|
||||
dict(zip(created, names)),
|
||||
created[default_index] if default_index is not None else None,
|
||||
actor,
|
||||
)
|
||||
return await self.get_organization(organization_id)
|
||||
|
||||
async def overall_budget_patch(
|
||||
self, current_id: str | None, payload: OrganizationPatch | TeamPatch, actor: str
|
||||
) -> str | None:
|
||||
if "budget" not in payload.model_fields_set:
|
||||
return current_id
|
||||
if payload.budget is None:
|
||||
return None
|
||||
return await self.update_budget(current_id, payload.budget, actor)
|
||||
|
||||
async def delete_released_budget(self, previous_id: str | None, current_id: str | None) -> None:
|
||||
if previous_id is not None and current_id is None:
|
||||
await self.budgets.delete(where={"budget_id": previous_id})
|
||||
|
||||
async def update_organization(
|
||||
self, organization_id: str, payload: OrganizationPatch, actor: str
|
||||
) -> Organization | Failure:
|
||||
organization: Final = await self.get_organization(organization_id)
|
||||
if isinstance(organization, Failure):
|
||||
return organization
|
||||
budget_id: Final = await self.overall_budget_patch(organization.budget_id, payload, actor)
|
||||
await self.organizations.update(
|
||||
where={"customer_org_id": organization_id},
|
||||
data={**alias_patch(payload), "budget_id": budget_id, "updated_by": actor},
|
||||
)
|
||||
await self.delete_released_budget(organization.budget_id, budget_id)
|
||||
return await self.get_organization(organization_id)
|
||||
|
||||
async def add_template(self, organization_id: str, payload: BudgetTemplate, actor: str) -> Organization | Failure:
|
||||
await self.lock_organization(organization_id)
|
||||
organization: Final = await self.get_organization(organization_id)
|
||||
if isinstance(organization, Failure):
|
||||
return organization
|
||||
names: Final = {template.budget_id: template.name for template in organization.default_budgets}
|
||||
if payload.name in names.values():
|
||||
return Failure(ErrorCode.INVALID_TEMPLATE_NAMES)
|
||||
budget_id: Final = await self.create_budget(
|
||||
Budget.model_validate(payload.model_dump(exclude={"name", "is_default"})), actor
|
||||
)
|
||||
default_id: Final = budget_id if payload.is_default or not names else organization.default_budget_id
|
||||
await self.write_template_metadata(organization_id, {**names, budget_id: payload.name}, default_id, actor)
|
||||
return await self.get_organization(organization_id)
|
||||
|
||||
async def update_template(
|
||||
self, organization_id: str, budget_id: str, payload: TemplatePatch, actor: str
|
||||
) -> Organization | Failure:
|
||||
await self.lock_organization(organization_id)
|
||||
organization: Final = await self.get_organization(organization_id)
|
||||
if isinstance(organization, Failure):
|
||||
return organization
|
||||
names: Final = {template.budget_id: template.name for template in organization.default_budgets}
|
||||
if budget_id not in names:
|
||||
return Failure(ErrorCode.TEMPLATE_NOT_IN_ORG)
|
||||
name: Final = payload.name if payload.name is not None else names[budget_id]
|
||||
if any(existing_name == name for existing_id, existing_name in names.items() if existing_id != budget_id):
|
||||
return Failure(ErrorCode.INVALID_TEMPLATE_NAMES)
|
||||
if payload.is_default is False and organization.default_budget_id == budget_id:
|
||||
return Failure(ErrorCode.INVALID_DEFAULT_BUDGET)
|
||||
patch: Final = Budget.model_validate(payload.model_dump(exclude_unset=True, exclude={"name", "is_default"}))
|
||||
await self.update_budget(budget_id, patch, actor)
|
||||
default_id: Final = budget_id if payload.is_default is True else organization.default_budget_id
|
||||
await self.write_template_metadata(organization_id, {**names, budget_id: name}, default_id, actor)
|
||||
return await self.get_organization(organization_id)
|
||||
|
||||
async def delete_template(self, organization_id: str, budget_id: str, actor: str) -> None | Failure:
|
||||
await self.lock_organization(organization_id)
|
||||
organization: Final = await self.get_organization(organization_id)
|
||||
if isinstance(organization, Failure):
|
||||
return organization
|
||||
names: Final = {template.budget_id: template.name for template in organization.default_budgets}
|
||||
if budget_id not in names:
|
||||
return Failure(ErrorCode.TEMPLATE_NOT_IN_ORG)
|
||||
if organization.default_budget_id == budget_id:
|
||||
return Failure(ErrorCode.TEMPLATE_IN_USE)
|
||||
if await self.members.count(where={"budget_template_id": budget_id}):
|
||||
return Failure(ErrorCode.TEMPLATE_IN_USE)
|
||||
remaining: Final = {key: name for key, name in names.items() if key != budget_id}
|
||||
await self.write_template_metadata(organization_id, remaining, organization.default_budget_id, actor)
|
||||
await self.budgets.delete(where={"budget_id": budget_id})
|
||||
return None
|
||||
|
||||
async def resolve_member_template(self, team_id: str, template_id: str | None) -> str | Failure | None:
|
||||
team: Final = await self.get_team(team_id)
|
||||
if isinstance(team, Failure):
|
||||
return Failure(ErrorCode.PARENT_NOT_FOUND)
|
||||
organization: Final = await self.get_organization(team.parent_org_id)
|
||||
if isinstance(organization, Failure):
|
||||
return Failure(ErrorCode.PARENT_NOT_FOUND)
|
||||
selected: Final = template_id if template_id is not None else organization.default_budget_id
|
||||
if selected is None:
|
||||
return None
|
||||
if not any(template.budget_id == selected for template in organization.default_budgets):
|
||||
return Failure(ErrorCode.TEMPLATE_NOT_IN_ORG)
|
||||
return selected
|
||||
|
||||
async def create_team(self, payload: TeamCreate, actor: str) -> Team | Failure:
|
||||
parent: Final = await self.organizations.find_unique(where={"customer_org_id": payload.parent_org_id})
|
||||
if parent is None:
|
||||
return Failure(ErrorCode.PARENT_NOT_FOUND)
|
||||
budget_id: Final = await self.create_budget(payload.budget, actor) if payload.budget is not None else None
|
||||
team_id: Final = payload.customer_team_id or str(uuid4())
|
||||
await self.teams.create(
|
||||
data={
|
||||
"customer_team_id": team_id,
|
||||
"alias": payload.alias,
|
||||
"parent_org_id": payload.parent_org_id,
|
||||
"budget_id": budget_id,
|
||||
"created_by": actor,
|
||||
"updated_by": actor,
|
||||
}
|
||||
)
|
||||
return await self.get_team(team_id)
|
||||
|
||||
async def update_team(self, team_id: str, payload: TeamPatch, actor: str) -> Team | Failure:
|
||||
team: Final = await self.get_team(team_id)
|
||||
if isinstance(team, Failure):
|
||||
return team
|
||||
budget_id: Final = await self.overall_budget_patch(team.budget_id, payload, actor)
|
||||
await self.teams.update(
|
||||
where={"customer_team_id": team_id},
|
||||
data={**alias_patch(payload), "budget_id": budget_id, "updated_by": actor},
|
||||
)
|
||||
await self.delete_released_budget(team.budget_id, budget_id)
|
||||
return await self.get_team(team_id)
|
||||
|
||||
async def create_member(self, payload: MemberCreate, actor: str) -> CustomerMember | Failure:
|
||||
template_id: Final = await self.resolve_member_template(payload.parent_team_id, payload.budget_template_id)
|
||||
if isinstance(template_id, Failure):
|
||||
return template_id
|
||||
if template_id is None and payload.budget_override is None:
|
||||
return Failure(ErrorCode.BUDGET_REQUIRED)
|
||||
budget_id: Final = (
|
||||
await self.create_budget(payload.budget_override, actor) if payload.budget_override is not None else None
|
||||
)
|
||||
member_id: Final = payload.customer_member_id or str(uuid4())
|
||||
await self.members.create(
|
||||
data={
|
||||
"customer_member_id": member_id,
|
||||
"alias": payload.alias,
|
||||
"parent_team_id": payload.parent_team_id,
|
||||
"budget_template_id": template_id,
|
||||
"budget_id": budget_id,
|
||||
"budget_override_expires_at": payload.budget_override_expires_at,
|
||||
"created_by": actor,
|
||||
"updated_by": actor,
|
||||
}
|
||||
)
|
||||
return await self.get_member(member_id)
|
||||
|
||||
async def member_template_patch(self, member: CustomerMember, payload: MemberPatch) -> str | Failure | None:
|
||||
if "budget_template_id" not in payload.model_fields_set:
|
||||
return member.budget_template_id
|
||||
if payload.budget_template_id is None:
|
||||
return None
|
||||
return await self.resolve_member_template(member.parent_team_id, payload.budget_template_id)
|
||||
|
||||
async def member_override_patch(self, member: CustomerMember, payload: MemberPatch, actor: str) -> str | None:
|
||||
if "budget_override" not in payload.model_fields_set:
|
||||
return member.budget_id
|
||||
if payload.budget_override is None:
|
||||
return None
|
||||
return await self.update_budget(member.budget_id, payload.budget_override, actor)
|
||||
|
||||
async def update_member(self, member_id: str, payload: MemberPatch, actor: str) -> CustomerMember | Failure:
|
||||
member: Final = await self.get_member(member_id)
|
||||
if isinstance(member, Failure):
|
||||
return member
|
||||
selected: Final = await self.member_template_patch(member, payload)
|
||||
if isinstance(selected, Failure):
|
||||
return selected
|
||||
override: Final = (
|
||||
payload.budget_override if "budget_override" in payload.model_fields_set else member.budget_override
|
||||
)
|
||||
if selected is None and override is None:
|
||||
return Failure(ErrorCode.BUDGET_REQUIRED)
|
||||
budget_id: Final = await self.member_override_patch(member, payload, actor)
|
||||
expiry: Final = (
|
||||
payload.budget_override_expires_at
|
||||
if "budget_override_expires_at" in payload.model_fields_set
|
||||
else member.budget_override_expires_at
|
||||
if budget_id is not None
|
||||
else None
|
||||
)
|
||||
await self.members.update(
|
||||
where={"customer_member_id": member_id},
|
||||
data={
|
||||
**alias_patch(payload),
|
||||
"budget_id": budget_id,
|
||||
"budget_template_id": selected,
|
||||
"budget_override_expires_at": expiry,
|
||||
"updated_by": actor,
|
||||
},
|
||||
)
|
||||
await self.delete_released_budget(member.budget_id, budget_id)
|
||||
return await self.get_member(member_id)
|
||||
|
||||
async def list_organizations(self, page: int, page_size: int) -> Page[Organization]:
|
||||
count: Final = await self.organizations.count()
|
||||
rows: Final = await self.organizations.find_many(
|
||||
skip=(page - 1) * page_size,
|
||||
take=page_size,
|
||||
order=[{"created_at": "asc"}, {"customer_org_id": "asc"}],
|
||||
include=_BUDGET_INCLUDE,
|
||||
)
|
||||
return Page(await self.organization_responses(rows), count)
|
||||
|
||||
async def list_teams(self, parent_org_id: str | None, page: int, page_size: int) -> Page[Team]:
|
||||
where: Final = {"parent_org_id": parent_org_id} if parent_org_id is not None else {}
|
||||
count: Final = await self.teams.count(where=where)
|
||||
rows: Final = await self.teams.find_many(
|
||||
where=where,
|
||||
skip=(page - 1) * page_size,
|
||||
take=page_size,
|
||||
order=[{"created_at": "asc"}, {"customer_team_id": "asc"}],
|
||||
include=_BUDGET_INCLUDE,
|
||||
)
|
||||
return Page(tuple(team_response(row) for row in rows), count)
|
||||
|
||||
async def list_members(self, parent_team_id: str | None, page: int, page_size: int) -> Page[CustomerMember]:
|
||||
where: Final = {"parent_team_id": parent_team_id} if parent_team_id is not None else {}
|
||||
count: Final = await self.members.count(where=where)
|
||||
rows: Final = await self.members.find_many(
|
||||
where=where,
|
||||
skip=(page - 1) * page_size,
|
||||
take=page_size,
|
||||
order=[{"created_at": "asc"}, {"customer_member_id": "asc"}],
|
||||
include=_MEMBER_INCLUDE,
|
||||
)
|
||||
return Page(tuple(member_response(row, self.now) for row in rows), count)
|
||||
|
||||
async def delete_member(self, member_id: str) -> None | Failure:
|
||||
member: Final = await self.get_member(member_id)
|
||||
if isinstance(member, Failure):
|
||||
return member
|
||||
await self.members.delete(where={"customer_member_id": member_id})
|
||||
if member.budget_id is not None:
|
||||
await self.budgets.delete(where={"budget_id": member.budget_id})
|
||||
return None
|
||||
|
||||
async def delete_team(self, team_id: str) -> None | Failure:
|
||||
team: Final = await self.get_team(team_id)
|
||||
if isinstance(team, Failure):
|
||||
return team
|
||||
if await self.members.count(where={"parent_team_id": team_id}):
|
||||
return Failure(ErrorCode.HAS_CHILDREN)
|
||||
await self.teams.delete(where={"customer_team_id": team_id})
|
||||
if team.budget_id is not None:
|
||||
await self.budgets.delete(where={"budget_id": team.budget_id})
|
||||
return None
|
||||
|
||||
async def delete_organization(self, organization_id: str) -> None | Failure:
|
||||
organization: Final = await self.get_organization(organization_id)
|
||||
if isinstance(organization, Failure):
|
||||
return organization
|
||||
if await self.teams.count(where={"parent_org_id": organization_id}):
|
||||
return Failure(ErrorCode.HAS_CHILDREN)
|
||||
await self.organizations.delete(where={"customer_org_id": organization_id})
|
||||
for template in organization.default_budgets:
|
||||
await self.budgets.delete(where={"budget_id": template.budget_id})
|
||||
if organization.budget_id is not None:
|
||||
await self.budgets.delete(where={"budget_id": organization.budget_id})
|
||||
return None
|
||||
|
||||
|
||||
class CustomerHierarchy:
|
||||
def __init__(self, db: TransactionSource, clock: Callable[[], datetime] = utc_now) -> None:
|
||||
self.db = db
|
||||
self.clock = clock
|
||||
|
||||
async def transaction(
|
||||
self,
|
||||
operation: Callable[[CustomerHierarchyStore], Awaitable[T]],
|
||||
*,
|
||||
operation_type: TransactionOperation = TransactionOperation.READ_WRITE,
|
||||
) -> T | Failure:
|
||||
from prisma.errors import ForeignKeyViolationError, RecordNotFoundError, UniqueViolationError
|
||||
|
||||
try:
|
||||
async with self.db.tx() as tx:
|
||||
return await operation(CustomerHierarchyStore(tx, self.clock()))
|
||||
except UniqueViolationError:
|
||||
return Failure(ErrorCode.ALREADY_EXISTS)
|
||||
except RecordNotFoundError:
|
||||
return Failure(ErrorCode.NOT_FOUND)
|
||||
except ForeignKeyViolationError:
|
||||
return Failure(
|
||||
ErrorCode.HAS_CHILDREN if operation_type is TransactionOperation.DELETE else ErrorCode.PARENT_NOT_FOUND
|
||||
)
|
||||
|
||||
async def create_organization(self, payload: OrganizationCreate, actor: str) -> Organization | Failure:
|
||||
return await self.transaction(lambda store: store.create_organization(payload, actor))
|
||||
|
||||
async def get_organization(self, entity_id: str) -> Organization | Failure:
|
||||
return await self.transaction(lambda store: store.get_organization(entity_id))
|
||||
|
||||
async def update_organization(
|
||||
self, entity_id: str, payload: OrganizationPatch, actor: str
|
||||
) -> Organization | Failure:
|
||||
return await self.transaction(lambda store: store.update_organization(entity_id, payload, actor))
|
||||
|
||||
async def delete_organization(self, entity_id: str) -> None | Failure:
|
||||
return await self.transaction(
|
||||
lambda store: store.delete_organization(entity_id), operation_type=TransactionOperation.DELETE
|
||||
)
|
||||
|
||||
async def create_team(self, payload: TeamCreate, actor: str) -> Team | Failure:
|
||||
return await self.transaction(lambda store: store.create_team(payload, actor))
|
||||
|
||||
async def get_team(self, entity_id: str) -> Team | Failure:
|
||||
return await self.transaction(lambda store: store.get_team(entity_id))
|
||||
|
||||
async def update_team(self, entity_id: str, payload: TeamPatch, actor: str) -> Team | Failure:
|
||||
return await self.transaction(lambda store: store.update_team(entity_id, payload, actor))
|
||||
|
||||
async def delete_team(self, entity_id: str) -> None | Failure:
|
||||
return await self.transaction(
|
||||
lambda store: store.delete_team(entity_id), operation_type=TransactionOperation.DELETE
|
||||
)
|
||||
|
||||
async def create_member(self, payload: MemberCreate, actor: str) -> CustomerMember | Failure:
|
||||
return await self.transaction(lambda store: store.create_member(payload, actor))
|
||||
|
||||
async def get_member(self, entity_id: str) -> CustomerMember | Failure:
|
||||
return await self.transaction(lambda store: store.get_member(entity_id))
|
||||
|
||||
async def update_member(self, entity_id: str, payload: MemberPatch, actor: str) -> CustomerMember | Failure:
|
||||
return await self.transaction(lambda store: store.update_member(entity_id, payload, actor))
|
||||
|
||||
async def delete_member(self, entity_id: str) -> None | Failure:
|
||||
return await self.transaction(
|
||||
lambda store: store.delete_member(entity_id), operation_type=TransactionOperation.DELETE
|
||||
)
|
||||
|
||||
async def list_organizations(self, page: int = 1, page_size: int = 50) -> Page[Organization] | Failure:
|
||||
return await self.transaction(lambda store: store.list_organizations(page, page_size))
|
||||
|
||||
async def list_teams(
|
||||
self, parent_org_id: str | None = None, page: int = 1, page_size: int = 50
|
||||
) -> Page[Team] | Failure:
|
||||
return await self.transaction(lambda store: store.list_teams(parent_org_id, page, page_size))
|
||||
|
||||
async def list_members(
|
||||
self, parent_team_id: str | None = None, page: int = 1, page_size: int = 50
|
||||
) -> Page[CustomerMember] | Failure:
|
||||
return await self.transaction(lambda store: store.list_members(parent_team_id, page, page_size))
|
||||
|
||||
async def add_template(self, organization_id: str, payload: BudgetTemplate, actor: str) -> Organization | Failure:
|
||||
return await self.transaction(lambda store: store.add_template(organization_id, payload, actor))
|
||||
|
||||
async def update_template(
|
||||
self, organization_id: str, budget_id: str, payload: TemplatePatch, actor: str
|
||||
) -> Organization | Failure:
|
||||
return await self.transaction(lambda store: store.update_template(organization_id, budget_id, payload, actor))
|
||||
|
||||
async def delete_template(self, organization_id: str, budget_id: str, actor: str) -> None | Failure:
|
||||
return await self.transaction(
|
||||
lambda store: store.delete_template(organization_id, budget_id, actor),
|
||||
operation_type=TransactionOperation.DELETE,
|
||||
)
|
||||
|
|
@ -11,6 +11,11 @@ generator client {
|
|||
|
||||
// Budget / Rate Limits for an org
|
||||
model LiteLLM_BudgetTable {
|
||||
customer_organization_defaults LiteLLM_CustomerOrganizationTable[] @relation("CustomerOrganizationDefaultBudget")
|
||||
customer_organizations LiteLLM_CustomerOrganizationTable[] @relation("CustomerOrganizationBudget")
|
||||
customer_member_templates LiteLLM_CustomerMemberTable[] @relation("CustomerMemberTemplate")
|
||||
customer_teams LiteLLM_CustomerTeamTable[] @relation("CustomerTeamBudget")
|
||||
customer_members LiteLLM_CustomerMemberTable[] @relation("CustomerMemberBudget")
|
||||
budget_id String @id @default(uuid())
|
||||
max_budget Float?
|
||||
soft_budget Float?
|
||||
|
|
@ -40,6 +45,61 @@ model LiteLLM_BudgetTable {
|
|||
organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization
|
||||
}
|
||||
|
||||
model LiteLLM_CustomerOrganizationTable {
|
||||
spend Float @default(0.0)
|
||||
budget_reset_at DateTime?
|
||||
budget_ids String[] @default([])
|
||||
budget_names Json @default("{}")
|
||||
budget_id String?
|
||||
budget LiteLLM_BudgetTable? @relation("CustomerOrganizationBudget", fields: [budget_id], references: [budget_id])
|
||||
customer_org_id String @id @default(uuid())
|
||||
alias String?
|
||||
default_budget_id String?
|
||||
default_budget LiteLLM_BudgetTable? @relation("CustomerOrganizationDefaultBudget", fields: [default_budget_id], references: [budget_id])
|
||||
teams LiteLLM_CustomerTeamTable[]
|
||||
created_at DateTime @default(now())
|
||||
created_by String
|
||||
updated_at DateTime @default(now()) @updatedAt
|
||||
updated_by String
|
||||
}
|
||||
|
||||
model LiteLLM_CustomerTeamTable {
|
||||
spend Float @default(0.0)
|
||||
budget_reset_at DateTime?
|
||||
customer_team_id String @id @default(uuid())
|
||||
alias String?
|
||||
parent_org_id String
|
||||
parent_org LiteLLM_CustomerOrganizationTable @relation(fields: [parent_org_id], references: [customer_org_id])
|
||||
budget_id String?
|
||||
budget LiteLLM_BudgetTable? @relation("CustomerTeamBudget", fields: [budget_id], references: [budget_id])
|
||||
members LiteLLM_CustomerMemberTable[]
|
||||
created_at DateTime @default(now())
|
||||
created_by String
|
||||
updated_at DateTime @default(now()) @updatedAt
|
||||
updated_by String
|
||||
@@index([parent_org_id])
|
||||
}
|
||||
|
||||
model LiteLLM_CustomerMemberTable {
|
||||
spend Float @default(0.0)
|
||||
budget_reset_at DateTime?
|
||||
customer_member_id String @id @default(uuid())
|
||||
alias String?
|
||||
parent_team_id String
|
||||
parent_team LiteLLM_CustomerTeamTable @relation(fields: [parent_team_id], references: [customer_team_id])
|
||||
budget_id String?
|
||||
budget LiteLLM_BudgetTable? @relation("CustomerMemberBudget", fields: [budget_id], references: [budget_id], onDelete: Restrict)
|
||||
budget_template_id String?
|
||||
budget_template LiteLLM_BudgetTable? @relation("CustomerMemberTemplate", fields: [budget_template_id], references: [budget_id], onDelete: Restrict)
|
||||
budget_override_expires_at DateTime?
|
||||
created_at DateTime @default(now())
|
||||
created_by String
|
||||
updated_at DateTime @default(now()) @updatedAt
|
||||
updated_by String
|
||||
@@index([parent_team_id])
|
||||
@@index([budget_template_id])
|
||||
}
|
||||
|
||||
// Models on proxy
|
||||
model LiteLLM_CredentialsTable {
|
||||
credential_id String @id @default(uuid())
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ Budget repository for database operations on LiteLLM_BudgetTable.
|
|||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Final
|
||||
from typing import TYPE_CHECKING, Final, TypeVar
|
||||
|
||||
from litellm.models.budget import LiteLLM_BudgetTable
|
||||
from litellm.repositories.base_repository import BaseRepository
|
||||
|
|
@ -14,6 +14,19 @@ if TYPE_CHECKING:
|
|||
from prisma import models as prisma_models
|
||||
|
||||
|
||||
_Row: Final = TypeVar("_Row")
|
||||
|
||||
|
||||
async def create_budget_row(table: TableActions[_Row], data: Mapping[str, object], actor: str) -> _Row:
|
||||
return await table.create(data={**data, "created_by": actor, "updated_by": actor})
|
||||
|
||||
|
||||
async def update_budget_row(
|
||||
table: TableActions[_Row], budget_id: str, data: Mapping[str, object], actor: str
|
||||
) -> _Row | None:
|
||||
return await table.update(where={"budget_id": budget_id}, data={**data, "updated_by": actor})
|
||||
|
||||
|
||||
class _BudgetTable(PrismaTableRepository["prisma_models.LiteLLM_BudgetTable"]):
|
||||
table_name = "litellm_budgettable"
|
||||
|
||||
|
|
@ -66,7 +79,10 @@ class BudgetRepository(BaseRepository[LiteLLM_BudgetTable]):
|
|||
if allowed_models is not None:
|
||||
data["allowed_models"] = allowed_models
|
||||
|
||||
return await self.create(data)
|
||||
record: Final = await create_budget_row(self.table, data, created_by)
|
||||
model: Final = self._to_model(record)
|
||||
assert model is not None
|
||||
return model
|
||||
|
||||
async def update_budget(
|
||||
self,
|
||||
|
|
@ -100,7 +116,8 @@ class BudgetRepository(BaseRepository[LiteLLM_BudgetTable]):
|
|||
if allowed_models is not None:
|
||||
data["allowed_models"] = allowed_models
|
||||
|
||||
return await self.update(budget_id, data, id_field="budget_id")
|
||||
record: Final = await update_budget_row(self.table, budget_id, data, updated_by)
|
||||
return self._to_model(record)
|
||||
|
||||
async def delete_budget(self, budget_id: str) -> LiteLLM_BudgetTable | None:
|
||||
"""Delete a budget record."""
|
||||
|
|
|
|||
|
|
@ -77,6 +77,10 @@ class PrismaTableRepository(Generic[RowT_co]):
|
|||
return wrap_table_actions_for_config_sync(actions=actions, table_name=self.table_name)
|
||||
|
||||
|
||||
class CustomerOrganizationRepository(PrismaTableRepository["prisma_models.LiteLLM_CustomerOrganizationTable"]):
|
||||
table_name = "litellm_customerorganizationtable"
|
||||
|
||||
|
||||
class PolicyRepository(PrismaTableRepository["prisma_models.LiteLLM_PolicyTable"]):
|
||||
table_name = "litellm_policytable"
|
||||
|
||||
|
|
|
|||
|
|
@ -11,6 +11,11 @@ generator client {
|
|||
|
||||
// Budget / Rate Limits for an org
|
||||
model LiteLLM_BudgetTable {
|
||||
customer_organization_defaults LiteLLM_CustomerOrganizationTable[] @relation("CustomerOrganizationDefaultBudget")
|
||||
customer_organizations LiteLLM_CustomerOrganizationTable[] @relation("CustomerOrganizationBudget")
|
||||
customer_member_templates LiteLLM_CustomerMemberTable[] @relation("CustomerMemberTemplate")
|
||||
customer_teams LiteLLM_CustomerTeamTable[] @relation("CustomerTeamBudget")
|
||||
customer_members LiteLLM_CustomerMemberTable[] @relation("CustomerMemberBudget")
|
||||
budget_id String @id @default(uuid())
|
||||
max_budget Float?
|
||||
soft_budget Float?
|
||||
|
|
@ -40,6 +45,61 @@ model LiteLLM_BudgetTable {
|
|||
organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization
|
||||
}
|
||||
|
||||
model LiteLLM_CustomerOrganizationTable {
|
||||
spend Float @default(0.0)
|
||||
budget_reset_at DateTime?
|
||||
budget_ids String[] @default([])
|
||||
budget_names Json @default("{}")
|
||||
budget_id String?
|
||||
budget LiteLLM_BudgetTable? @relation("CustomerOrganizationBudget", fields: [budget_id], references: [budget_id])
|
||||
customer_org_id String @id @default(uuid())
|
||||
alias String?
|
||||
default_budget_id String?
|
||||
default_budget LiteLLM_BudgetTable? @relation("CustomerOrganizationDefaultBudget", fields: [default_budget_id], references: [budget_id])
|
||||
teams LiteLLM_CustomerTeamTable[]
|
||||
created_at DateTime @default(now())
|
||||
created_by String
|
||||
updated_at DateTime @default(now()) @updatedAt
|
||||
updated_by String
|
||||
}
|
||||
|
||||
model LiteLLM_CustomerTeamTable {
|
||||
spend Float @default(0.0)
|
||||
budget_reset_at DateTime?
|
||||
customer_team_id String @id @default(uuid())
|
||||
alias String?
|
||||
parent_org_id String
|
||||
parent_org LiteLLM_CustomerOrganizationTable @relation(fields: [parent_org_id], references: [customer_org_id])
|
||||
budget_id String?
|
||||
budget LiteLLM_BudgetTable? @relation("CustomerTeamBudget", fields: [budget_id], references: [budget_id])
|
||||
members LiteLLM_CustomerMemberTable[]
|
||||
created_at DateTime @default(now())
|
||||
created_by String
|
||||
updated_at DateTime @default(now()) @updatedAt
|
||||
updated_by String
|
||||
@@index([parent_org_id])
|
||||
}
|
||||
|
||||
model LiteLLM_CustomerMemberTable {
|
||||
spend Float @default(0.0)
|
||||
budget_reset_at DateTime?
|
||||
customer_member_id String @id @default(uuid())
|
||||
alias String?
|
||||
parent_team_id String
|
||||
parent_team LiteLLM_CustomerTeamTable @relation(fields: [parent_team_id], references: [customer_team_id])
|
||||
budget_id String?
|
||||
budget LiteLLM_BudgetTable? @relation("CustomerMemberBudget", fields: [budget_id], references: [budget_id], onDelete: Restrict)
|
||||
budget_template_id String?
|
||||
budget_template LiteLLM_BudgetTable? @relation("CustomerMemberTemplate", fields: [budget_template_id], references: [budget_id], onDelete: Restrict)
|
||||
budget_override_expires_at DateTime?
|
||||
created_at DateTime @default(now())
|
||||
created_by String
|
||||
updated_at DateTime @default(now()) @updatedAt
|
||||
updated_by String
|
||||
@@index([parent_team_id])
|
||||
@@index([budget_template_id])
|
||||
}
|
||||
|
||||
// Models on proxy
|
||||
model LiteLLM_CredentialsTable {
|
||||
credential_id String @id @default(uuid())
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ count them per rule the same delta-vs-base way as the other lint gates. The rule
|
|||
OA001 no request body schema (the handler reads ``request.body()`` and never told FastAPI
|
||||
what is in it, see ``inline_request_body`` in custom_openapi_spec.py for the fix)
|
||||
OA002 no description (the handler has no docstring and the route no ``description=``)
|
||||
OA003 the 200 response has no schema (no ``response_model`` and no ``responses={200: ...}``)
|
||||
OA003 neither a 200 nor a 201 response has a schema
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -45,7 +45,8 @@ def _operation_codes(operation: Mapping[str, object]) -> Iterator[str]:
|
|||
description: Final = operation.get("description")
|
||||
if not isinstance(description, str) or not description.strip():
|
||||
yield "OA002"
|
||||
if not _content_schema(_as_mapping(operation.get("responses")).get("200")):
|
||||
responses: Final = _as_mapping(operation.get("responses"))
|
||||
if not any(_content_schema(responses.get(status)) for status in ("200", "201")):
|
||||
yield "OA003"
|
||||
|
||||
|
||||
|
|
|
|||
114
litellm/types/proxy/management_endpoints/customer_hierarchy.py
Normal file
114
litellm/types/proxy/management_endpoints/customer_hierarchy.py
Normal file
|
|
@ -0,0 +1,114 @@
|
|||
from datetime import datetime
|
||||
from typing import Final
|
||||
|
||||
from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, field_validator
|
||||
|
||||
from litellm.proxy.common_utils.timezone_utils import budget_duration_error
|
||||
|
||||
|
||||
class Budget(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
max_budget: float | None = Field(default=None, ge=0, allow_inf_nan=False)
|
||||
soft_budget: float | None = Field(default=None, ge=0, allow_inf_nan=False)
|
||||
tpm_limit: int | None = Field(default=None, ge=0)
|
||||
rpm_limit: int | None = Field(default=None, ge=0)
|
||||
allowed_models: tuple[str, ...] | None = None
|
||||
budget_duration: str | None = None
|
||||
|
||||
@field_validator("budget_duration")
|
||||
@classmethod
|
||||
def valid_budget_duration(cls, value: str | None) -> str | None:
|
||||
error: Final = budget_duration_error(value)
|
||||
if error is not None:
|
||||
raise ValueError(error)
|
||||
return value
|
||||
|
||||
|
||||
class BudgetTemplate(Budget):
|
||||
name: str = Field(min_length=1)
|
||||
is_default: bool = False
|
||||
|
||||
|
||||
class TemplatePatch(Budget):
|
||||
name: str | None = Field(default=None, min_length=1)
|
||||
is_default: bool | None = None
|
||||
|
||||
|
||||
class OrganizationBudget(BudgetTemplate):
|
||||
budget_id: str
|
||||
|
||||
|
||||
class OrganizationPatch(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
alias: str | None = None
|
||||
budget: Budget | None = None
|
||||
|
||||
|
||||
class OrganizationCreate(OrganizationPatch):
|
||||
customer_org_id: str | None = None
|
||||
default_budgets: tuple[BudgetTemplate, ...] = ()
|
||||
|
||||
|
||||
class TeamPatch(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
alias: str | None = None
|
||||
budget: Budget | None = None
|
||||
|
||||
|
||||
class TeamCreate(TeamPatch):
|
||||
customer_team_id: str | None = None
|
||||
parent_org_id: str
|
||||
|
||||
|
||||
class MemberPatch(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
alias: str | None = None
|
||||
budget_override: Budget | None = None
|
||||
budget_template_id: str | None = None
|
||||
budget_override_expires_at: AwareDatetime | None = None
|
||||
|
||||
|
||||
class MemberCreate(MemberPatch):
|
||||
customer_member_id: str | None = None
|
||||
parent_team_id: str
|
||||
|
||||
|
||||
class AuditFields(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
alias: str | None = None
|
||||
spend: float = 0.0
|
||||
budget_reset_at: datetime | None = None
|
||||
created_at: datetime
|
||||
created_by: str
|
||||
updated_at: datetime
|
||||
updated_by: str
|
||||
|
||||
|
||||
class Organization(AuditFields):
|
||||
budget_id: str | None = None
|
||||
budget: Budget | None = None
|
||||
customer_org_id: str
|
||||
default_budget_id: str | None = None
|
||||
default_budgets: tuple[OrganizationBudget, ...] = ()
|
||||
|
||||
|
||||
class Team(AuditFields):
|
||||
customer_team_id: str
|
||||
parent_org_id: str
|
||||
budget_id: str | None = None
|
||||
budget: Budget | None = None
|
||||
|
||||
|
||||
class CustomerMember(AuditFields):
|
||||
customer_member_id: str
|
||||
parent_team_id: str
|
||||
budget_id: str | None = None
|
||||
budget_template_id: str | None = None
|
||||
budget_override: Budget | None = None
|
||||
budget_override_expires_at: datetime | None = None
|
||||
effective_budget: Budget
|
||||
1052
litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
1052
litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
File diff suppressed because it is too large
Load diff
|
|
@ -392,6 +392,10 @@ def test_virtual_key_allowed_routes_with_litellm_routes_member_name_allowed(grou
|
|||
[
|
||||
("/model/update", "POST"),
|
||||
("/model/123e4567-e89b-12d3-a456-426614174000/update", "PATCH"),
|
||||
("/management/v1/customer_organizations", "POST"),
|
||||
("/management/v1/customer_organizations/org-1/default_budgets/budget-1", "DELETE"),
|
||||
("/management/v1/customer_teams/team-1", "PATCH"),
|
||||
("/management/v1/customer_members/member-1", "GET"),
|
||||
],
|
||||
)
|
||||
def test_virtual_key_management_routes_allows_model_update(route: str, method: str) -> None:
|
||||
|
|
|
|||
|
|
@ -0,0 +1,541 @@
|
|||
from typing import TYPE_CHECKING, Final, cast
|
||||
|
||||
import pytest
|
||||
from fastapi import FastAPI, Request
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.list_api.common import ManagementProblem, problem_response
|
||||
from litellm.proxy.management_endpoints.management_v1 import router
|
||||
from litellm.proxy.management_endpoints.management_v1.customers import get_customer_hierarchy
|
||||
from litellm.proxy.management_helpers.customer_hierarchy import CustomerHierarchy
|
||||
from tests.unit.proxy.management_helpers.test_customer_hierarchy import NOW, Database
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from prisma import Prisma
|
||||
|
||||
|
||||
def client(role: LitellmUserRoles = LitellmUserRoles.PROXY_ADMIN) -> TestClient:
|
||||
app: Final = FastAPI()
|
||||
app.include_router(router)
|
||||
|
||||
@app.exception_handler(ManagementProblem)
|
||||
async def problem(request: Request, exc: ManagementProblem):
|
||||
return problem_response(exc.problem)
|
||||
|
||||
auth: Final = UserAPIKeyAuth(user_id="admin", user_role=role)
|
||||
hierarchy: Final = CustomerHierarchy(
|
||||
cast("Prisma", Database()), clock=lambda: NOW
|
||||
) # cast-ok: in-memory Prisma edge
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: auth
|
||||
app.dependency_overrides[get_customer_hierarchy] = lambda: hierarchy
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
def family(http: TestClient) -> dict[str, object]:
|
||||
org: Final = http.post(
|
||||
"/management/v1/customer_organizations",
|
||||
json={
|
||||
"customer_org_id": "org",
|
||||
"budget": {"max_budget": 1000},
|
||||
"default_budgets": [
|
||||
{
|
||||
"name": "standard",
|
||||
"max_budget": 100,
|
||||
"rpm_limit": 10,
|
||||
"allowed_models": ["model-a"],
|
||||
"is_default": True,
|
||||
},
|
||||
{"name": "small", "max_budget": 20, "rpm_limit": 8},
|
||||
],
|
||||
},
|
||||
)
|
||||
assert org.status_code == 201, org.text
|
||||
team: Final = http.post(
|
||||
"/management/v1/customer_teams",
|
||||
json={
|
||||
"customer_team_id": "team",
|
||||
"parent_org_id": "org",
|
||||
"budget": {"max_budget": 999},
|
||||
},
|
||||
)
|
||||
assert team.status_code == 201, team.text
|
||||
return org.json()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("role", (LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY))
|
||||
@pytest.mark.parametrize(
|
||||
"method,path,body",
|
||||
(
|
||||
("POST", "/customer_organizations", {}),
|
||||
("GET", "/customer_organizations", None),
|
||||
("GET", "/customer_organizations/id", None),
|
||||
("PATCH", "/customer_organizations/id", {}),
|
||||
("DELETE", "/customer_organizations/id", None),
|
||||
("POST", "/customer_teams", {"parent_org_id": "id"}),
|
||||
("GET", "/customer_teams", None),
|
||||
("GET", "/customer_teams/id", None),
|
||||
("PATCH", "/customer_teams/id", {}),
|
||||
("DELETE", "/customer_teams/id", None),
|
||||
("POST", "/customer_members", {"parent_team_id": "id"}),
|
||||
("GET", "/customer_members", None),
|
||||
("GET", "/customer_members/id", None),
|
||||
("PATCH", "/customer_members/id", {}),
|
||||
("DELETE", "/customer_members/id", None),
|
||||
("POST", "/customer_organizations/id/default_budgets", {"name": "new"}),
|
||||
("PATCH", "/customer_organizations/id/default_budgets/budget", {}),
|
||||
("DELETE", "/customer_organizations/id/default_budgets/budget", None),
|
||||
),
|
||||
)
|
||||
def test_every_dispatched_method_is_admin_only(
|
||||
role: LitellmUserRoles,
|
||||
method: str,
|
||||
path: str,
|
||||
body: dict[str, object] | None,
|
||||
) -> None:
|
||||
response: Final = client(role).request(method, "/management/v1" + path, json=body)
|
||||
assert response.status_code == 403, response.text
|
||||
assert response.json()["status"] == 403
|
||||
|
||||
|
||||
def test_template_edits_propagate_and_only_live_override_fields_win() -> None:
|
||||
http: Final = client()
|
||||
org: Final = family(http)
|
||||
templates: Final = org["default_budgets"]
|
||||
default_id: Final = templates[0]["budget_id"]
|
||||
other_id: Final = templates[1]["budget_id"]
|
||||
default: Final = http.post(
|
||||
"/management/v1/customer_members",
|
||||
json={
|
||||
"customer_member_id": "default",
|
||||
"parent_team_id": "team",
|
||||
},
|
||||
)
|
||||
assert default.status_code == 201, default.text
|
||||
assert default.json()["budget_id"] is None
|
||||
assert default.json()["budget_template_id"] == default_id
|
||||
override: Final = http.post(
|
||||
"/management/v1/customer_members",
|
||||
json={
|
||||
"customer_member_id": "override",
|
||||
"parent_team_id": "team",
|
||||
"budget_template_id": other_id,
|
||||
"budget_override": {"max_budget": 3, "allowed_models": []},
|
||||
},
|
||||
)
|
||||
assert override.status_code == 201, override.text
|
||||
expired: Final = http.post(
|
||||
"/management/v1/customer_members",
|
||||
json={
|
||||
"customer_member_id": "expired",
|
||||
"parent_team_id": "team",
|
||||
"budget_override": {"max_budget": 1},
|
||||
"budget_override_expires_at": "2025-01-01T00:00:00Z",
|
||||
},
|
||||
)
|
||||
assert expired.status_code == 201, expired.text
|
||||
first: Final = http.patch(
|
||||
f"/management/v1/customer_organizations/org/default_budgets/{default_id}",
|
||||
json={
|
||||
"max_budget": 200,
|
||||
"rpm_limit": 15,
|
||||
"name": "renamed",
|
||||
},
|
||||
)
|
||||
assert first.status_code == 200, first.text
|
||||
second: Final = http.patch(
|
||||
f"/management/v1/customer_organizations/org/default_budgets/{other_id}",
|
||||
json={
|
||||
"max_budget": 40,
|
||||
"rpm_limit": 12,
|
||||
"allowed_models": ["model-b"],
|
||||
},
|
||||
)
|
||||
assert second.status_code == 200, second.text
|
||||
fetched: Final = http.get("/management/v1/customer_members/default")
|
||||
assert fetched.status_code == 200, fetched.text
|
||||
assert fetched.json()["effective_budget"]["max_budget"] == 200
|
||||
assert fetched.json()["effective_budget"]["rpm_limit"] == 15
|
||||
selected: Final = http.get("/management/v1/customer_members/override")
|
||||
assert selected.status_code == 200, selected.text
|
||||
assert selected.json()["effective_budget"] == {
|
||||
"max_budget": 3,
|
||||
"soft_budget": None,
|
||||
"rpm_limit": 12,
|
||||
"tpm_limit": None,
|
||||
"allowed_models": ["model-b"],
|
||||
"budget_duration": None,
|
||||
}
|
||||
fallback: Final = http.get("/management/v1/customer_members/expired")
|
||||
assert fallback.status_code == 200, fallback.text
|
||||
assert fallback.json()["effective_budget"] == fetched.json()["effective_budget"]
|
||||
assert fallback.json()["budget_override"]["max_budget"] == 1
|
||||
assert [template["name"] for template in second.json()["default_budgets"]] == ["renamed", "small"]
|
||||
assert second.json()["budget"]["max_budget"] == 1000
|
||||
|
||||
|
||||
def test_override_can_be_reexpired_cleared_and_template_switched() -> None:
|
||||
http: Final = client()
|
||||
org: Final = family(http)
|
||||
created: Final = http.post(
|
||||
"/management/v1/customer_members",
|
||||
json={
|
||||
"customer_member_id": "member",
|
||||
"parent_team_id": "team",
|
||||
"budget_override": {"max_budget": 5},
|
||||
},
|
||||
)
|
||||
assert created.status_code == 201, created.text
|
||||
expired: Final = http.patch(
|
||||
"/management/v1/customer_members/member",
|
||||
json={
|
||||
"budget_override_expires_at": "2025-01-01T00:00:00Z",
|
||||
},
|
||||
)
|
||||
assert expired.status_code == 200, expired.text
|
||||
assert expired.json()["effective_budget"]["max_budget"] == 100
|
||||
renewed: Final = http.patch("/management/v1/customer_members/member", json={"budget_override_expires_at": None})
|
||||
assert renewed.status_code == 200, renewed.text
|
||||
assert renewed.json()["effective_budget"]["max_budget"] == 5
|
||||
switched: Final = http.patch(
|
||||
"/management/v1/customer_members/member",
|
||||
json={
|
||||
"budget_template_id": org["default_budgets"][1]["budget_id"],
|
||||
"budget_override": None,
|
||||
},
|
||||
)
|
||||
assert switched.status_code == 200, switched.text
|
||||
assert switched.json()["budget_id"] is None
|
||||
assert switched.json()["budget_override"] is None
|
||||
assert switched.json()["effective_budget"]["max_budget"] == 20
|
||||
cleared: Final = http.patch("/management/v1/customer_members/member", json={"budget_template_id": None})
|
||||
assert cleared.status_code == 422, cleared.text
|
||||
after: Final = http.get("/management/v1/customer_members/member")
|
||||
assert after.json() == switched.json()
|
||||
|
||||
|
||||
def test_clearing_an_override_also_clears_its_expiry() -> None:
|
||||
http: Final = client()
|
||||
family(http)
|
||||
created: Final = http.post(
|
||||
"/management/v1/customer_members",
|
||||
json={
|
||||
"customer_member_id": "member",
|
||||
"parent_team_id": "team",
|
||||
"budget_override": {"max_budget": 5},
|
||||
"budget_override_expires_at": "2025-01-01T00:00:00Z",
|
||||
},
|
||||
)
|
||||
assert created.status_code == 201, created.text
|
||||
assert created.json()["effective_budget"]["max_budget"] == 100
|
||||
cleared: Final = http.patch("/management/v1/customer_members/member", json={"budget_override": None})
|
||||
assert cleared.status_code == 200, cleared.text
|
||||
assert cleared.json()["budget_override"] is None
|
||||
assert cleared.json()["budget_override_expires_at"] is None
|
||||
renewed: Final = http.patch("/management/v1/customer_members/member", json={"budget_override": {"max_budget": 7}})
|
||||
assert renewed.status_code == 200, renewed.text
|
||||
assert renewed.json()["budget_override_expires_at"] is None
|
||||
assert renewed.json()["effective_budget"]["max_budget"] == 7
|
||||
|
||||
|
||||
def test_null_budget_clears_the_overall_budget_of_an_organization_and_a_team() -> None:
|
||||
http: Final = client()
|
||||
family(http)
|
||||
org_cleared: Final = http.patch("/management/v1/customer_organizations/org", json={"budget": None})
|
||||
assert org_cleared.status_code == 200, org_cleared.text
|
||||
assert org_cleared.json()["budget"] is None
|
||||
assert org_cleared.json()["budget_id"] is None
|
||||
team_cleared: Final = http.patch("/management/v1/customer_teams/team", json={"budget": None})
|
||||
assert team_cleared.status_code == 200, team_cleared.text
|
||||
assert team_cleared.json()["budget"] is None
|
||||
assert team_cleared.json()["budget_id"] is None
|
||||
omitted: Final = http.patch("/management/v1/customer_organizations/org", json={"alias": "renamed"})
|
||||
assert omitted.status_code == 200, omitted.text
|
||||
assert omitted.json()["budget"] is None
|
||||
restored: Final = http.patch("/management/v1/customer_organizations/org", json={"budget": {"max_budget": 5}})
|
||||
assert restored.status_code == 200, restored.text
|
||||
assert restored.json()["budget"]["max_budget"] == 5
|
||||
|
||||
|
||||
def test_template_deletion_guards_names_and_references() -> None:
|
||||
http: Final = client()
|
||||
org: Final = family(http)
|
||||
default_id: Final = org["default_budget_id"]
|
||||
selected_id: Final = org["default_budgets"][1]["budget_id"]
|
||||
member: Final = http.post(
|
||||
"/management/v1/customer_members",
|
||||
json={
|
||||
"customer_member_id": "member",
|
||||
"parent_team_id": "team",
|
||||
"budget_template_id": selected_id,
|
||||
},
|
||||
)
|
||||
assert member.status_code == 201, member.text
|
||||
for template_id in (default_id, selected_id):
|
||||
response: Final = http.delete(f"/management/v1/customer_organizations/org/default_budgets/{template_id}")
|
||||
assert response.status_code == 409, response.text
|
||||
assert response.json()["type"].endswith("template-in-use")
|
||||
duplicate: Final = http.post("/management/v1/customer_organizations/org/default_budgets", json={"name": "small"})
|
||||
assert duplicate.status_code == 422, duplicate.text
|
||||
rename: Final = http.patch(
|
||||
f"/management/v1/customer_organizations/org/default_budgets/{selected_id}", json={"name": "standard"}
|
||||
)
|
||||
assert rename.status_code == 422, rename.text
|
||||
removed: Final = http.delete("/management/v1/customer_members/member")
|
||||
assert removed.status_code == 204, removed.text
|
||||
fetched: Final = http.get("/management/v1/customer_organizations/org")
|
||||
assert fetched.json() == org
|
||||
deleted: Final = http.delete(f"/management/v1/customer_organizations/org/default_budgets/{selected_id}")
|
||||
assert deleted.status_code == 204, deleted.text
|
||||
added: Final = http.post(
|
||||
"/management/v1/customer_organizations/org/default_budgets", json={"name": "new", "is_default": True}
|
||||
)
|
||||
assert added.status_code == 201, added.text
|
||||
assert added.json()["default_budget_id"] != default_id
|
||||
old_default: Final = http.delete(f"/management/v1/customer_organizations/org/default_budgets/{default_id}")
|
||||
assert old_default.status_code == 204, old_default.text
|
||||
|
||||
|
||||
def test_missing_template_requires_custom_budget_and_wrong_org_is_rejected() -> None:
|
||||
http: Final = client()
|
||||
org: Final = family(http)
|
||||
empty: Final = http.post("/management/v1/customer_organizations", json={"customer_org_id": "empty"})
|
||||
assert empty.status_code == 201, empty.text
|
||||
team: Final = http.post(
|
||||
"/management/v1/customer_teams", json={"customer_team_id": "empty-team", "parent_org_id": "empty"}
|
||||
)
|
||||
assert team.status_code == 201, team.text
|
||||
missing: Final = http.post("/management/v1/customer_members", json={"parent_team_id": "empty-team"})
|
||||
assert missing.status_code == 422, missing.text
|
||||
wrong: Final = http.post(
|
||||
"/management/v1/customer_members",
|
||||
json={
|
||||
"parent_team_id": "empty-team",
|
||||
"budget_template_id": org["default_budget_id"],
|
||||
"budget_override": {"max_budget": 2},
|
||||
},
|
||||
)
|
||||
assert wrong.status_code == 422, wrong.text
|
||||
custom: Final = http.post(
|
||||
"/management/v1/customer_members", json={"parent_team_id": "empty-team", "budget_override": {"max_budget": 2}}
|
||||
)
|
||||
assert custom.status_code == 201, custom.text
|
||||
assert custom.json()["budget_template_id"] is None
|
||||
assert custom.json()["effective_budget"]["max_budget"] == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"entity,body",
|
||||
(
|
||||
("customer_organizations", {}),
|
||||
("customer_teams", {"parent_org_id": "org"}),
|
||||
("customer_members", {"parent_team_id": "team"}),
|
||||
),
|
||||
)
|
||||
@pytest.mark.parametrize("method", ("POST", "PATCH"))
|
||||
@pytest.mark.parametrize("field,value", (("spend", 1), ("budget_reset_at", "2025-01-01T00:00:00Z")))
|
||||
def test_spend_and_reset_fields_are_read_only(
|
||||
entity: str, body: dict[str, str], method: str, field: str, value: object
|
||||
) -> None:
|
||||
http: Final = client()
|
||||
family(http)
|
||||
path: Final = f"/management/v1/{entity}" + ("/id" if method == "PATCH" else "")
|
||||
response: Final = http.request(method, path, json={**body, field: value})
|
||||
assert response.status_code == 422, response.text
|
||||
|
||||
|
||||
def test_new_entities_report_zero_spend_and_no_reset_and_parent_deletion_conflicts() -> None:
|
||||
http: Final = client()
|
||||
org: Final = family(http)
|
||||
member: Final = http.post(
|
||||
"/management/v1/customer_members", json={"customer_member_id": "member", "parent_team_id": "team"}
|
||||
)
|
||||
assert member.status_code == 201, member.text
|
||||
for path in ("customer_organizations/org", "customer_teams/team", "customer_members/member"):
|
||||
fetched: Final = http.get("/management/v1/" + path)
|
||||
assert fetched.status_code == 200, fetched.text
|
||||
assert fetched.json()["spend"] == 0
|
||||
assert fetched.json()["budget_reset_at"] is None
|
||||
for path in ("customer_organizations/org", "customer_teams/team"):
|
||||
blocked: Final = http.delete("/management/v1/" + path)
|
||||
assert blocked.status_code == 409, blocked.text
|
||||
for path in ("customer_members/member", "customer_teams/team", "customer_organizations/org"):
|
||||
deleted: Final = http.delete("/management/v1/" + path)
|
||||
assert deleted.status_code == 204, deleted.text
|
||||
absent: Final = http.get("/management/v1/" + path)
|
||||
assert absent.status_code == 404, absent.text
|
||||
assert org["budget_id"] is not None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"templates",
|
||||
(
|
||||
[{"name": "a"}, {"name": "b"}],
|
||||
[{"name": "a", "is_default": True}, {"name": "b", "is_default": True}],
|
||||
[{"name": "a"}, {"name": "a", "is_default": True}],
|
||||
),
|
||||
)
|
||||
def test_invalid_template_configuration_does_not_create_an_organization(templates: list[dict[str, object]]) -> None:
|
||||
http: Final = client()
|
||||
response: Final = http.post(
|
||||
"/management/v1/customer_organizations", json={"customer_org_id": "org", "default_budgets": templates}
|
||||
)
|
||||
assert response.status_code == 422, response.text
|
||||
missing: Final = http.get("/management/v1/customer_organizations/org")
|
||||
assert missing.status_code == 404, missing.text
|
||||
|
||||
|
||||
def test_single_template_defaults_and_org_patch_cannot_replace_templates() -> None:
|
||||
http: Final = client()
|
||||
created: Final = http.post(
|
||||
"/management/v1/customer_organizations",
|
||||
json={
|
||||
"customer_org_id": "org",
|
||||
"default_budgets": [{"name": "only", "max_budget": 10}],
|
||||
},
|
||||
)
|
||||
assert created.status_code == 201, created.text
|
||||
assert created.json()["default_budget_id"] == created.json()["default_budgets"][0]["budget_id"]
|
||||
patched: Final = http.patch("/management/v1/customer_organizations/org", json={"default_budgets": []})
|
||||
assert patched.status_code == 422, patched.text
|
||||
fetched: Final = http.get("/management/v1/customer_organizations/org")
|
||||
assert fetched.json() == created.json()
|
||||
|
||||
|
||||
def test_lists_filter_parents_and_link_pages_and_reject_unknown_parameters() -> None:
|
||||
http: Final = client()
|
||||
family(http)
|
||||
other: Final = http.post("/management/v1/customer_organizations", json={"customer_org_id": "other"})
|
||||
assert other.status_code == 201, other.text
|
||||
second_team: Final = http.post(
|
||||
"/management/v1/customer_teams", json={"customer_team_id": "z-team", "parent_org_id": "org"}
|
||||
)
|
||||
assert second_team.status_code == 201, second_team.text
|
||||
listed: Final = http.get("/management/v1/customer_teams?parent_org_id=org&page_size=1")
|
||||
assert listed.status_code == 200, listed.text
|
||||
assert listed.json()["meta"]["total_count"] == 2
|
||||
assert listed.json()["data"][0]["customer_team_id"] == "team"
|
||||
next_page: Final = http.get(listed.json()["links"]["next"])
|
||||
assert next_page.status_code == 200, next_page.text
|
||||
assert next_page.json()["data"][0]["customer_team_id"] == "z-team"
|
||||
empty: Final = http.get("/management/v1/customer_members?parent_team_id=team")
|
||||
assert empty.status_code == 200, empty.text
|
||||
assert empty.json()["data"] == []
|
||||
bad: Final = http.get("/management/v1/customer_teams?unknown=1")
|
||||
assert bad.status_code == 400, bad.text
|
||||
repeated: Final = http.get("/management/v1/customer_teams?page=1&page=2")
|
||||
assert repeated.status_code == 400, repeated.text
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"entity,payload",
|
||||
(
|
||||
("customer_organizations", {"customer_org_id": "duplicate"}),
|
||||
("customer_teams", {"customer_team_id": "duplicate", "parent_org_id": "org"}),
|
||||
("customer_members", {"customer_member_id": "duplicate", "parent_team_id": "team"}),
|
||||
),
|
||||
)
|
||||
def test_duplicate_id_conflicts(entity: str, payload: dict[str, str]) -> None:
|
||||
http: Final = client()
|
||||
family(http)
|
||||
first: Final = http.post("/management/v1/" + entity, json=payload)
|
||||
assert first.status_code == 201, first.text
|
||||
duplicate: Final = http.post("/management/v1/" + entity, json=payload)
|
||||
assert duplicate.status_code == 409, duplicate.text
|
||||
|
||||
|
||||
@pytest.mark.parametrize("entity", ("customer_teams", "customer_members"))
|
||||
def test_missing_parent_is_404(entity: str) -> None:
|
||||
field: Final = "parent_org_id" if entity == "customer_teams" else "parent_team_id"
|
||||
response: Final = client().post("/management/v1/" + entity, json={field: "missing"})
|
||||
assert response.status_code == 404, response.text
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"budget",
|
||||
(
|
||||
{"max_budget": -1},
|
||||
{"max_budget": "NaN"},
|
||||
{"max_budget": "Infinity"},
|
||||
{"tpm_limit": -1},
|
||||
{"rpm_limit": -1},
|
||||
{"budget_duration": "0s"},
|
||||
),
|
||||
)
|
||||
def test_invalid_budget_values_fail_at_the_http_boundary(budget: dict[str, object]) -> None:
|
||||
response: Final = client().post("/management/v1/customer_organizations", json={"budget": budget})
|
||||
assert response.status_code == 422, response.text
|
||||
|
||||
|
||||
def test_partial_member_override_updates_preserve_other_fields_and_never_change_the_template() -> None:
|
||||
http: Final = client()
|
||||
organization: Final = family(http)
|
||||
created: Final = http.post(
|
||||
"/management/v1/customer_members",
|
||||
json={
|
||||
"customer_member_id": "member",
|
||||
"parent_team_id": "team",
|
||||
},
|
||||
)
|
||||
assert created.status_code == 201, created.text
|
||||
patched: Final = http.patch(
|
||||
"/management/v1/customer_members/member",
|
||||
json={
|
||||
"budget_override": {"max_budget": 5, "soft_budget": 4, "tpm_limit": 2, "allowed_models": ["model-b"]},
|
||||
"budget_override_expires_at": "2027-01-01T00:00:00Z",
|
||||
},
|
||||
)
|
||||
assert patched.status_code == 200, patched.text
|
||||
assert patched.json()["budget_id"] != organization["default_budget_id"]
|
||||
assert patched.json()["effective_budget"]["max_budget"] == 5
|
||||
inherited: Final = http.patch(
|
||||
"/management/v1/customer_members/member",
|
||||
json={
|
||||
"budget_override": {"max_budget": None, "allowed_models": []},
|
||||
},
|
||||
)
|
||||
assert inherited.status_code == 200, inherited.text
|
||||
assert inherited.json()["budget_id"] == patched.json()["budget_id"]
|
||||
assert inherited.json()["effective_budget"]["max_budget"] == 100
|
||||
assert inherited.json()["effective_budget"]["soft_budget"] == 4
|
||||
assert inherited.json()["effective_budget"]["tpm_limit"] == 2
|
||||
assert inherited.json()["effective_budget"]["allowed_models"] == ["model-a"]
|
||||
unchanged: Final = http.get("/management/v1/customer_organizations/org")
|
||||
assert unchanged.status_code == 200, unchanged.text
|
||||
assert unchanged.json() == organization
|
||||
|
||||
|
||||
def test_first_added_template_is_default_and_default_cannot_be_unset_without_replacement() -> None:
|
||||
http: Final = client()
|
||||
created: Final = http.post("/management/v1/customer_organizations", json={"customer_org_id": "org"})
|
||||
assert created.status_code == 201, created.text
|
||||
first: Final = http.post("/management/v1/customer_organizations/org/default_budgets", json={"name": "first"})
|
||||
assert first.status_code == 201, first.text
|
||||
first_id: Final = first.json()["default_budget_id"]
|
||||
second: Final = http.post("/management/v1/customer_organizations/org/default_budgets", json={"name": "second"})
|
||||
assert second.status_code == 201, second.text
|
||||
assert second.json()["default_budget_id"] == first_id
|
||||
invalid: Final = http.patch(
|
||||
f"/management/v1/customer_organizations/org/default_budgets/{first_id}", json={"is_default": False}
|
||||
)
|
||||
assert invalid.status_code == 422, invalid.text
|
||||
still_default: Final = http.get("/management/v1/customer_organizations/org")
|
||||
assert still_default.json() == second.json()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"method,path,body",
|
||||
(
|
||||
("POST", "/customer_organizations/missing/default_budgets", {"name": "template"}),
|
||||
("PATCH", "/customer_organizations/missing/default_budgets/budget", {}),
|
||||
("DELETE", "/customer_organizations/missing/default_budgets/budget", None),
|
||||
),
|
||||
)
|
||||
def test_template_routes_return_404_for_a_missing_organization(
|
||||
method: str,
|
||||
path: str,
|
||||
body: dict[str, object] | None,
|
||||
) -> None:
|
||||
response: Final = client().request(method, "/management/v1" + path, json=body)
|
||||
assert response.status_code == 404, response.text
|
||||
|
|
@ -514,3 +514,37 @@ async def test_update_budget_serializes_model_max_budget_for_prisma(
|
|||
"gpt4o": {"max_budget": 5.0, "budget_duration": "1d"},
|
||||
"glm-5.2": {"max_budget": 7.5, "budget_duration": "30d"},
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def admin_client(client_and_mocks):
|
||||
client, mock_prisma, mock_table = client_and_mocks
|
||||
mock_table.delete = AsyncMock(return_value={"budget_id": "tpl-1"})
|
||||
mock_prisma.db.litellm_customerorganizationtable = MagicMock()
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN
|
||||
)
|
||||
return client, mock_prisma, mock_table
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_budget_refuses_customer_organization_default_budget(admin_client):
|
||||
client, mock_prisma, mock_table = admin_client
|
||||
mock_prisma.db.litellm_customerorganizationtable.count = AsyncMock(return_value=1)
|
||||
|
||||
resp: Final = client.post("/budget/delete", json={"id": "tpl-1"})
|
||||
|
||||
assert resp.status_code == 409, resp.text
|
||||
assert "/default_budgets/" in resp.json()["detail"]["error"]
|
||||
mock_table.delete.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_budget_deletes_budget_no_customer_organization_uses(admin_client):
|
||||
client, mock_prisma, mock_table = admin_client
|
||||
mock_prisma.db.litellm_customerorganizationtable.count = AsyncMock(return_value=0)
|
||||
|
||||
resp: Final = client.post("/budget/delete", json={"id": "tpl-1"})
|
||||
|
||||
assert resp.status_code == 200, resp.text
|
||||
mock_table.delete.assert_awaited_once_with(where={"budget_id": "tpl-1"})
|
||||
|
|
|
|||
187
tests/unit/proxy/management_helpers/test_customer_hierarchy.py
Normal file
187
tests/unit/proxy/management_helpers/test_customer_hierarchy.py
Normal file
|
|
@ -0,0 +1,187 @@
|
|||
from collections.abc import AsyncIterator, Mapping, Sequence
|
||||
from contextlib import asynccontextmanager
|
||||
from datetime import datetime, timezone
|
||||
from functools import reduce
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from prisma import Json
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter
|
||||
|
||||
from litellm.proxy.management_helpers.customer_hierarchy import effective_member_budget, related_budget
|
||||
from litellm.types.proxy.management_endpoints.customer_hierarchy import Budget
|
||||
|
||||
NOW: Final = datetime(2026, 1, 1, tzinfo=timezone.utc)
|
||||
|
||||
|
||||
class Record(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="allow")
|
||||
|
||||
|
||||
class Table:
|
||||
def __init__(self, id_field: str, database: "Database") -> None:
|
||||
self.id_field = id_field
|
||||
self.database = database
|
||||
self.rows: Mapping[str, Record] = {}
|
||||
|
||||
async def include(self, row: Record, include: Mapping[str, object] | None) -> Record:
|
||||
if not include:
|
||||
return row
|
||||
relations: Final = {
|
||||
name: self.database.litellm_budgettable.rows.get(str(row.model_dump().get(name + "_id")))
|
||||
for name in include
|
||||
}
|
||||
return row.model_copy(
|
||||
update={name: budget.model_dump() if budget is not None else None for name, budget in relations.items()}
|
||||
)
|
||||
|
||||
async def find_unique(
|
||||
self, where: Mapping[str, object], include: Mapping[str, object] | None = None
|
||||
) -> Record | None:
|
||||
row: Final = self.rows.get(str(where[self.id_field]))
|
||||
return await self.include(row, include) if row else None
|
||||
|
||||
async def create(self, data: Mapping[str, object], include: Mapping[str, object] | None = None) -> Record:
|
||||
from prisma.errors import UniqueViolationError
|
||||
|
||||
if str(data[self.id_field]) in self.rows:
|
||||
raise UniqueViolationError({})
|
||||
row: Final = Record.model_validate(
|
||||
{
|
||||
"spend": 0.0,
|
||||
"budget_reset_at": None,
|
||||
"budget_ids": (),
|
||||
"budget_names": {},
|
||||
"created_at": NOW,
|
||||
"updated_at": NOW,
|
||||
**data,
|
||||
}
|
||||
)
|
||||
self.rows = {**self.rows, str(data[self.id_field]): row}
|
||||
return await self.include(row, include)
|
||||
|
||||
async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> Record | None:
|
||||
row_id: Final = str(where[self.id_field])
|
||||
if row_id not in self.rows:
|
||||
return None
|
||||
values: Final = {key: value.data if isinstance(value, Json) else value for key, value in data.items()}
|
||||
row: Final = self.rows[row_id].model_copy(update=values)
|
||||
self.rows = {**self.rows, row_id: row}
|
||||
return row
|
||||
|
||||
async def find_many(
|
||||
self,
|
||||
where: Mapping[str, object] | None = None,
|
||||
order: Sequence[Mapping[str, object]] | None = None,
|
||||
include: Mapping[str, object] | None = None,
|
||||
skip: int = 0,
|
||||
take: int | None = None,
|
||||
) -> tuple[Record, ...]:
|
||||
matches: Final = tuple(
|
||||
row
|
||||
for row in self.rows.values()
|
||||
if all(
|
||||
row.model_dump().get(key) in value["in"]
|
||||
if isinstance(value, Mapping)
|
||||
else row.model_dump().get(key) == value
|
||||
for key, value in (where or {}).items()
|
||||
)
|
||||
)
|
||||
ordered: Final = reduce(
|
||||
lambda rows, key: tuple(sorted(rows, key=lambda row: str(row.model_dump().get(next(iter(key)))))),
|
||||
reversed(order or ()),
|
||||
matches,
|
||||
)
|
||||
selected: Final = ordered[skip : skip + take] if take is not None else ordered[skip:]
|
||||
return tuple([await self.include(row, include) for row in selected])
|
||||
|
||||
async def count(self, where: Mapping[str, object] | None = None) -> int:
|
||||
return len(await self.find_many(where=where))
|
||||
|
||||
async def delete(self, where: Mapping[str, object]) -> Record | None:
|
||||
row_id: Final = TypeAdapter(str).validate_python(where[self.id_field])
|
||||
row: Final = self.rows.get(row_id)
|
||||
self.rows = {key: value for key, value in self.rows.items() if key != row_id}
|
||||
return row
|
||||
|
||||
async def delete_many(self, where: Mapping[str, object]) -> int:
|
||||
matched: Final = await self.find_many(where=where)
|
||||
self.rows = {key: row for key, row in self.rows.items() if row not in matched}
|
||||
return len(matched)
|
||||
|
||||
|
||||
class Database:
|
||||
def __init__(self) -> None:
|
||||
self.litellm_customerorganizationtable = Table("customer_org_id", self)
|
||||
self.litellm_customerteamtable = Table("customer_team_id", self)
|
||||
self.litellm_customermembertable = Table("customer_member_id", self)
|
||||
self.litellm_budgettable = Table("budget_id", self)
|
||||
self.locked_organization_ids: tuple[str, ...] = ()
|
||||
|
||||
async def query_raw(self, query: str, *args: object) -> tuple[()]:
|
||||
assert "FOR UPDATE" in query
|
||||
self.locked_organization_ids = (*self.locked_organization_ids, *(str(arg) for arg in args))
|
||||
return ()
|
||||
|
||||
@asynccontextmanager
|
||||
async def tx(self) -> AsyncIterator["Database"]:
|
||||
tables: Final = (
|
||||
self.litellm_customerorganizationtable,
|
||||
self.litellm_customerteamtable,
|
||||
self.litellm_customermembertable,
|
||||
self.litellm_budgettable,
|
||||
)
|
||||
snapshots: Final = tuple(table.rows for table in tables)
|
||||
try:
|
||||
yield self
|
||||
except Exception:
|
||||
for table, rows in zip(tables, snapshots):
|
||||
table.rows = rows
|
||||
raise
|
||||
|
||||
|
||||
@pytest.mark.parametrize("expires_at", (None, datetime(2027, 1, 1, tzinfo=timezone.utc)))
|
||||
def test_effective_budget_merges_non_null_fields_and_empty_models_inherit(expires_at: datetime | None) -> None:
|
||||
template: Final = Budget(max_budget=100, rpm_limit=10, allowed_models=("model-a",))
|
||||
override: Final = Budget(max_budget=0, rpm_limit=None, allowed_models=())
|
||||
assert effective_member_budget(template, override, expires_at, NOW) == Budget(
|
||||
max_budget=0,
|
||||
rpm_limit=10,
|
||||
allowed_models=("model-a",),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("expires_at", (NOW, datetime(2025, 1, 1), datetime(2025, 1, 1, tzinfo=timezone.utc)))
|
||||
def test_expired_override_is_ignored_at_the_expiration_boundary(expires_at: datetime) -> None:
|
||||
template: Final = Budget(max_budget=100)
|
||||
assert effective_member_budget(template, Budget(max_budget=1), expires_at, NOW) == template
|
||||
|
||||
|
||||
def test_custom_budget_and_missing_sources_resolve_without_a_template() -> None:
|
||||
assert effective_member_budget(None, Budget(max_budget=4), None, NOW) == Budget(max_budget=4, allowed_models=())
|
||||
assert effective_member_budget(None, None, None, NOW) == Budget(allowed_models=())
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"expiry, expected",
|
||||
(
|
||||
(datetime(2027, 1, 1, tzinfo=timezone.utc), 15),
|
||||
(NOW, 10),
|
||||
),
|
||||
)
|
||||
def test_existing_budget_model_applies_temporary_increase_using_injected_now(expiry: datetime, expected: float) -> None:
|
||||
resolved: Final = related_budget(
|
||||
{
|
||||
"max_budget": 10,
|
||||
"temp_budget_increase": 5,
|
||||
"temp_budget_expiry": expiry,
|
||||
},
|
||||
NOW,
|
||||
)
|
||||
assert resolved == Budget(max_budget=expected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("expiry, expected", ((datetime(2027, 1, 1), 105), (NOW, 100)))
|
||||
def test_member_temporary_increase_is_added_to_its_inherited_maximum(expiry: datetime, expected: float) -> None:
|
||||
resolved: Final = related_budget({"temp_budget_increase": 5, "temp_budget_expiry": expiry}, NOW, 100)
|
||||
assert resolved == Budget(max_budget=expected)
|
||||
|
|
@ -29,6 +29,26 @@ def test_a_fully_documented_post_has_no_findings(request_type: str, response_typ
|
|||
assert checker.violations({"paths": {"/v1/thing": {"post": operation}}}) == ()
|
||||
|
||||
|
||||
def test_a_documented_created_response_satisfies_the_success_schema_rule() -> None:
|
||||
operation: Final = {**_DOCUMENTED, "responses": {"201": _RESPONSES["200"]}}
|
||||
assert checker.violations({"paths": {"/v1/thing": {"post": operation}}}) == ()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("responses", (
|
||||
{},
|
||||
{"201": {"description": "Created"}},
|
||||
{"201": {"content": {"application/json": {"schema": {}}}}},
|
||||
{"422": _RESPONSES["200"]},
|
||||
))
|
||||
def test_an_error_schema_or_an_untyped_created_response_cannot_replace_a_success_schema(
|
||||
responses: Mapping[str, object],
|
||||
) -> None:
|
||||
operation: Final = {**_DOCUMENTED, "responses": responses}
|
||||
assert checker.violations({"paths": {"/v1/thing": {"post": operation}}}) == (
|
||||
checker.Violation("POST", "/v1/thing", "OA003"),
|
||||
)
|
||||
|
||||
|
||||
def test_a_schema_under_another_media_type_still_counts_as_undocumented() -> None:
|
||||
operation: Final = _post_with_media_types("application/octet-stream", "audio/wav")
|
||||
assert {v.code for v in checker.violations({"paths": {"/v1/thing": {"post": operation}}})} == {"OA001", "OA003"}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue