diff --git a/litellm/integrations/otel/model/spans.py b/litellm/integrations/otel/model/spans.py index a2d55aaaa06..d1ed9f7b62b 100644 --- a/litellm/integrations/otel/model/spans.py +++ b/litellm/integrations/otel/model/spans.py @@ -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", diff --git a/litellm/proxy-extras/litellm_proxy_extras/migrations/20261112000000_customer_hierarchy/migration.sql b/litellm/proxy-extras/litellm_proxy_extras/migrations/20261112000000_customer_hierarchy/migration.sql new file mode 100644 index 00000000000..90358d05893 --- /dev/null +++ b/litellm/proxy-extras/litellm_proxy_extras/migrations/20261112000000_customer_hierarchy/migration.sql @@ -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 $$; diff --git a/litellm/proxy-extras/litellm_proxy_extras/schema.prisma b/litellm/proxy-extras/litellm_proxy_extras/schema.prisma index df8134f2b0e..2c0246bb94e 100644 --- a/litellm/proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm/proxy-extras/litellm_proxy_extras/schema.prisma @@ -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()) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d4b6579523c..3ee3ea92d7b 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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 diff --git a/litellm/proxy/management_endpoints/budget_management_endpoints.py b/litellm/proxy/management_endpoints/budget_management_endpoints.py index 0ab559ce009..2502c77ccdb 100644 --- a/litellm/proxy/management_endpoints/budget_management_endpoints.py +++ b/litellm/proxy/management_endpoints/budget_management_endpoints.py @@ -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 diff --git a/litellm/proxy/management_endpoints/management_v1/__init__.py b/litellm/proxy/management_endpoints/management_v1/__init__.py index a02d989cc28..c42a786db0d 100644 --- a/litellm/proxy/management_endpoints/management_v1/__init__.py +++ b/litellm/proxy/management_endpoints/management_v1/__init__.py @@ -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) diff --git a/litellm/proxy/management_endpoints/management_v1/customers.py b/litellm/proxy/management_endpoints/management_v1/customers.py new file mode 100644 index 00000000000..438aa26f75b --- /dev/null +++ b/litellm/proxy/management_endpoints/management_v1/customers.py @@ -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) diff --git a/litellm/proxy/management_helpers/customer_hierarchy.py b/litellm/proxy/management_helpers/customer_hierarchy.py new file mode 100644 index 00000000000..1e89161025e --- /dev/null +++ b/litellm/proxy/management_helpers/customer_hierarchy.py @@ -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, + ) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index df8134f2b0e..2c0246bb94e 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -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()) diff --git a/litellm/repositories/budget_repository.py b/litellm/repositories/budget_repository.py index af6c8cfe6f5..3e623c86c0e 100644 --- a/litellm/repositories/budget_repository.py +++ b/litellm/repositories/budget_repository.py @@ -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.""" diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index 99d2ef97634..7833f58384d 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -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" diff --git a/litellm/schema.prisma b/litellm/schema.prisma index df8134f2b0e..2c0246bb94e 100644 --- a/litellm/schema.prisma +++ b/litellm/schema.prisma @@ -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()) diff --git a/litellm/scripts/check_openapi_docs.py b/litellm/scripts/check_openapi_docs.py index e88594f5a51..acd4fcb6a25 100755 --- a/litellm/scripts/check_openapi_docs.py +++ b/litellm/scripts/check_openapi_docs.py @@ -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" diff --git a/litellm/types/proxy/management_endpoints/customer_hierarchy.py b/litellm/types/proxy/management_endpoints/customer_hierarchy.py new file mode 100644 index 00000000000..5866630be12 --- /dev/null +++ b/litellm/types/proxy/management_endpoints/customer_hierarchy.py @@ -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 diff --git a/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts b/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts index 3076fe8a4e2..c49a206f429 100644 --- a/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -9584,6 +9584,206 @@ export interface paths { patch?: never; trace?: never; }; + "/management/v1/customer_members": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * List Customer Members + * @description List customer members with page-number pagination and resolved budgets. Proxy admins only. + */ + get: operations["list_customer_members_management_v1_customer_members_get"]; + put?: never; + /** + * Create Customer Member + * @description Create a customer member with resolved budgets. Proxy admins only. + */ + post: operations["create_customer_member_management_v1_customer_members_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/management/v1/customer_members/{customer_member_id}": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Customer Member + * @description Get a customer member with resolved budgets. Proxy admins only. + */ + get: operations["get_customer_member_management_v1_customer_members__customer_member_id__get"]; + put?: never; + post?: never; + /** + * Delete Customer Member + * @description Delete a customer member and its owned budgets. Parents with children return 409. Proxy admins only. + */ + delete: operations["delete_customer_member_management_v1_customer_members__customer_member_id__delete"]; + options?: never; + head?: never; + /** + * Update Customer Member + * @description Update a customer member and its budgets. Proxy admins only. + */ + patch: operations["update_customer_member_management_v1_customer_members__customer_member_id__patch"]; + trace?: never; + }; + "/management/v1/customer_organizations": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * List Customer Organizations + * @description List customer organizations with page-number pagination and resolved budgets. Proxy admins only. + */ + get: operations["list_customer_organizations_management_v1_customer_organizations_get"]; + put?: never; + /** + * Create Customer Organization + * @description Create a customer organization with resolved budgets. Proxy admins only. + */ + post: operations["create_customer_organization_management_v1_customer_organizations_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/management/v1/customer_organizations/{customer_org_id}": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Customer Organization + * @description Get a customer organization with resolved budgets. Proxy admins only. + */ + get: operations["get_customer_organization_management_v1_customer_organizations__customer_org_id__get"]; + put?: never; + post?: never; + /** + * Delete Customer Organization + * @description Delete a customer organization and its owned budgets. Parents with children return 409. Proxy admins only. + */ + delete: operations["delete_customer_organization_management_v1_customer_organizations__customer_org_id__delete"]; + options?: never; + head?: never; + /** + * Update Customer Organization + * @description Update a customer organization alias or overall budget. Proxy admins only. + */ + patch: operations["update_customer_organization_management_v1_customer_organizations__customer_org_id__patch"]; + trace?: never; + }; + "/management/v1/customer_organizations/{customer_org_id}/default_budgets": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Add Customer Budget Template + * @description Add a named budget template to a customer organization. Proxy admins only. + */ + post: operations["add_customer_budget_template_management_v1_customer_organizations__customer_org_id__default_budgets_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/management/v1/customer_organizations/{customer_org_id}/default_budgets/{budget_id}": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + post?: never; + /** + * Delete Customer Budget Template + * @description Delete an unreferenced, non-default customer budget template. Proxy admins only. + */ + delete: operations["delete_customer_budget_template_management_v1_customer_organizations__customer_org_id__default_budgets__budget_id__delete"]; + options?: never; + head?: never; + /** + * Update Customer Budget Template + * @description Edit a budget template in place. Members inherit the updated values. Proxy admins only. + */ + patch: operations["update_customer_budget_template_management_v1_customer_organizations__customer_org_id__default_budgets__budget_id__patch"]; + trace?: never; + }; + "/management/v1/customer_teams": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * List Customer Teams + * @description List customer teams with page-number pagination and resolved budgets. Proxy admins only. + */ + get: operations["list_customer_teams_management_v1_customer_teams_get"]; + put?: never; + /** + * Create Customer Team + * @description Create a customer team with resolved budgets. Proxy admins only. + */ + post: operations["create_customer_team_management_v1_customer_teams_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/management/v1/customer_teams/{customer_team_id}": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Customer Team + * @description Get a customer team with resolved budgets. Proxy admins only. + */ + get: operations["get_customer_team_management_v1_customer_teams__customer_team_id__get"]; + put?: never; + post?: never; + /** + * Delete Customer Team + * @description Delete a customer team and its owned budgets. Parents with children return 409. Proxy admins only. + */ + delete: operations["delete_customer_team_management_v1_customer_teams__customer_team_id__delete"]; + options?: never; + head?: never; + /** + * Update Customer Team + * @description Update a customer team and its budgets. Proxy admins only. + */ + patch: operations["update_customer_team_management_v1_customer_teams__customer_team_id__patch"]; + trace?: never; + }; "/management/v1/spend_logs/end_users": { parameters: { query?: never; @@ -29158,6 +29358,21 @@ export interface components { [key: string]: components["schemas"]["MetricWithMetadata"]; } | null; }; + /** Budget */ + Budget: { + /** Allowed Models */ + allowed_models?: string[] | null; + /** Budget Duration */ + budget_duration?: string | null; + /** Max Budget */ + max_budget?: number | null; + /** Rpm Limit */ + rpm_limit?: number | null; + /** Soft Budget */ + soft_budget?: number | null; + /** Tpm Limit */ + tpm_limit?: number | null; + }; /** BudgetConfig */ BudgetConfig: { /** Budget Duration */ @@ -29296,6 +29511,28 @@ export interface components { /** Budgets */ budgets: string[]; }; + /** BudgetTemplate */ + BudgetTemplate: { + /** Allowed Models */ + allowed_models?: string[] | null; + /** Budget Duration */ + budget_duration?: string | null; + /** + * Is Default + * @default false + */ + is_default: boolean; + /** Max Budget */ + max_budget?: number | null; + /** Name */ + name: string; + /** Rpm Limit */ + rpm_limit?: number | null; + /** Soft Budget */ + soft_budget?: number | null; + /** Tpm Limit */ + tpm_limit?: number | null; + }; /** * BulkDeleteUserRequest * @description Body of `POST /management/v1/users/bulk_delete`. @@ -32811,6 +33048,44 @@ export interface components { */ type: "custom"; }; + /** CustomerMember */ + CustomerMember: { + /** Alias */ + alias?: string | null; + /** Budget Id */ + budget_id?: string | null; + budget_override?: components["schemas"]["Budget"] | null; + /** Budget Override Expires At */ + budget_override_expires_at?: string | null; + /** Budget Reset At */ + budget_reset_at?: string | null; + /** Budget Template Id */ + budget_template_id?: string | null; + /** + * Created At + * Format: date-time + */ + created_at: string; + /** Created By */ + created_by: string; + /** Customer Member Id */ + customer_member_id: string; + effective_budget: components["schemas"]["Budget"]; + /** Parent Team Id */ + parent_team_id: string; + /** + * Spend + * @default 0 + */ + spend: number; + /** + * Updated At + * Format: date-time + */ + updated_at: string; + /** Updated By */ + updated_by: string; + }; /** * CustomerResponse * @description Customer object returned by the /customer read+write endpoints. @@ -35840,6 +36115,13 @@ export interface components { links: components["schemas"]["ListLinks"]; meta: components["schemas"]["ListMeta"]; }; + /** ListResponse[CustomerMember] */ + ListResponse_CustomerMember_: { + /** Data */ + data: components["schemas"]["CustomerMember"][]; + links: components["schemas"]["ListLinks"]; + meta: components["schemas"]["ListMeta"]; + }; /** ListResponse[ModelGroupInfoProxy] */ ListResponse_ModelGroupInfoProxy_: { /** Data */ @@ -35847,6 +36129,20 @@ export interface components { links: components["schemas"]["ListLinks"]; meta: components["schemas"]["ListMeta"]; }; + /** ListResponse[Organization] */ + ListResponse_Organization_: { + /** Data */ + data: components["schemas"]["Organization"][]; + links: components["schemas"]["ListLinks"]; + meta: components["schemas"]["ListMeta"]; + }; + /** ListResponse[Team] */ + ListResponse_Team_: { + /** Data */ + data: components["schemas"]["Team"][]; + links: components["schemas"]["ListLinks"]; + meta: components["schemas"]["ListMeta"]; + }; /** * ListRunsResponse * @description Response from listing runs @@ -40023,6 +40319,30 @@ export interface components { */ user_id?: string | null; }; + /** MemberCreate */ + MemberCreate: { + /** Alias */ + alias?: string | null; + budget_override?: components["schemas"]["Budget"] | null; + /** Budget Override Expires At */ + budget_override_expires_at?: string | null; + /** Budget Template Id */ + budget_template_id?: string | null; + /** Customer Member Id */ + customer_member_id?: string | null; + /** Parent Team Id */ + parent_team_id: string; + }; + /** MemberPatch */ + MemberPatch: { + /** Alias */ + alias?: string | null; + budget_override?: components["schemas"]["Budget"] | null; + /** Budget Override Expires At */ + budget_override_expires_at?: string | null; + /** Budget Template Id */ + budget_template_id?: string | null; + }; /** MemoryCreateRequest */ MemoryCreateRequest: { /** @@ -41885,6 +42205,44 @@ export interface components { */ user_id?: string | null; }; + /** Organization */ + Organization: { + /** Alias */ + alias?: string | null; + budget?: components["schemas"]["Budget"] | null; + /** Budget Id */ + budget_id?: string | null; + /** Budget Reset At */ + budget_reset_at?: string | null; + /** + * Created At + * Format: date-time + */ + created_at: string; + /** Created By */ + created_by: string; + /** Customer Org Id */ + customer_org_id: string; + /** Default Budget Id */ + default_budget_id?: string | null; + /** + * Default Budgets + * @default [] + */ + default_budgets: components["schemas"]["OrganizationBudget"][]; + /** + * Spend + * @default 0 + */ + spend: number; + /** + * Updated At + * Format: date-time + */ + updated_at: string; + /** Updated By */ + updated_by: string; + }; /** OrganizationAddMemberResponse */ OrganizationAddMemberResponse: { /** Organization Id */ @@ -41894,6 +42252,43 @@ export interface components { /** Updated Users */ updated_users: components["schemas"]["LiteLLM_UserTable"][]; }; + /** OrganizationBudget */ + OrganizationBudget: { + /** Allowed Models */ + allowed_models?: string[] | null; + /** Budget Duration */ + budget_duration?: string | null; + /** Budget Id */ + budget_id: string; + /** + * Is Default + * @default false + */ + is_default: boolean; + /** Max Budget */ + max_budget?: number | null; + /** Name */ + name: string; + /** Rpm Limit */ + rpm_limit?: number | null; + /** Soft Budget */ + soft_budget?: number | null; + /** Tpm Limit */ + tpm_limit?: number | null; + }; + /** OrganizationCreate */ + OrganizationCreate: { + /** Alias */ + alias?: string | null; + budget?: components["schemas"]["Budget"] | null; + /** Customer Org Id */ + customer_org_id?: string | null; + /** + * Default Budgets + * @default [] + */ + default_budgets: components["schemas"]["BudgetTemplate"][]; + }; /** OrganizationMemberAddRequest */ OrganizationMemberAddRequest: { /** Max Budget In Organization */ @@ -41924,6 +42319,12 @@ export interface components { /** User Id */ user_id?: string | null; }; + /** OrganizationPatch */ + OrganizationPatch: { + /** Alias */ + alias?: string | null; + budget?: components["schemas"]["Budget"] | null; + }; /** OrganizationRequest */ OrganizationRequest: { /** Organizations */ @@ -48731,6 +49132,39 @@ export interface components { /** Tpm Limit */ tpm_limit?: number | null; }; + /** Team */ + Team: { + /** Alias */ + alias?: string | null; + budget?: components["schemas"]["Budget"] | null; + /** Budget Id */ + budget_id?: string | null; + /** Budget Reset At */ + budget_reset_at?: string | null; + /** + * Created At + * Format: date-time + */ + created_at: string; + /** Created By */ + created_by: string; + /** Customer Team Id */ + customer_team_id: string; + /** Parent Org Id */ + parent_org_id: string; + /** + * Spend + * @default 0 + */ + spend: number; + /** + * Updated At + * Format: date-time + */ + updated_at: string; + /** Updated By */ + updated_by: string; + }; /** TeamAddMemberResponse */ TeamAddMemberResponse: { /** Access Group Ids */ @@ -48854,6 +49288,16 @@ export interface components { /** Team Id */ team_id: string; }; + /** TeamCreate */ + TeamCreate: { + /** Alias */ + alias?: string | null; + budget?: components["schemas"]["Budget"] | null; + /** Customer Team Id */ + customer_team_id?: string | null; + /** Parent Org Id */ + parent_org_id: string; + }; /** * TeamListCallerMembership * @description The calling user's own spend and effective member budget in a team. @@ -49333,6 +49777,12 @@ export interface components { /** Team Id */ team_id: string; }; + /** TeamPatch */ + TeamPatch: { + /** Alias */ + alias?: string | null; + budget?: components["schemas"]["Budget"] | null; + }; /** TeamUserSpendResponse */ TeamUserSpendResponse: { /** End Date */ @@ -49455,6 +49905,25 @@ export interface components { /** Groups */ groups: string[]; }; + /** TemplatePatch */ + TemplatePatch: { + /** Allowed Models */ + allowed_models?: string[] | null; + /** Budget Duration */ + budget_duration?: string | null; + /** Is Default */ + is_default?: boolean | null; + /** Max Budget */ + max_budget?: number | null; + /** Name */ + name?: string | null; + /** Rpm Limit */ + rpm_limit?: number | null; + /** Soft Budget */ + soft_budget?: number | null; + /** Tpm Limit */ + tpm_limit?: number | null; + }; /** * TestCustomCodeGuardrailRequest * @description Request model for testing custom code guardrails. @@ -67533,6 +68002,589 @@ export interface operations { }; }; }; + list_customer_members_management_v1_customer_members_get: { + parameters: { + query?: { + page?: number; + page_size?: number; + parent_team_id?: string | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["ListResponse_CustomerMember_"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + create_customer_member_management_v1_customer_members_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["MemberCreate"]; + }; + }; + responses: { + /** @description Successful Response */ + 201: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["CustomerMember"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + get_customer_member_management_v1_customer_members__customer_member_id__get: { + parameters: { + query?: never; + header?: never; + path: { + customer_member_id: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["CustomerMember"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + delete_customer_member_management_v1_customer_members__customer_member_id__delete: { + parameters: { + query?: never; + header?: never; + path: { + customer_member_id: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 204: { + headers: { + [name: string]: unknown; + }; + content?: never; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + update_customer_member_management_v1_customer_members__customer_member_id__patch: { + parameters: { + query?: never; + header?: never; + path: { + customer_member_id: string; + }; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["MemberPatch"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["CustomerMember"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + list_customer_organizations_management_v1_customer_organizations_get: { + parameters: { + query?: { + page?: number; + page_size?: number; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["ListResponse_Organization_"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + create_customer_organization_management_v1_customer_organizations_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["OrganizationCreate"]; + }; + }; + responses: { + /** @description Successful Response */ + 201: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["Organization"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + get_customer_organization_management_v1_customer_organizations__customer_org_id__get: { + parameters: { + query?: never; + header?: never; + path: { + customer_org_id: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["Organization"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + delete_customer_organization_management_v1_customer_organizations__customer_org_id__delete: { + parameters: { + query?: never; + header?: never; + path: { + customer_org_id: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 204: { + headers: { + [name: string]: unknown; + }; + content?: never; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + update_customer_organization_management_v1_customer_organizations__customer_org_id__patch: { + parameters: { + query?: never; + header?: never; + path: { + customer_org_id: string; + }; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["OrganizationPatch"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["Organization"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + add_customer_budget_template_management_v1_customer_organizations__customer_org_id__default_budgets_post: { + parameters: { + query?: never; + header?: never; + path: { + customer_org_id: string; + }; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["BudgetTemplate"]; + }; + }; + responses: { + /** @description Successful Response */ + 201: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["Organization"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + delete_customer_budget_template_management_v1_customer_organizations__customer_org_id__default_budgets__budget_id__delete: { + parameters: { + query?: never; + header?: never; + path: { + customer_org_id: string; + budget_id: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 204: { + headers: { + [name: string]: unknown; + }; + content?: never; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + update_customer_budget_template_management_v1_customer_organizations__customer_org_id__default_budgets__budget_id__patch: { + parameters: { + query?: never; + header?: never; + path: { + customer_org_id: string; + budget_id: string; + }; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["TemplatePatch"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["Organization"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + list_customer_teams_management_v1_customer_teams_get: { + parameters: { + query?: { + page?: number; + page_size?: number; + parent_org_id?: string | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["ListResponse_Team_"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + create_customer_team_management_v1_customer_teams_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["TeamCreate"]; + }; + }; + responses: { + /** @description Successful Response */ + 201: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["Team"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + get_customer_team_management_v1_customer_teams__customer_team_id__get: { + parameters: { + query?: never; + header?: never; + path: { + customer_team_id: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["Team"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + delete_customer_team_management_v1_customer_teams__customer_team_id__delete: { + parameters: { + query?: never; + header?: never; + path: { + customer_team_id: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 204: { + headers: { + [name: string]: unknown; + }; + content?: never; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + update_customer_team_management_v1_customer_teams__customer_team_id__patch: { + parameters: { + query?: never; + header?: never; + path: { + customer_team_id: string; + }; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["TeamPatch"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["Team"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; list_spend_log_end_users_management_v1_spend_logs_end_users_get: { parameters: { query: { diff --git a/tests/unit/proxy/auth/test_route_checks.py b/tests/unit/proxy/auth/test_route_checks.py index 988e3a49ea0..3883514e34a 100644 --- a/tests/unit/proxy/auth/test_route_checks.py +++ b/tests/unit/proxy/auth/test_route_checks.py @@ -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: diff --git a/tests/unit/proxy/management_endpoints/management_v1/test_customers.py b/tests/unit/proxy/management_endpoints/management_v1/test_customers.py new file mode 100644 index 00000000000..baa89ae342b --- /dev/null +++ b/tests/unit/proxy/management_endpoints/management_v1/test_customers.py @@ -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 diff --git a/tests/unit/proxy/management_endpoints/test_budget_endpoints.py b/tests/unit/proxy/management_endpoints/test_budget_endpoints.py index eb4fa2106c4..1e5e134c0b0 100644 --- a/tests/unit/proxy/management_endpoints/test_budget_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_budget_endpoints.py @@ -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"}) diff --git a/tests/unit/proxy/management_helpers/test_customer_hierarchy.py b/tests/unit/proxy/management_helpers/test_customer_hierarchy.py new file mode 100644 index 00000000000..1e79655de8e --- /dev/null +++ b/tests/unit/proxy/management_helpers/test_customer_hierarchy.py @@ -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) diff --git a/tests/unit/test_check_openapi_docs.py b/tests/unit/test_check_openapi_docs.py index a277eeb6932..cd74923c57b 100644 --- a/tests/unit/test_check_openapi_docs.py +++ b/tests/unit/test_check_openapi_docs.py @@ -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"}