diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260829000000_add_model_access_group_budget_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260829000000_add_model_access_group_budget_table/migration.sql new file mode 100644 index 00000000000..62398da7f04 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260829000000_add_model_access_group_budget_table/migration.sql @@ -0,0 +1,20 @@ +-- CreateTable +CREATE TABLE IF NOT EXISTS "LiteLLM_ModelAccessGroupBudgetTable" ( + "access_group_name" TEXT NOT NULL, + "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "budget_id" TEXT, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "created_by" TEXT, + "updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_by" TEXT, + + CONSTRAINT "LiteLLM_ModelAccessGroupBudgetTable_pkey" PRIMARY KEY ("access_group_name") +); + +-- AddForeignKey +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_ModelAccessGroupBudgetTable_budget_id_fkey') THEN + ALTER TABLE "LiteLLM_ModelAccessGroupBudgetTable" ADD CONSTRAINT "LiteLLM_ModelAccessGroupBudgetTable_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE; + END IF; +END $$; diff --git a/litellm/constants.py b/litellm/constants.py index fc88086805f..24340d7bd55 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1681,6 +1681,7 @@ DEFAULT_MCP_NAMESPACE_CSV_MAX_TOKENS: Final = 16 # Ceilings on the cached auth registries; larger tables fall back to per-row lookups # instead of holding an unbounded id set in every worker. TAG_REGISTRY_MAX_SIZE: Final = 5000 +ACCESS_GROUP_REGISTRY_MAX_SIZE: Final = 5000 END_USER_RESTRICTED_REGISTRY_MAX_SIZE: Final = 5000 # How long a failed registry load is remembered as "unusable", so a degraded Postgres # is not re-scanned on every request on top of the per-id lookups it falls back to. diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index f40b0632398..4d86d258788 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -241,6 +241,7 @@ class Litellm_EntityType(enum.Enum): PROJECT = "project" TAG = "tag" AGENT = "agent" + ACCESS_GROUP = "access_group" # global proxy level entity PROXY = "proxy" @@ -2886,6 +2887,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob ), ) budget_reservation: dict[str, Any] | None = Field(default=None, exclude=True) + matched_access_groups: list[str] | None = Field(default=None, exclude=True) budget_throttle_pct: float | None = Field(default=None, exclude=True) user: Any | None = None # Expanded user object when expand=user is used created_by_user: Any | None = None # Expanded created_by user when expand=user is used @@ -4918,6 +4920,7 @@ class DBSpendUpdateTransactions(TypedDict): org_list_transactions: dict[str, float] | None tag_list_transactions: dict[str, float] | None agent_list_transactions: dict[str, float] | None + access_group_list_transactions: dict[str, float] | None class SpendUpdateQueueItem(TypedDict, total=False): diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index b8df0105b7b..96152240181 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -185,6 +185,21 @@ def tag_registry_cache_key() -> str: return "tag_registry" +#: Cached under ``access_group_registry_cache_key`` when the table exceeds +#: ``ACCESS_GROUP_REGISTRY_MAX_SIZE``: registry unusable, fall back to the per-group lookup. +ACCESS_GROUP_REGISTRY_OVERFLOW_SENTINEL: Final = "__access_group_registry_overflow__" + + +def access_group_cache_key(access_group_name: str) -> str: + """Cache key one model access group budget row is stored under; shared so auth, spend tracking and the management endpoints cannot drift.""" + return f"access_group:{access_group_name}" + + +def access_group_registry_cache_key() -> str: + """Cache key for the set of model access group names that have a budget row.""" + return "access_group_registry" + + #: Cached under ``end_user_restricted_registry_cache_key`` when the restricted set exceeds #: ``END_USER_RESTRICTED_REGISTRY_MAX_SIZE``: registry unusable, fall back to the per-id fetch. END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL: Final = "__end_user_restricted_registry_overflow__" diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index e02f652caf6..de13e5aeb6c 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -100,6 +100,12 @@ class TagRepository(PrismaTableRepository["prisma_models.LiteLLM_TagTable"]): table_name = "litellm_tagtable" +class ModelAccessGroupBudgetRepository( + PrismaTableRepository["prisma_models.LiteLLM_ModelAccessGroupBudgetTable"] +): + table_name = "litellm_modelaccessgroupbudgettable" + + class InvitationLinkRepository(PrismaTableRepository["prisma_models.LiteLLM_InvitationLink"]): table_name = "litellm_invitationlink" diff --git a/schema.prisma b/schema.prisma index 2bb850139a2..8ddbee63973 100644 --- a/schema.prisma +++ b/schema.prisma @@ -29,6 +29,7 @@ model LiteLLM_BudgetTable { keys LiteLLM_VerificationToken[] // multiple keys can have the same budget end_users LiteLLM_EndUserTable[] // multiple end-users can have the same budget tags LiteLLM_TagTable[] // multiple tags can have the same budget + model_access_groups LiteLLM_ModelAccessGroupBudgetTable[] // multiple model access groups can have the same budget team_membership LiteLLM_TeamMembership[] // budgets of Users within a Team organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization } @@ -585,6 +586,17 @@ model LiteLLM_EndUserTable { blocked Boolean @default(false) } +model LiteLLM_ModelAccessGroupBudgetTable { + access_group_name String @id + spend Float @default(0.0) + budget_id String? + litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) + created_at DateTime @default(now()) @map("created_at") + created_by String? + updated_at DateTime @default(now()) @updatedAt + updated_by String? +} + // Track tags with budgets and spend model LiteLLM_TagTable { tag_name String @id