diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 95e69bf7f1d..83969d8dedf 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -3,7 +3,7 @@ "limit": 17270 }, "reportArgumentType": { - "limit": 2539 + "limit": 2538 }, "reportAssignmentType": { "limit": 319 diff --git a/codecov.yaml b/codecov.yaml index bc0b3604329..4d93c18f3ac 100644 --- a/codecov.yaml +++ b/codecov.yaml @@ -25,6 +25,8 @@ flag_management: carryforward: false - name: proxy-db-schema-migration carryforward: false + - name: circleci + carryforward: false component_management: individual_components: 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-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 7818e3acd56..60223265211 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/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,20 @@ model LiteLLM_EndUserTable { blocked Boolean @default(false) } +// Budget and shared spend for a model access group. The groups themselves are not rows anywhere: +// they are free-text strings in LiteLLM_ProxyModelTable.model_info.access_groups, so a row here +// exists only once someone gives that group a budget. +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 diff --git a/litellm/constants.py b/litellm/constants.py index c1f087316cc..cc6db6c10cc 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1683,6 +1683,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 +MODEL_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/litellm_core_utils/internal_call_metadata.py b/litellm/litellm_core_utils/internal_call_metadata.py index 34d5797a6d8..4d043701f40 100644 --- a/litellm/litellm_core_utils/internal_call_metadata.py +++ b/litellm/litellm_core_utils/internal_call_metadata.py @@ -25,6 +25,13 @@ from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, Inter BUDGET_RESERVATION_METADATA_KEYS: Final = frozenset({"user_api_key_budget_reservation"}) +MODEL_ACCESS_GROUP_METADATA_KEY: Final = "user_api_key_matched_model_access_groups" +"""Where auth records the model access groups that authorized the request, for the spend writer. + +The ``user_api_key`` prefix is load-bearing, not cosmetic: when a request carries both +``metadata`` and ``litellm_metadata``, ``get_litellm_metadata_from_kwargs`` returns the latter and +copies a key across only when ``user_api_key`` appears in its name.""" + _USER_API_KEY_AUTH_KEY: Final = "user_api_key_auth" FORWARDABLE_IDENTITY_METADATA_KEYS: Final = frozenset( diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index ff9a5d41c2d..c27822d0479 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -64,7 +64,10 @@ from litellm.integrations.mlflow import MlflowLogger from litellm.integrations.sqs import SQSLogger from litellm.litellm_core_utils.core_helpers import is_expected_client_error, reconstruct_model_name from litellm.litellm_core_utils.get_litellm_params import get_litellm_params -from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call +from litellm.litellm_core_utils.internal_call_metadata import ( + MODEL_ACCESS_GROUP_METADATA_KEY, + is_unbilled_non_inference_call, +) from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import ( cost_breakdown_with_guardrail, guardrail_information_cost, @@ -5051,6 +5054,42 @@ def is_valid_sha256_hash(value: str) -> bool: return bool(re.fullmatch(r"[a-fA-F0-9]{64}", value)) +def coerce_model_access_groups(value: object) -> tuple[str, ...]: + """Model access group names out of untrusted request metadata, deduped and order preserving.""" + if not isinstance(value, (list, tuple)): + return () + return tuple(dict.fromkeys(group for group in value if isinstance(group, str) and group)) + + +def _model_access_groups_on_auth_object(user_api_key_auth: object) -> object: + if isinstance(user_api_key_auth, Mapping): + return user_api_key_auth.get("matched_model_access_groups") + return getattr(user_api_key_auth, "matched_model_access_groups", None) + + +def _model_access_groups_from_metadata(metadata: Mapping[str, object]) -> tuple[str, ...]: + stamped: Final = coerce_model_access_groups(metadata.get(MODEL_ACCESS_GROUP_METADATA_KEY)) + if stamped: + return stamped + return coerce_model_access_groups(_model_access_groups_on_auth_object(metadata.get("user_api_key_auth"))) + + +def request_model_access_groups_from_litellm_params(litellm_params: Mapping[str, object]) -> tuple[str, ...]: + """Access groups the auth layer stamped onto this request, from whichever metadata field carries them. + + Detached internal sub-calls only inherit the identity keys, so the auth object is the + fallback there, exactly as _get_budget_reservation_from_metadata does for reservations. + """ + for metadata_variable_name in ("metadata", "litellm_metadata"): + metadata = litellm_params.get(metadata_variable_name) + if not isinstance(metadata, Mapping): + continue + model_access_groups = _model_access_groups_from_metadata(metadata) + if model_access_groups: + return model_access_groups + return () + + class StandardLoggingPayloadSetup: @staticmethod def cleanup_timestamps( @@ -5902,6 +5941,7 @@ def get_standard_logging_object_payload( request_tags: Final = StandardLoggingPayloadSetup._get_request_tags( litellm_params=litellm_params, proxy_server_request=proxy_server_request ) + request_model_access_groups: Final = request_model_access_groups_from_litellm_params(litellm_params) # cleanup timestamps ( @@ -6064,6 +6104,7 @@ def get_standard_logging_object_payload( prompt_tokens=usage_dict.get("prompt_tokens", 0), completion_tokens=usage_dict.get("completion_tokens", 0), request_tags=request_tags, + request_model_access_groups=request_model_access_groups, end_user=end_user_id, api_base=StandardLoggingPayloadSetup.strip_trailing_slash(litellm_params.get("api_base", "")) or "", model_group=_model_group, @@ -6277,6 +6318,7 @@ def create_dummy_standard_logging_payload() -> StandardLoggingPayload: cache_key=None, saved_cache_cost=saved_cache_cost, request_tags=[], + request_model_access_groups=(), end_user=None, requester_ip_address="127.0.0.1", messages=messages, diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 08a35b41126..1d1b7f057cb 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -461,6 +461,145 @@ "access_groups": { "components": { "schemas": { + "AccessGroupBudget": { + "properties": { + "budget_duration": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Budget Duration" + }, + "budget_id": { + "title": "Budget Id", + "type": "string" + }, + "budget_reset_at": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Budget Reset At" + }, + "max_budget": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Max Budget" + }, + "soft_budget": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Soft Budget" + } + }, + "required": [ + "budget_id" + ], + "title": "AccessGroupBudget", + "type": "object" + }, + "AccessGroupBudgetRequest": { + "additionalProperties": false, + "properties": { + "budget_duration": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Budget Duration" + }, + "budget_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Budget Id" + }, + "max_budget": { + "anyOf": [ + { + "minimum": 0.0, + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Max Budget" + }, + "soft_budget": { + "anyOf": [ + { + "minimum": 0.0, + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Soft Budget" + } + }, + "title": "AccessGroupBudgetRequest", + "type": "object" + }, + "AccessGroupBudgetResponse": { + "properties": { + "access_group": { + "title": "Access Group", + "type": "string" + }, + "budget": { + "anyOf": [ + { + "$ref": "#/components/schemas/AccessGroupBudget" + }, + { + "type": "null" + } + ] + }, + "spend": { + "title": "Spend", + "type": "number" + } + }, + "required": [ + "access_group", + "spend" + ], + "title": "AccessGroupBudgetResponse", + "type": "object" + }, "AccessGroupCreateRequest": { "properties": { "access_agent_ids": { @@ -561,6 +700,16 @@ "title": "Access Group", "type": "string" }, + "budget": { + "anyOf": [ + { + "$ref": "#/components/schemas/AccessGroupBudget" + }, + { + "type": "null" + } + ] + }, "deployment_count": { "title": "Deployment Count", "type": "integer" @@ -571,6 +720,17 @@ }, "title": "Model Names", "type": "array" + }, + "spend": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Spend" } }, "required": [ @@ -782,6 +942,29 @@ "title": "AccessGroupUpdateRequest", "type": "object" }, + "DeleteAccessGroupBudgetResponse": { + "properties": { + "access_group": { + "title": "Access Group", + "type": "string" + }, + "budget_deleted": { + "title": "Budget Deleted", + "type": "boolean" + }, + "message": { + "title": "Message", + "type": "string" + } + }, + "required": [ + "access_group", + "budget_deleted", + "message" + ], + "title": "DeleteAccessGroupBudgetResponse", + "type": "object" + }, "DeleteModelGroupResponse": { "properties": { "access_group": { @@ -1072,6 +1255,156 @@ ] } }, + "/access_group/{access_group}/budget": { + "delete": { + "description": "Clear the shared budget of an access group, leaving the group itself in place.\n\nExample:\n```bash\ncurl -X DELETE 'http://localhost:4000/access_group/production-models/budget' \\\n -H 'Authorization: Bearer sk-1234'\n```\n\nParameters:\n- access_group: str - The access group name (URL path parameter)\n\nReturns:\n- DeleteAccessGroupBudgetResponse; budget_deleted is false when there was nothing to clear\n\nRaises:\n- HTTPException 404: If access group not found", + "operationId": "delete_access_group_budget_access_group__access_group__budget_delete", + "parameters": [ + { + "in": "path", + "name": "access_group", + "required": true, + "schema": { + "title": "Access Group", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/DeleteAccessGroupBudgetResponse" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Delete Access Group Budget", + "tags": [ + "access_groups" + ] + }, + "get": { + "description": "Get the shared budget of an access group, and the spend drawn against it.\n\nExample:\n```bash\ncurl -X GET 'http://localhost:4000/access_group/production-models/budget' \\\n -H 'Authorization: Bearer sk-1234'\n```\n\nParameters:\n- access_group: str - The access group name (URL path parameter)\n\nReturns:\n- AccessGroupBudgetResponse; budget is null when the group has no budget set\n\nRaises:\n- HTTPException 404: If access group not found", + "operationId": "get_access_group_budget_access_group__access_group__budget_get", + "parameters": [ + { + "in": "path", + "name": "access_group", + "required": true, + "schema": { + "title": "Access Group", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AccessGroupBudgetResponse" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Get Access Group Budget", + "tags": [ + "access_groups" + ] + }, + "put": { + "description": "Set or replace the shared budget of an access group. Idempotent.\n\nEvery key that can reach a model in the group draws from this one budget.\n\nExample:\n```bash\ncurl -X PUT 'http://localhost:4000/access_group/production-models/budget' \\\n -H 'Authorization: Bearer sk-1234' \\\n -H 'Content-Type: application/json' \\\n -d '{\n \"max_budget\": 100.0,\n \"budget_duration\": \"30d\"\n }'\n```\n\nParameters:\n- access_group: str - The access group name (URL path parameter)\n- max_budget: Optional[float] - Requests fail once the group's shared spend exceeds this\n- soft_budget: Optional[float] - Fires an alert when reached; requests still succeed\n- budget_duration: Optional[str] - Frequency of resetting the group's spend (e.g. '30d')\n- budget_id: Optional[str] - Link an existing budget instead of creating one\n\nReturns:\n- AccessGroupBudgetResponse with the stored budget and current spend\n\nRaises:\n- HTTPException 400: If no budget field is given, or budget_duration cannot be parsed\n- HTTPException 404: If access group not found", + "operationId": "set_access_group_budget_access_group__access_group__budget_put", + "parameters": [ + { + "in": "path", + "name": "access_group", + "required": true, + "schema": { + "title": "Access Group", + "type": "string" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AccessGroupBudgetRequest" + } + } + }, + "required": true + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AccessGroupBudgetResponse" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Set Access Group Budget", + "tags": [ + "access_groups" + ] + } + }, "/access_group/{access_group}/delete": { "delete": { "description": "Delete an access group.\n\nRemoves the access group from all deployments that have it.\n\nExample:\n```bash\ncurl -X DELETE 'http://localhost:4000/access_group/production-models/delete' \\\n -H 'Authorization: Bearer sk-1234'\n```\n\nParameters:\n- access_group: str - The access group name (URL path parameter)\n\nReturns:\n- DeleteModelGroupResponse with deletion details\n\nRaises:\n- HTTPException 404: If access group not found", @@ -1122,7 +1455,7 @@ }, "/access_group/{access_group}/info": { "get": { - "description": "Get information about a specific access group.\n\nExample:\n```bash\ncurl -X GET 'http://localhost:4000/access_group/production-models/info' \\\n -H 'Authorization: Bearer sk-1234'\n```\n\nParameters:\n- access_group: str - The access group name (URL path parameter)\n\nReturns:\n- AccessGroupInfo with the access group details\n\nRaises:\n- HTTPException 404: If access group not found", + "description": "Get information about a specific access group.\n\nExample:\n```bash\ncurl -X GET 'http://localhost:4000/access_group/production-models/info' \\\n -H 'Authorization: Bearer sk-1234'\n```\n\nParameters:\n- access_group: str - The access group name (URL path parameter)\n\nReturns:\n- AccessGroupInfo with the access group details, its shared budget and its spend\n\nRaises:\n- HTTPException 404: If access group not found", "operationId": "get_access_group_info_access_group__access_group__info_get", "parameters": [ { diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 89f79548845..cbd3571ec47 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" + MODEL_ACCESS_GROUP = "model_access_group" # global proxy level entity PROXY = "proxy" @@ -2887,6 +2888,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob ), ) budget_reservation: dict[str, Any] | None = Field(default=None, exclude=True) + matched_model_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 @@ -4921,6 +4923,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 + model_access_group_list_transactions: ReadOnly[dict[str, float] | None] class SpendUpdateQueueItem(TypedDict, total=False): diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 657ca4dbf0b..5703c6cd5e8 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -32,6 +32,7 @@ from litellm.constants import ( DEFAULT_MAX_RECURSE_DEPTH, EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE, END_USER_RESTRICTED_REGISTRY_MAX_SIZE, + MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE, REGISTRY_ERROR_NEGATIVE_CACHE_TTL, TAG_REGISTRY_MAX_SIZE, ) @@ -78,11 +79,15 @@ from litellm.proxy.common_utils.http_parsing_utils import ( from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.common_utils.user_api_key_cache import ( END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL, + MODEL_ACCESS_GROUP_REGISTRY_OVERFLOW_SENTINEL, TAG_REGISTRY_OVERFLOW_SENTINEL, UserApiKeyCache, end_user_cache_key, end_user_restricted_registry_cache_key, get_management_object_ttl, + model_access_group_cache_key, + model_access_group_registry_cache_key, + model_access_group_spend_counter_key, object_permission_cache_key, tag_cache_key, tag_registry_cache_key, @@ -107,12 +112,14 @@ from litellm.repositories.table_repositories import ( EndUserRepository, JWTKeyMappingRepository, ManagedVectorStoresRepository, + ModelAccessGroupBudgetRepository, TagRepository, TeamMembershipRepository, ) from litellm.repositories.team_repository import TeamRepository from litellm.repositories.user_repository import UserRepository from litellm.router import Router +from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget from litellm.utils import get_utc_datetime from .auth_checks_organization import ( @@ -251,6 +258,43 @@ def _end_user_table(repo: _PrismaTableHolder[_PrismaEndUserRow]) -> _PrismaAuthT return repo.table +class _PrismaMaxBudgetRow(Protocol): + @property + def max_budget(self) -> float | None: ... + + +class _PrismaModelAccessGroupBudgetRow(Protocol): + access_group_name: str + + @property + def spend(self) -> float | None: ... + + @property + def litellm_budget_table(self) -> _PrismaMaxBudgetRow | None: ... + + +def _model_access_group_budget_table( + repo: _PrismaTableHolder[_PrismaModelAccessGroupBudgetRow], +) -> _PrismaAuthTable[_PrismaModelAccessGroupBudgetRow]: + return repo.table + + +class _MemberModelScope(Protocol): + @property + def allowed_models(self) -> Sequence[str] | None: ... + + +class _TeamMembershipModelScope(Protocol): + @property + def litellm_budget_table(self) -> _MemberModelScope | None: ... + + +def _member_allowed_models(membership: _TeamMembershipModelScope) -> Sequence[str]: + """The member's own model scope, read through a narrowed view of the membership row.""" + budget_table: Final = membership.litellm_budget_table + return () if budget_table is None else (budget_table.allowed_models or ()) + + class _RawCacheRead(Protocol): async def async_get_cache(self, *, key: str) -> object: ... @@ -807,6 +851,7 @@ async def common_checks( 1.1. If project is blocked 2. If team can call model 2.2 If project can call model + 2.3 Which model access groups authorized this request 3. If team is in budget 3.0.2. If project is in budget 3.0.3. If project is over soft budget (alert only) @@ -925,6 +970,18 @@ async def common_checks( proxy_logging_obj=proxy_logging_obj, ) + # 2.3 Which model access groups authorized this request + matched_model_access_groups: Final = await stamp_matched_model_access_groups( + model=_model, + valid_token=valid_token, + team_object=team_object, + project_object=project_object, + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + # Run before apply_key_tags_pre_auth injects key metadata.tags into request_body. _reject_clientside_metadata_tags_check(general_settings, request_body, route) @@ -1004,6 +1061,13 @@ async def common_checks( proxy_logging_obj=proxy_logging_obj, valid_token=valid_token, ), + _model_access_group_max_budget_check( + matched_model_access_groups=matched_model_access_groups, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + if matched_model_access_groups + else None, _user_max_budget_check(), _check_team_member_budget( team_object=team_object, @@ -1444,6 +1508,7 @@ _REGISTRY_NOT_CACHED: Final = _RegistryNotCached() #: One lock per registry; module-level because the stampede to collapse is worker-wide. _TAG_REGISTRY_LOAD_LOCK: Final = asyncio.Lock() _END_USER_REGISTRY_LOAD_LOCK: Final = asyncio.Lock() +_MODEL_ACCESS_GROUP_REGISTRY_LOAD_LOCK: Final = asyncio.Lock() async def _cached_registry( @@ -1836,6 +1901,105 @@ async def _load_tag_registry( ) +async def _load_model_access_group_registry( + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, +) -> frozenset[str] | None: + """The set of model access group names that have a row in ``LiteLLM_ModelAccessGroupBudgetTable``.""" + + async def fetch_ids() -> tuple[str, ...]: + registry_rows: Final = await _model_access_group_budget_table( + ModelAccessGroupBudgetRepository(prisma_client) + ).find_many(take=MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE + 1) + return tuple(row.access_group_name for row in registry_rows) + + return await _load_bounded_registry( + cache_key=model_access_group_registry_cache_key(), + overflow_sentinel=MODEL_ACCESS_GROUP_REGISTRY_OVERFLOW_SENTINEL, + max_size=MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE, + load_lock=_MODEL_ACCESS_GROUP_REGISTRY_LOAD_LOCK, + fetch_ids=fetch_ids, + user_api_key_cache=user_api_key_cache, + ) + + +async def _fetch_uncached_model_access_group_budgets( + uncached_groups: Sequence[str], + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, +) -> tuple[tuple[str, ModelAccessGroupBudget], ...]: + """Budget rows for the groups a cache probe missed. + + No registry gate here, unlike the tag path: the names only ever come from + ``matched_model_access_groups``, which :func:`collect_matched_model_access_groups` already + intersected with the registry, so a name that has no row cannot reach this. + """ + if not uncached_groups: + return () + + try: + db_rows: Final = await _model_access_group_budget_table( + ModelAccessGroupBudgetRepository(prisma_client) + ).find_many( + where={"access_group_name": {"in": list(uncached_groups)}}, + include={"litellm_budget_table": True}, + ) + fetched: Final = tuple((row.access_group_name, _model_access_group_budget(row)) for row in db_rows) + for fetched_name, fetched_obj in fetched: + await user_api_key_cache.async_set_cache( + key=model_access_group_cache_key(fetched_name), + value=fetched_obj, + model_type=ModelAccessGroupBudget, + ttl=get_management_object_ttl(user_api_key_cache), + ) + except Exception as e: # noqa: BLE001 # fail-safe: a budget fetch error must yield "no budget rows", never break auth + verbose_proxy_logger.debug("Error batch fetching model access group budgets from database: %s", e) + return () + else: + return fetched + + +def _model_access_group_budget(row: _PrismaModelAccessGroupBudgetRow) -> ModelAccessGroupBudget: + budget_table: Final = row.litellm_budget_table + return ModelAccessGroupBudget( + access_group_name=row.access_group_name, + spend=row.spend or 0.0, + max_budget=None if budget_table is None else budget_table.max_budget, + ) + + +@log_db_metrics +async def get_model_access_group_budgets_batch( + access_group_names: Sequence[str], + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, +) -> dict[str, ModelAccessGroupBudget]: + """Budget rows for the given model access groups, served from cache where possible. + + Shared by the two enforcement paths so they read one row per group per request: the + reservation counters when reservations are on, and :func:`_model_access_group_max_budget_check` + when ``disable_budget_reservation`` turns them off. + """ + if prisma_client is None or not access_group_names: + return {} + + probed: Final = [ + ( + group, + await user_api_key_cache.async_get_cache( + key=model_access_group_cache_key(group), model_type=ModelAccessGroupBudget + ), + ) + for group in access_group_names + ] + fetched: Final = await _fetch_uncached_model_access_group_budgets( + uncached_groups=tuple(group for group, budget in probed if budget is None), + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + return {group: budget for group, budget in (*probed, *fetched) if budget is not None} + + async def _fetch_uncached_tags( uncached_tags: Sequence[str], prisma_client: PrismaClient, @@ -3882,6 +4046,192 @@ def _resolve_key_models_for_auth_check(valid_token: UserAPIKeyAuth) -> list[str] return models +def _model_access_groups_serving_model( + model: str | Sequence[str], + llm_router: Router, + team_id: str | None, +) -> frozenset[str]: + """Every model access group whose deployments serve the requested model(s).""" + requested: Final = (model,) if isinstance(model, str) else tuple(model) + return frozenset( + group + for requested_model in requested + for group in llm_router.get_model_access_groups(model_name=requested_model, team_id=team_id) + ) + + +async def _team_member_granted_models( + valid_token: UserAPIKeyAuth, + team_object: LiteLLM_TeamTable | None, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> Sequence[str]: + """The member's own ``allowed_models`` scope; empty when the member is not narrowed below the team.""" + if team_object is None or valid_token.user_id is None: + return () + + team_membership: Final = await get_team_membership( + user_id=valid_token.user_id, + team_id=team_object.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + return () if team_membership is None else _member_allowed_models(team_membership) + + +async def _org_granted_models( + valid_token: UserAPIKeyAuth, + team_object: LiteLLM_TeamTable | None, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> Sequence[str]: + """The org allowlist reached through the key, or through its team when the key names no org.""" + org_id: Final = valid_token.org_id or (team_object.organization_id if team_object is not None else None) + if org_id is None: + return () + + try: + org_object: Final = await get_org_object( + org_id=org_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception as e: # noqa: BLE001 # fail-safe: attribution degrades to "no org grant", it must never break auth + verbose_proxy_logger.debug("access group attribution: org lookup failed: %s", e) + return () + return org_object.models if org_object is not None else () + + +async def _granted_model_lists( + valid_token: UserAPIKeyAuth, + team_object: LiteLLM_TeamTable | None, + project_object: LiteLLM_ProjectTableCachedObj | None, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> tuple[Sequence[str], ...]: + """One model allowlist per level that participates in authorizing the request.""" + return ( + _resolve_key_models_for_auth_check(valid_token=valid_token), + team_object.models if team_object is not None else (), + await _team_member_granted_models( + valid_token=valid_token, + team_object=team_object, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ), + project_object.models if project_object is not None else (), + await _org_granted_models( + valid_token=valid_token, + team_object=team_object, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ), + ) + + +async def collect_matched_model_access_groups( + model: str | Sequence[str] | None, + valid_token: UserAPIKeyAuth | None, + team_object: LiteLLM_TeamTable | None, + project_object: LiteLLM_ProjectTableCachedObj | None, + llm_router: Router | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> tuple[str, ...]: + """ + The budgeted model access groups that authorized this request, sorted and deduplicated. + + A group is charged only when its name appears on an allowlist the caller was granted -- key, + team, team-member scope, project or org -- *and* that group serves the requested model. Asking + for a model that merely belongs to a group attributes nothing, because nothing about the caller + named the group. + + Levels are unioned, never ranked: a team granted ``*`` whose member is scoped to one group is + still a caller gated by that group. An unrestricted allowlist (empty, ``*``) names no group and + so contributes nothing. + + The whole walk is gated on the budget registry, because collecting every match costs a full scan + of each allowlist where the plain access check stops at the first hit. An empty registry means no + group carries a budget, so there is nothing to attribute and no work worth doing. + """ + if model is None or valid_token is None or llm_router is None or prisma_client is None: + return () + + registry: Final = await _load_model_access_group_registry( + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + if registry is not None and not registry: + return () + + covering_groups: Final = _model_access_groups_serving_model( + model=model, + llm_router=llm_router, + team_id=valid_token.team_id, + ) + budgeted_groups: Final = covering_groups if registry is None else covering_groups & registry + if not budgeted_groups: + return () + + granted: Final = frozenset( + granted_model + for granted_models in await _granted_model_lists( + valid_token=valid_token, + team_object=team_object, + project_object=project_object, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + for granted_model in granted_models + ) + return tuple(sorted(budgeted_groups & granted)) + + +async def stamp_matched_model_access_groups( + model: str | Sequence[str] | None, + valid_token: UserAPIKeyAuth | None, + team_object: LiteLLM_TeamTable | None, + project_object: LiteLLM_ProjectTableCachedObj | None, + llm_router: Router | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> tuple[str, ...]: + """Record the groups that authorized this request on its auth object, for the post-call spend + writer and the reservation counters, and hand them back for the budget check.""" + if valid_token is None: + return () + + try: + matched: Final = await collect_matched_model_access_groups( + model=model, + valid_token=valid_token, + team_object=team_object, + project_object=project_object, + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception as e: # noqa: BLE001 # fail-safe: attribution is spend telemetry, it must never break auth + verbose_proxy_logger.debug("model access group attribution failed: %s", e) + return () + if not matched: + return () + matched_groups: Final = list(matched) # mutable-ok: the auth field is typed list[str] | None + valid_token.matched_model_access_groups = matched_groups # rebind-ok: request-scoped carrier for the writer + return matched + + async def can_key_call_model( model: str | list[str], llm_model_list: list | None, @@ -5258,6 +5608,61 @@ async def _tag_max_budget_check( ) +async def _model_access_group_max_budget_check( + matched_model_access_groups: Sequence[str], + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, +) -> None: + """Block the request when a model access group that authorized it is over its max budget. + + Only the groups auth already matched are charged and therefore only they are checked, so a + request that no budgeted group authorized costs nothing here. + + Like the tag check this is a plain read with no reservation, so concurrent requests can + overshoot the ceiling slightly. The reservation counters are the precise path; this one covers + the ``disable_budget_reservation`` case. + + The ceiling is exclusive, unlike the tag check it otherwise mirrors: a pool whose recorded + spend has reached ``max_budget`` has nothing left to give, so the next request is refused. + Keys and organizations already draw the line there. A non-positive budget means no budget, + matching what the reservation path treats as unbudgeted. + + Raises: + BudgetExceededError if a matched group is over its max budget. + """ + if prisma_client is None or not matched_model_access_groups: + return + + budgets: Final = await get_model_access_group_budgets_batch( + access_group_names=matched_model_access_groups, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + + from litellm.proxy.proxy_server import get_current_spend + + for group in matched_model_access_groups: + budget = budgets.get(group) + if budget is None or budget.max_budget is None or budget.max_budget <= 0: + continue + + group_spend = await get_current_spend( + counter_key=model_access_group_spend_counter_key(group), + fallback_spend=budget.spend, + max_budget=budget.max_budget, + fallback_authoritative=True, + ) + if group_spend < budget.max_budget: + continue + raise litellm.BudgetExceededError( + current_cost=group_spend, + max_budget=budget.max_budget, + message=f"Budget has been exceeded! Model access group={group} Current cost: {group_spend}, Max budget: {budget.max_budget}", + entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP.value, + entity_id=group, + ) + + def is_model_allowed_by_pattern(model: str, allowed_model_pattern: str) -> bool: """ Check if a model matches an allowed pattern. diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 2bf5b7b9283..1682cf12f4e 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -35,7 +35,11 @@ from litellm.proxy.common_utils.timezone_utils import ( compute_budget_reset_at, get_budget_reset_settings, ) -from litellm.proxy.common_utils.user_api_key_cache import tag_cache_key +from litellm.proxy.common_utils.user_api_key_cache import ( + model_access_group_cache_key, + model_access_group_spend_counter_key, + tag_cache_key, +) from litellm.proxy.db.budget_window_spend_writer import roll_window_spend_row from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry @@ -44,6 +48,7 @@ from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.prisma_protocols import SpendLinkedTable from litellm.repositories.table_repositories import ( EndUserRepository, + ModelAccessGroupBudgetRepository, TagRepository, TeamMembershipRepository, ) @@ -95,6 +100,11 @@ class _TagRow(_BudgetLinkedRow, Protocol): def tag_name(self) -> str: ... +class _ModelAccessGroupRow(_BudgetLinkedRow, Protocol): + @property + def access_group_name(self) -> str: ... + + class _EndUserRow(_BudgetLinkedRow, Protocol): @property def user_id(self) -> str: ... @@ -157,6 +167,14 @@ def _tag_cache_keys(row: _TagRow) -> tuple[str, ...]: return (tag_cache_key(row.tag_name),) +def _model_access_group_counter_key(row: _ModelAccessGroupRow) -> str: + return model_access_group_spend_counter_key(row.access_group_name) + + +def _model_access_group_cache_keys(row: _ModelAccessGroupRow) -> tuple[str, ...]: + return (model_access_group_cache_key(row.access_group_name),) + + def _budget_link_where( budget_ids: Sequence[str], extra: Mapping[str, object] = MappingProxyType({}), @@ -616,6 +634,11 @@ class ResetBudgetJob: where=_budget_link_where(budget_ids, _SPENT_ROWS_WHERE), log_subject="tags", ) + model_access_groups: Final[tuple[_ModelAccessGroupRow, ...]] = await self._fetch_linked_rows( + table=ModelAccessGroupBudgetRepository(self.prisma_client).table, + where=_budget_link_where(budget_ids, _SPENT_ROWS_WHERE), + log_subject="model access groups", + ) rollover_caps: Final[Mapping[str, float]] = MappingProxyType( { # mutable-ok: MappingProxyType wraps a one-shot dict comprehension b.budget_id: cap @@ -645,6 +668,10 @@ class ResetBudgetJob: *((_key_counter_key(row), _row_carried_spend(row, rollover_caps)) for row in keys), *((_org_counter_key(row), _row_carried_spend(row, rollover_caps)) for row in orgs), *((_tag_counter_key(row), _row_carried_spend(row, rollover_caps)) for row in tags), + *( + (_model_access_group_counter_key(row), _row_carried_spend(row, rollover_caps)) + for row in model_access_groups + ), ), rollover_caps=rollover_caps, cache_keys=( @@ -652,6 +679,7 @@ class ResetBudgetJob: *(key for row in keys for key in _key_cache_keys(row)), *(key for row in orgs for key in _org_cache_keys(row)), *(key for row in tags for key in _tag_cache_keys(row)), + *(key for row in model_access_groups for key in _model_access_group_cache_keys(row)), ), ) @@ -677,6 +705,7 @@ class ResetBudgetJob: _queue_budget_linked_resets(uow.keys, cascade, extra=_LINKED_KEYS_WHERE) _queue_budget_linked_resets(uow.organizations, cascade, extra=_SPENT_ROWS_WHERE) _queue_budget_linked_resets(uow.tags, cascade, extra=_SPENT_ROWS_WHERE) + _queue_budget_linked_resets(uow.model_access_groups, cascade, extra=_SPENT_ROWS_WHERE) _queue_enduser_resets(uow.endusers, cascade) for budget_id, budget_reset_at in cascade.budget_resets: uow.budgets.queue_window_advance(budget_id=budget_id, budget_reset_at=budget_reset_at) @@ -720,7 +749,8 @@ class ResetBudgetJob: async def reset_budget_for_litellm_budget_table(self) -> None: """ Resets the spend a budget tier gates (end users, team members, keys, - orgs, tags) and advances the tier's budget_reset_at, atomically. + orgs, tags, model access groups) and advances the tier's + budget_reset_at, atomically. Caches are invalidated only after the transaction commits, so a failed run cannot leave a zeroed counter in front of an un-reset DB row. @@ -751,8 +781,9 @@ class ResetBudgetJob: return _ChunkOutcome(fetched=len(cascade.budgets), advanced=advanced) case _BudgetCascadeFailed(cascade=cascade, error=error): verbose_proxy_logger.exception( - "Failed to reset the budget table cascade (team member, enduser, org and tag spend, plus " - "budget_reset_at); nothing was committed and the budgets stay due for the next run: %s", + "Failed to reset the budget table cascade (team member, enduser, org, tag and model access " + "group spend, plus budget_reset_at); nothing was committed and the budgets stay due for the " + "next run: %s", error, exc_info=error, ) diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index b8df0105b7b..589d8fe68d1 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -185,6 +185,32 @@ def tag_registry_cache_key() -> str: return "tag_registry" +#: Cached under ``model_access_group_registry_cache_key`` when the table exceeds +#: ``MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE``: registry unusable, fall back to the per-group lookup. +MODEL_ACCESS_GROUP_REGISTRY_OVERFLOW_SENTINEL: Final = "__model_access_group_registry_overflow__" + + +def model_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"model_access_group:{access_group_name}" + + +def model_access_group_registry_cache_key() -> str: + """Cache key for the set of model access group names that have a budget row.""" + return "model_access_group_registry" + + +def model_access_group_spend_counter_key(access_group_name: str) -> str: + """Spend counter key for one model access group; shared so its four owners cannot drift. + + The reservation path writes it up front, the cost callback writes it after the call, auth + reads it to enforce ``max_budget``, and the reset job clears it on rollover. A copy that + drifts in any one of them silently resets or reads a counter nobody else touches, which shows + up as a budget that never trips or never resets. + """ + return f"spend:model_access_group:{access_group_name}" + + #: 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/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 1ce7a959fec..641c07914d9 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -12,7 +12,7 @@ import os import random import time import traceback -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, overload @@ -24,6 +24,7 @@ from litellm.constants import ( DB_SPEND_UPDATE_JOB_NAME, INTERNAL_CALL_ORIGIN_METADATA_KEY, ) +from litellm.litellm_core_utils.litellm_logging import coerce_model_access_groups from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.proxy._types import ( DB_RETRY_SAFE_ERROR_TYPES, @@ -91,6 +92,7 @@ class _SpendBatch(Protocol): litellm_organizationtable: BatchTable litellm_tagtable: BatchTable litellm_agentstable: BatchTable + litellm_modelaccessgroupbudgettable: BatchTable class _SpendBatchManager(Protocol): @@ -114,7 +116,7 @@ def _spend_update_tx(prisma_client: PrismaClient) -> _SpendTransactionManager: return tx -def _get_llm_router(): +def get_llm_router(): """The proxy's router, or None outside a running proxy. Injected rather than imported where it is used, so the savings computation stays @@ -128,6 +130,52 @@ def _get_llm_router(): return None +class _DeploymentLookup(Protocol): + def get_model_info(self, id: str) -> Mapping[str, object] | None: ... + + +def _served_model_access_groups( + router: _DeploymentLookup | None, + served_model_id: str | None, +) -> frozenset[str] | None: + """Access groups declared by the deployment that actually served the request. + + None when the served deployment cannot be identified, in which case the set + attributed at auth time stands unchanged. + """ + if router is None or not served_model_id: + return None + deployment: Final = router.get_model_info(id=served_model_id) + if deployment is None: + return None + model_info: Final = deployment.get("model_info") + if not isinstance(model_info, Mapping): + return None + declared: Final = model_info.get("access_groups") + if not isinstance(declared, (list, tuple)): + return frozenset() + return frozenset(group for group in declared if isinstance(group, str)) + + +def debitable_model_access_groups( + attributed: Sequence[str] | None, + served_model_id: str | None, + router: _DeploymentLookup | None, +) -> tuple[str, ...]: + """Groups to debit: the set attributed at auth time, narrowed to those the served model belongs to. + + The router may fall back to a model outside the pool auth reserved against, so the + attributed set is the hard upper bound: a group absent from it is never debited. + """ + ordered: Final = coerce_model_access_groups(attributed) + if not ordered: + return () + served: Final = _served_model_access_groups(router=router, served_model_id=served_model_id) + if served is None: + return ordered + return tuple(group for group in ordered if group in served) + + class DBSpendUpdateWriter: """ Module responsible for @@ -197,6 +245,7 @@ class DBSpendUpdateWriter: ## CREATE SPEND LOG PAYLOAD ## from litellm.proxy.spend_tracking.spend_tracking_utils import ( get_logging_payload, + get_request_model_access_groups, ) payload: Final = get_logging_payload( @@ -249,6 +298,7 @@ class DBSpendUpdateWriter: prisma_client=prisma_client, litellm_proxy_budget_name=litellm_proxy_budget_name, payload=payload, + request_model_access_groups=get_request_model_access_groups(kwargs), ) ) @@ -332,7 +382,7 @@ class DBSpendUpdateWriter: routing_decision=metadata.get("routing_decision"), usage_object=usage_object_raw if isinstance(usage_object_raw, dict) else None, model_id=payload.get("model_id"), - llm_router=_get_llm_router, + llm_router=get_llm_router, cost_breakdown=metadata.get("cost_breakdown"), recorded_autorouter_savings=metadata.get("autorouter_savings"), ) @@ -443,9 +493,10 @@ class DBSpendUpdateWriter: prisma_client: PrismaClient | None, litellm_proxy_budget_name: str | None, payload: SpendLogsPayload, + request_model_access_groups: Sequence[str] = (), ): """ - Runs all 11 spend-update helpers sequentially inside a single asyncio task. + Runs all 13 spend-update helpers sequentially inside a single asyncio task. Each helper is wrapped in try/except so one failure doesn't prevent the others. @@ -517,6 +568,14 @@ class DBSpendUpdateWriter: traceback.format_exc(), ) + await self._update_model_access_group_db( + response_cost=response_cost, + request_model_access_groups=request_model_access_groups, + served_model_id=payload_copy.get("model_id"), + prisma_client=prisma_client, + router=get_llm_router(), + ) + _agent_id_for_spend: Final = payload_copy.get("agent_id") try: await self._update_agent_db( @@ -826,6 +885,50 @@ class DBSpendUpdateWriter: ) raise e + async def _update_model_access_group_db( + self, + response_cost: float | None, + request_model_access_groups: Sequence[str] | None, + served_model_id: str | None, + prisma_client: PrismaClient | None, + router: _DeploymentLookup | None = None, + ) -> None: + """ + Update spend for every model access group this request is billed against. + + Args: + response_cost: Cost of the request, charged in full to each group + request_model_access_groups: Groups attributed at auth time, the upper bound on what may be debited + served_model_id: Deployment id actually served, used to narrow the attributed set + prisma_client: Prisma client instance + router: Deployment lookup used to re-resolve groups after a fallback + """ + try: + if prisma_client is None: + return + + for model_access_group in debitable_model_access_groups( + attributed=request_model_access_groups, + served_model_id=served_model_id, + router=router, + ): + await self.spend_update_queue.add_update( + update=SpendUpdateQueueItem( + entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP, + entity_id=model_access_group, + response_cost=response_cost, + ) + ) + except Exception as e: # noqa: BLE001 # isolation: a helper failure must not stop the batch + spend_log_error( + "Spend tracking - failed to enqueue model access group spend update. " + "model_access_groups=%s, response_cost=%s - %s", + request_model_access_groups, + response_cost, + str(e), + exc=e, + ) + async def _insert_spend_log_to_db( self, payload: dict | SpendLogsPayload, @@ -942,7 +1045,8 @@ class DBSpendUpdateWriter: if db_spend_update_transactions is not None: verbose_proxy_logger.info( "Spend tracking - committing spend updates from Redis to DB: " - "keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, tags=%d, agents=%d", + "keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, tags=%d, agents=%d, " + "model_access_groups=%d", len(db_spend_update_transactions.get("key_list_transactions") or {}), len(db_spend_update_transactions.get("user_list_transactions") or {}), len(db_spend_update_transactions.get("team_list_transactions") or {}), @@ -951,6 +1055,7 @@ class DBSpendUpdateWriter: len(db_spend_update_transactions.get("team_member_list_transactions") or {}), len(db_spend_update_transactions.get("tag_list_transactions") or {}), len(db_spend_update_transactions.get("agent_list_transactions") or {}), + len(db_spend_update_transactions.get("model_access_group_list_transactions") or {}), ) await self._commit_spend_updates_to_db( prisma_client=prisma_client, @@ -1497,6 +1602,20 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, ) + ### UPDATE MODEL ACCESS GROUP TABLE ### + model_access_group_list_transactions: Final = db_spend_update_transactions.get( + "model_access_group_list_transactions" + ) + await DBSpendUpdateWriter._update_entity_spend_in_db( + entity_name="Model access group", + transactions=model_access_group_list_transactions, + table_accessor="litellm_modelaccessgroupbudgettable", + where_field="access_group_name", + n_retry_times=n_retry_times, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + ) + ### UPDATE AGENT TABLE ### agent_list_transactions: Final = db_spend_update_transactions["agent_list_transactions"] await DBSpendUpdateWriter._update_entity_spend_in_db( @@ -1513,7 +1632,7 @@ class DBSpendUpdateWriter: async def _update_entity_spend_in_db( entity_name: str, transactions: dict[str, float] | None, - table_accessor: Literal["litellm_tagtable", "litellm_agentstable"], + table_accessor: Literal["litellm_tagtable", "litellm_agentstable", "litellm_modelaccessgroupbudgettable"], where_field: str, n_retry_times: int, prisma_client: PrismaClient, @@ -1948,7 +2067,7 @@ class DBSpendUpdateWriter: gateway_injected_cache=marks_gateway_injection(_metadata, payload.get("model_id")), routing_decision=_metadata.get("routing_decision"), model_id=payload.get("model_id"), - llm_router=_get_llm_router, + llm_router=get_llm_router, usage_object=usage_obj, cost_breakdown=_metadata.get("cost_breakdown"), recorded_autorouter_savings=_metadata.get("autorouter_savings"), diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index fa096fa0bf2..4dd23270bf8 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -70,6 +70,7 @@ _SpendTransactionField: TypeAlias = Literal[ "org_list_transactions", "tag_list_transactions", "agent_list_transactions", + "model_access_group_list_transactions", ] _SPEND_TRANSACTION_FIELDS: Final[tuple[_SpendTransactionField, ...]] = ( @@ -81,6 +82,7 @@ _SPEND_TRANSACTION_FIELDS: Final[tuple[_SpendTransactionField, ...]] = ( "org_list_transactions", "tag_list_transactions", "agent_list_transactions", + "model_access_group_list_transactions", ) _ValueT = TypeVar("_ValueT") @@ -417,6 +419,10 @@ class RedisUpdateBuffer: Litellm_EntityType.AGENT, db_spend_update_transactions.get("agent_list_transactions"), ), + ( + Litellm_EntityType.MODEL_ACCESS_GROUP, + db_spend_update_transactions.get("model_access_group_list_transactions"), + ), ] for entity_type, entities in entity_entries: if not entities: @@ -866,6 +872,9 @@ class RedisUpdateBuffer: org_list_transactions=_merged_entity_transactions(list_of_transactions, "org_list_transactions"), tag_list_transactions=_merged_entity_transactions(list_of_transactions, "tag_list_transactions"), agent_list_transactions=_merged_entity_transactions(list_of_transactions, "agent_list_transactions"), + model_access_group_list_transactions=_merged_entity_transactions( + list_of_transactions, "model_access_group_list_transactions" + ), ) async def _emit_new_item_added_to_redis_buffer_event( diff --git a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py index 57cb5e73b64..8c0076b10c1 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py @@ -139,6 +139,7 @@ class SpendUpdateQueue(BaseUpdateQueue): org_list_transactions={}, tag_list_transactions={}, agent_list_transactions={}, + model_access_group_list_transactions={}, ) # Map entity types to their corresponding transaction dictionary keys @@ -151,6 +152,7 @@ class SpendUpdateQueue(BaseUpdateQueue): Litellm_EntityType.ORGANIZATION: "org_list_transactions", Litellm_EntityType.TAG: "tag_list_transactions", Litellm_EntityType.AGENT: "agent_list_transactions", + Litellm_EntityType.MODEL_ACCESS_GROUP: "model_access_group_list_transactions", } for update in updates: @@ -190,6 +192,8 @@ class SpendUpdateQueue(BaseUpdateQueue): transactions_dict = db_spend_update_transactions["tag_list_transactions"] elif dict_key == "agent_list_transactions": transactions_dict = db_spend_update_transactions["agent_list_transactions"] + elif dict_key == "model_access_group_list_transactions": + transactions_dict = db_spend_update_transactions["model_access_group_list_transactions"] else: continue diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index a8ffb650ff0..f02901f0e97 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -1,5 +1,6 @@ import asyncio import traceback +from collections.abc import Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final, cast @@ -20,6 +21,10 @@ from litellm.proxy.auth.auth_checks import ( log_db_metrics, ) from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.db.db_spend_update_writer import ( + debitable_model_access_groups, + get_llm_router, +) from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.spend_tracking.spend_log_error_logger import ( should_suppress_spend_log_tracebacks, @@ -27,6 +32,7 @@ from litellm.proxy.spend_tracking.spend_log_error_logger import ( ) from litellm.proxy.spend_tracking.spend_tracking_utils import ( _sanitize_error_information_for_spend_logs, + get_request_model_access_groups, ) from litellm.proxy.utils import ProxyUpdateSpend from litellm.types.utils import ( @@ -258,6 +264,11 @@ class _ProxyDBLogger(CustomLogger): sl_object=sl_object, metadata=metadata, ) + model_access_groups: Final = debitable_model_access_groups( + attributed=get_request_model_access_groups(kwargs), + served_model_id=sl_object.get("model_id") if sl_object is not None else None, + router=get_llm_router(), + ) if response_cost is not None: user_api_key: Final = metadata.get("user_api_key", None) @@ -296,6 +307,7 @@ class _ProxyDBLogger(CustomLogger): response_cost=response_cost, budget_reservation=budget_reservation, request_tags=tags, + model_access_groups=model_access_groups, ) # update cache (fire-and-forget for backward compat: @@ -572,6 +584,7 @@ async def _update_database_and_spend_counters( response_cost: float, budget_reservation: dict | None, request_tags: list[str] | None = None, + model_access_groups: Sequence[str] | None = None, ) -> None: try: spend_log_request_id = await proxy_logging_obj.db_spend_update_writer.update_database( @@ -612,6 +625,7 @@ async def _update_database_and_spend_counters( tags=request_tags, request_id=spend_log_request_id, request_started_at=start_time, + model_access_groups=model_access_groups, ) except Exception: if budget_reservation is not None: diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 222ba45b059..ae55b7ab906 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -29,6 +29,7 @@ from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( _request_blocked_callback_params, iter_client_callback_metadata_dicts, ) +from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.url_utils import ( is_url_destination_allowed_by_host, @@ -1378,6 +1379,10 @@ class LiteLLMProxyRequestSetup: ) if user_api_key_dict.budget_reservation is not None: data[_metadata_variable_name]["user_api_key_budget_reservation"] = user_api_key_dict.budget_reservation + if user_api_key_dict.matched_model_access_groups: + data[_metadata_variable_name][MODEL_ACCESS_GROUP_METADATA_KEY] = ( + user_api_key_dict.matched_model_access_groups + ) # UserAPIKeyAuth object for MCP server access control data[_metadata_variable_name]["user_api_key_auth"] = user_api_key_dict.model_copy( update={ diff --git a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py index e1a7645e988..dabb9334b16 100644 --- a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py @@ -2,18 +2,33 @@ Allow proxy admin to manage model access groups Endpoints here: -- POST /model_group/new - Create a new access group with multiple model names +- POST /access_group/new - Create a new access group with multiple model names +- GET /access_group/list - List every access group +- GET /access_group/{access_group}/info - Read one access group, including its budget +- PUT /access_group/{access_group}/update - Replace an access group's deployments +- DELETE /access_group/{access_group}/delete - Delete an access group and its budget +- GET /access_group/{access_group}/budget - Read an access group's shared budget and spend +- PUT /access_group/{access_group}/budget - Set or replace an access group's shared budget +- DELETE /access_group/{access_group}/budget - Clear an access group's shared budget """ import json from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Protocol +from datetime import datetime +from typing import TYPE_CHECKING, Annotated, Any, Final, Protocol from fastapi import APIRouter, Depends, HTTPException +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.user_api_key_cache import ( + UserApiKeyCache, + model_access_group_cache_key, + model_access_group_registry_cache_key, +) +from litellm.proxy.management_endpoints.common_utils import validate_budget_duration # Clear cache and reload models to pick up the access group changes from litellm.proxy.management_endpoints.model_management_endpoints import ( @@ -22,10 +37,16 @@ from litellm.proxy.management_endpoints.model_management_endpoints import ( model_info_as_mapping, reload_serving_verdict, ) +from litellm.proxy.management_helpers.utils import handle_budget_for_entity from litellm.proxy.utils import PrismaClient from litellm.repositories.model_repository import ModelRepository +from litellm.repositories.table_repositories import ModelAccessGroupBudgetRepository from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + AccessGroupBudget, + AccessGroupBudgetRequest, + AccessGroupBudgetResponse, AccessGroupInfo, + DeleteAccessGroupBudgetResponse, DeleteModelGroupResponse, ListAccessGroupsResponse, NewModelGroupRequest, @@ -36,7 +57,43 @@ from litellm.types.proxy.management_endpoints.model_management_endpoints import if TYPE_CHECKING: from litellm import Router -router: Final = APIRouter() +router: Final = APIRouter(tags=["model management"]) + +_AUTH_DEPENDENCIES: Final = (Depends(user_api_key_auth),) + + +class _ErrorDetail(TypedDict): + error: ReadOnly[str] + + +class _ModelAccessGroupWhere(TypedDict): + access_group_name: ReadOnly[str] + + +class _BudgetInclude(TypedDict): + litellm_budget_table: ReadOnly[bool] + + +class _ModelAccessGroupBudgetCreate(TypedDict): + access_group_name: ReadOnly[str] + budget_id: ReadOnly[str | None] + created_by: ReadOnly[str] + updated_by: ReadOnly[str] + + +class _ModelAccessGroupBudgetUpdate(TypedDict): + budget_id: ReadOnly[str | None] + updated_by: ReadOnly[str] + + +class _ModelAccessGroupBudgetUpsert(TypedDict): + create: ReadOnly[_ModelAccessGroupBudgetCreate] + update: ReadOnly[_ModelAccessGroupBudgetUpdate] + + +def _http_error(status_code: int, message: str) -> HTTPException: + detail: Final[_ErrorDetail] = {"error": message} + return HTTPException(status_code=status_code, detail=detail) class _DeploymentRow(Protocol): @@ -58,10 +115,140 @@ class _ModelTableClient(Protocol): async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> object: ... +class _BudgetRow(Protocol): + @property + def budget_id(self) -> str: ... + + @property + def max_budget(self) -> float | None: ... + + @property + def soft_budget(self) -> float | None: ... + + @property + def budget_duration(self) -> str | None: ... + + @property + def budget_reset_at(self) -> datetime | None: ... + + +class _ModelAccessGroupBudgetRow(Protocol): + @property + def spend(self) -> float: ... + + @property + def budget_id(self) -> str | None: ... + + @property + def litellm_budget_table(self) -> _BudgetRow | None: ... + + +class _ModelAccessGroupBudgetTableClient(Protocol): + async def find_unique( + self, *, where: Mapping[str, object], include: Mapping[str, object] | None = None + ) -> _ModelAccessGroupBudgetRow | None: ... + + async def upsert( + self, + *, + where: Mapping[str, object], + data: Mapping[str, object], + include: Mapping[str, object] | None = None, + ) -> _ModelAccessGroupBudgetRow: ... + + async def delete(self, *, where: Mapping[str, object]) -> _ModelAccessGroupBudgetRow | None: ... + + def _model_table(prisma_client: PrismaClient) -> _ModelTableClient: return ModelRepository(prisma_client).table +def _model_access_group_budget_table(prisma_client: PrismaClient) -> _ModelAccessGroupBudgetTableClient: + return ModelAccessGroupBudgetRepository(prisma_client).table + + +def _prisma_client_or_500() -> PrismaClient: + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise _http_error(500, "Database not connected.") + return prisma_client + + +def _auth_cache() -> UserApiKeyCache: + from litellm.proxy.proxy_server import user_api_key_cache + + return user_api_key_cache + + +async def _evict_model_access_group_cache_keys(access_group: str, auth_cache: UserApiKeyCache) -> None: + """ + Every endpoint that writes an access group budget row must call this, or the budget stays + unenforced until the TTL expires: auth gates the feature on a cached registry of the groups + that have a budget row, read cache-first with no freshness check. + """ + from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( + evict_and_broadcast, + ) + + await evict_and_broadcast( + cache_keys=(model_access_group_cache_key(access_group), model_access_group_registry_cache_key()), + user_api_key_cache=auth_cache, + ) + + +async def _model_access_group_budget_row( + access_group: str, prisma_client: PrismaClient +) -> _ModelAccessGroupBudgetRow | None: + where: Final[_ModelAccessGroupWhere] = {"access_group_name": access_group} + include: Final[_BudgetInclude] = {"litellm_budget_table": True} + return await _model_access_group_budget_table(prisma_client).find_unique(where=where, include=include) + + +def _budget_or_none(row: _ModelAccessGroupBudgetRow | None) -> AccessGroupBudget | None: + budget: Final = row.litellm_budget_table if row is not None else None + if budget is None: + return None + return AccessGroupBudget( + budget_id=budget.budget_id, + max_budget=budget.max_budget, + soft_budget=budget.soft_budget, + budget_duration=budget.budget_duration, + budget_reset_at=budget.budget_reset_at, + ) + + +def _budget_response(access_group: str, row: _ModelAccessGroupBudgetRow | None) -> AccessGroupBudgetResponse: + return AccessGroupBudgetResponse( + access_group=access_group, + spend=row.spend if row is not None else 0.0, + budget=_budget_or_none(row), + ) + + +async def _delete_model_access_group_budget_row( + access_group: str, prisma_client: PrismaClient, auth_cache: UserApiKeyCache +) -> bool: + """ + Drop the group's budget row only, matching /tag/delete: the LiteLLM_BudgetTable row survives + because the link is ON DELETE SET NULL and a budget_id an admin passed in may be shared with + other entities. + + Evicts unconditionally: a group with no row of its own can still be sitting in the cached + registry, so skipping the eviction when nothing was deleted would leave that stale. + """ + where: Final[_ModelAccessGroupWhere] = {"access_group_name": access_group} + row: Final = await _model_access_group_budget_table(prisma_client).delete(where=where) + await _evict_model_access_group_cache_keys(access_group, auth_cache) + return row is not None + + +async def _raise_404_if_model_access_group_missing(access_group: str, prisma_client: PrismaClient) -> None: + access_groups_map: Final = await get_all_access_groups_from_db(prisma_client=prisma_client) + if access_group not in access_groups_map: + raise _http_error(404, f"Access group '{access_group}' not found") + + def validate_models_exist(model_names: Sequence[str], llm_router: "Router | None") -> tuple[bool, Sequence[str]]: """ Validate that all requested model names exist in the router. @@ -356,13 +543,12 @@ async def get_all_access_groups_from_db( @router.post( "/access_group/new", - tags=["model management"], - dependencies=[Depends(user_api_key_auth)], + dependencies=_AUTH_DEPENDENCIES, response_model=NewModelGroupResponse, ) async def create_model_group( data: NewModelGroupRequest, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], ): """ Create a new access group containing multiple model names. @@ -503,12 +689,11 @@ async def create_model_group( @router.get( "/access_group/list", - tags=["model management"], - dependencies=[Depends(user_api_key_auth)], + dependencies=_AUTH_DEPENDENCIES, response_model=ListAccessGroupsResponse, ) async def list_access_groups( - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], ): """ List all access groups. @@ -553,13 +738,12 @@ async def list_access_groups( @router.get( "/access_group/{access_group}/info", - tags=["model management"], - dependencies=[Depends(user_api_key_auth)], + dependencies=_AUTH_DEPENDENCIES, response_model=AccessGroupInfo, ) async def get_access_group_info( access_group: str, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], ): """ Get information about a specific access group. @@ -574,7 +758,7 @@ async def get_access_group_info( - access_group: str - The access group name (URL path parameter) Returns: - - AccessGroupInfo with the access group details + - AccessGroupInfo with the access group details, its shared budget and its spend Raises: - HTTPException 404: If access group not found @@ -596,7 +780,15 @@ async def get_access_group_info( detail={"error": f"Access group '{access_group}' not found"}, ) - return access_groups_map[access_group] + info: Final = access_groups_map[access_group] + budget_row: Final = await _model_access_group_budget_row(access_group, prisma_client) + return AccessGroupInfo( + access_group=info.access_group, + model_names=info.model_names, + deployment_count=info.deployment_count, + spend=budget_row.spend if budget_row is not None else 0.0, + budget=_budget_or_none(budget_row), + ) except HTTPException: raise @@ -610,14 +802,13 @@ async def get_access_group_info( @router.put( "/access_group/{access_group}/update", - tags=["model management"], - dependencies=[Depends(user_api_key_auth)], + dependencies=_AUTH_DEPENDENCIES, response_model=NewModelGroupResponse, ) async def update_access_group( access_group: str, data: UpdateModelGroupRequest, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], ): """ Update an access group's model names. @@ -765,13 +956,13 @@ async def update_access_group( @router.delete( "/access_group/{access_group}/delete", - tags=["model management"], - dependencies=[Depends(user_api_key_auth)], + dependencies=_AUTH_DEPENDENCIES, response_model=DeleteModelGroupResponse, ) async def delete_access_group( access_group: str, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + auth_cache: Annotated[UserApiKeyCache, Depends(_auth_cache)], ): """ Delete an access group. @@ -835,6 +1026,13 @@ async def delete_access_group( removed_pairs: Final = tuple(pair for pair in removed if pair is not None) models_updated: Final = len(removed_pairs) + # Budget last, deliberately: failing here strands a budget row for a group already on no + # deployment (clutter), where the reverse order can leave a live group enforcing nothing. + # The LiteLLM_BudgetTable row it linked is left alone, as /tag/delete leaves a tag's. + await _delete_model_access_group_budget_row( + access_group=access_group, prisma_client=prisma_client, auth_cache=auth_cache + ) + # Clear cache and reload models to pick up the access group changes live_before_reload: Final = live_model_ids_snapshot() reload_outcome: Final = await clear_cache() @@ -864,3 +1062,162 @@ async def delete_access_group( status_code=500, detail={"error": f"Failed to delete access group: {e}"}, ) + + +@router.get( + "/access_group/{access_group}/budget", + dependencies=_AUTH_DEPENDENCIES, + response_model=AccessGroupBudgetResponse, +) +async def get_access_group_budget( + access_group: str, +) -> AccessGroupBudgetResponse: + """ + Get the shared budget of an access group, and the spend drawn against it. + + Example: + ```bash + curl -X GET 'http://localhost:4000/access_group/production-models/budget' \\ + -H 'Authorization: Bearer sk-1234' + ``` + + Parameters: + - access_group: str - The access group name (URL path parameter) + + Returns: + - AccessGroupBudgetResponse; budget is null when the group has no budget set + + Raises: + - HTTPException 404: If access group not found + """ + prisma_client: Final = _prisma_client_or_500() + await _raise_404_if_model_access_group_missing(access_group=access_group, prisma_client=prisma_client) + + return _budget_response( + access_group=access_group, + row=await _model_access_group_budget_row(access_group, prisma_client), + ) + + +@router.put( + "/access_group/{access_group}/budget", + dependencies=_AUTH_DEPENDENCIES, + response_model=AccessGroupBudgetResponse, +) +async def set_access_group_budget( + access_group: str, + data: AccessGroupBudgetRequest, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + auth_cache: Annotated[UserApiKeyCache, Depends(_auth_cache)], +) -> AccessGroupBudgetResponse: + """ + Set or replace the shared budget of an access group. Idempotent. + + Every key that can reach a model in the group draws from this one budget. + + Example: + ```bash + curl -X PUT 'http://localhost:4000/access_group/production-models/budget' \\ + -H 'Authorization: Bearer sk-1234' \\ + -H 'Content-Type: application/json' \\ + -d '{ + "max_budget": 100.0, + "budget_duration": "30d" + }' + ``` + + Parameters: + - access_group: str - The access group name (URL path parameter) + - max_budget: Optional[float] - Requests fail once the group's shared spend exceeds this + - soft_budget: Optional[float] - Fires an alert when reached; requests still succeed + - budget_duration: Optional[str] - Frequency of resetting the group's spend (e.g. '30d') + - budget_id: Optional[str] - Link an existing budget instead of creating one + + Returns: + - AccessGroupBudgetResponse with the stored budget and current spend + + Raises: + - HTTPException 400: If no budget field is given, or budget_duration cannot be parsed + - HTTPException 404: If access group not found + """ + from litellm.proxy.proxy_server import litellm_proxy_admin_name + + prisma_client: Final = _prisma_client_or_500() + if not data.model_dump(exclude_none=True): + raise _http_error(400, "One of max_budget, soft_budget, budget_duration or budget_id is required") + validate_budget_duration(data.budget_duration) + await _raise_404_if_model_access_group_missing(access_group=access_group, prisma_client=prisma_client) + + existing_row: Final = await _model_access_group_budget_row(access_group, prisma_client) + budget_id: Final = await handle_budget_for_entity( + data=data, + existing_budget_id=existing_row.budget_id if existing_row is not None else None, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ) + actor: Final = user_api_key_dict.user_id or litellm_proxy_admin_name + upsert_data: Final[_ModelAccessGroupBudgetUpsert] = { + "create": { + "access_group_name": access_group, + "budget_id": budget_id, + "created_by": actor, + "updated_by": actor, + }, + "update": {"budget_id": budget_id, "updated_by": actor}, + } + where: Final[_ModelAccessGroupWhere] = {"access_group_name": access_group} + include: Final[_BudgetInclude] = {"litellm_budget_table": True} + row: Final = await _model_access_group_budget_table(prisma_client).upsert( + where=where, data=upsert_data, include=include + ) + await _evict_model_access_group_cache_keys(access_group, auth_cache) + + verbose_proxy_logger.info("Set budget %s on access group '%s'", budget_id, access_group) + return _budget_response(access_group=access_group, row=row) + + +@router.delete( + "/access_group/{access_group}/budget", + dependencies=_AUTH_DEPENDENCIES, + response_model=DeleteAccessGroupBudgetResponse, +) +async def delete_access_group_budget( + access_group: str, + auth_cache: Annotated[UserApiKeyCache, Depends(_auth_cache)], +) -> DeleteAccessGroupBudgetResponse: + """ + Clear the shared budget of an access group, leaving the group itself in place. + + Example: + ```bash + curl -X DELETE 'http://localhost:4000/access_group/production-models/budget' \\ + -H 'Authorization: Bearer sk-1234' + ``` + + Parameters: + - access_group: str - The access group name (URL path parameter) + + Returns: + - DeleteAccessGroupBudgetResponse; budget_deleted is false when there was nothing to clear + + Raises: + - HTTPException 404: If access group not found + """ + prisma_client: Final = _prisma_client_or_500() + await _raise_404_if_model_access_group_missing(access_group=access_group, prisma_client=prisma_client) + + budget_deleted: Final = await _delete_model_access_group_budget_row( + access_group=access_group, + prisma_client=prisma_client, + auth_cache=auth_cache, + ) + return DeleteAccessGroupBudgetResponse( + access_group=access_group, + budget_deleted=budget_deleted, + message=( + f"Budget for access group '{access_group}' deleted successfully" + if budget_deleted + else f"Access group '{access_group}' has no budget to delete" + ), + ) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 3d60f4f5f3a..09d3dedaafa 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -47,6 +47,7 @@ from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, get_or_create_metadata_bucket, ) +from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -577,6 +578,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): _metadata["user_api_key"] = user_api_key_dict.api_key _metadata["litellm_parent_otel_span"] = user_api_key_dict.parent_otel_span _metadata["user_api_key_budget_reservation"] = user_api_key_dict.budget_reservation + _metadata[MODEL_ACCESS_GROUP_METADATA_KEY] = user_api_key_dict.matched_model_access_groups # The per-model budget counters are keyed off these. get_sanitized_user_information_from_key # returns StandardLoggingUserAPIKeyMetadata, which carries no budget field, so without this # the post-call increment finds nothing and every passthrough request goes untracked and diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6ed3c86e2da..654fe4a3f2e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -382,6 +382,8 @@ from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, end_user_cache_key, get_management_object_ttl, + model_access_group_cache_key, + model_access_group_spend_counter_key, tag_cache_key, ) from litellm.proxy.config_resolvers import resolve_fields @@ -2658,6 +2660,7 @@ async def increment_spend_counters( tags: list[str] | None = None, request_id: str | None = None, request_started_at: datetime | None = None, + model_access_groups: Sequence[str] | None = None, ): """ Atomically increment spend counters for budget enforcement. @@ -2813,6 +2816,13 @@ async def increment_spend_counters( ) if end_user_id is not None or tags is not None else None, + _increment_model_access_group_spend_counters( + model_access_groups=model_access_groups, + response_cost=cost, + reserved_counter_keys=reserved_counter_keys, + ) + if model_access_groups + else None, _increment_org_spend_counter( org_id=org_id, response_cost=cost, @@ -2901,6 +2911,33 @@ async def _increment_end_user_and_tag_spend_counters( ) +async def _increment_model_access_group_spend_counters( + model_access_groups: Sequence[object], + response_cost: float, + reserved_counter_keys: set[str], +) -> None: + """Charge the model access groups that authorized this request. + + Without this the counter auth reads is written only by the reservation path, so + ``disable_budget_reservation`` would leave ``_model_access_group_max_budget_check`` enforcing + against the DB row's spend, which lags by up to the cache TTL. + + Typed ``object`` rather than ``str`` because the names reach the cost callback out of request + metadata, which the coercion upstream filters to a list but not to strings. A non-string that + slipped through would build a counter key nothing else ever reads. + """ + unique_groups: Final = tuple( + dict.fromkeys(group for group in model_access_groups if group and isinstance(group, str)) + ) + for group in unique_groups: + await _init_and_increment_unreserved_spend_counter( + counter_key=model_access_group_spend_counter_key(group), + source_cache_key=model_access_group_cache_key(group), + increment=response_cost, + reserved_counter_keys=reserved_counter_keys, + ) + + async def _increment_org_spend_counter( org_id: str | None, response_cost: float, diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 7818e3acd56..60223265211 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/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,20 @@ model LiteLLM_EndUserTable { blocked Boolean @default(false) } +// Budget and shared spend for a model access group. The groups themselves are not rows anywhere: +// they are free-text strings in LiteLLM_ProxyModelTable.model_info.access_groups, so a row here +// exists only once someone gives that group a budget. +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 diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 2ba8cfb8b37..8b2a5dd9312 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -12,7 +12,6 @@ from fastapi import HTTPException, status import litellm from litellm._logging import verbose_proxy_logger -from litellm.caching import DualCache from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.litellm_core_utils.llm_cost_calc.tiered_pricing import select_tier_for_input, tier_rate from litellm.proxy._types import ( @@ -26,12 +25,16 @@ from litellm.proxy.auth.auth_utils import get_model_from_request from litellm.proxy.auth.budget_throttle import should_throttle_budget_exceeded from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.user_api_key_cache import ( + UserApiKeyCache, end_user_cache_key, + model_access_group_cache_key, + model_access_group_spend_counter_key, tag_cache_key, team_membership_reservation_cache_key, ) from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.router import Router +from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget @dataclass @@ -54,6 +57,7 @@ _COUNTER_ENTITY_TYPES: Final[Mapping[str, str]] = { "User": Litellm_EntityType.USER.value, "EndUser": Litellm_EntityType.END_USER.value, "Tag": Litellm_EntityType.TAG.value, + "Model access group": Litellm_EntityType.MODEL_ACCESS_GROUP.value, "Organization": Litellm_EntityType.ORGANIZATION.value, } @@ -159,7 +163,7 @@ async def reserve_budget_for_request( team_object: LiteLLM_TeamTable | None, user_object: LiteLLM_UserTable | None, prisma_client: PrismaClient | None, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, end_user_id: str | None = None, end_user_object: object = None, @@ -349,7 +353,7 @@ async def _get_budget_counters( team_object: LiteLLM_TeamTable | None, user_object: LiteLLM_UserTable | None, prisma_client: PrismaClient | None, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, end_user_id: str | None = None, end_user_object: object = None, @@ -438,6 +442,14 @@ async def _get_budget_counters( ) ) + counters.extend( + await _get_model_access_group_budget_counters( + valid_token=valid_token, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + ) + team_member_counter: Final = await _get_team_member_budget_counter( valid_token=valid_token, team_object=team_object, @@ -492,7 +504,7 @@ async def _get_end_user_budget_counter( async def _get_tag_budget_counters( request_body: dict, prisma_client: PrismaClient | None, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging, ) -> list[_BudgetCounter]: from litellm.proxy.auth.auth_checks import get_tag_objects_batch @@ -531,6 +543,46 @@ async def _get_tag_budget_counters( return counters +async def _get_model_access_group_budget_counters( + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, +) -> list[_BudgetCounter]: + """Reservation counters for the model access groups that authorized this request. + + The names come off the auth object rather than the request body: ``common_checks`` already + resolved which granted groups serve the requested model, and re-deriving that here would both + duplicate the walk and risk disagreeing with what the spend writer attributes. + """ + from litellm.proxy.auth.auth_checks import get_model_access_group_budgets_batch + + group_names: Final = tuple(dict.fromkeys(valid_token.matched_model_access_groups or ())) + if not group_names: + return [] + + budgets: Final = await get_model_access_group_budgets_batch( + access_group_names=group_names, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + candidates: Final = (_model_access_group_counter(group, budgets.get(group)) for group in group_names) + return [counter for counter in candidates if counter is not None] + + +def _model_access_group_counter(group: str, budget: ModelAccessGroupBudget | None) -> _BudgetCounter | None: + """A counter for one group, or nothing when the group carries no budget to reserve against.""" + if budget is None or budget.max_budget is None or budget.max_budget <= 0: + return None + return _BudgetCounter( + counter_key=model_access_group_spend_counter_key(group), + source_cache_key=model_access_group_cache_key(group), + max_budget=budget.max_budget, + fallback_spend=budget.spend, + entity_type="Model access group", + entity_id=group, + ) + + def _dedupe_tags(tags: list[str]) -> list[str]: seen: Final = set() deduped_tags: Final = [] @@ -546,7 +598,7 @@ async def _get_team_member_budget_counter( valid_token: UserAPIKeyAuth, team_object: LiteLLM_TeamTable | None, user_object: LiteLLM_UserTable | None, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, ) -> _BudgetCounter | None: if team_object is None or team_object.team_id is None or user_object is None or valid_token.user_id is None: return None @@ -589,7 +641,7 @@ async def _get_team_member_budget_counter( async def _get_org_budget_counter( valid_token: UserAPIKeyAuth, team_object: LiteLLM_TeamTable | None, - user_api_key_cache: DualCache, + user_api_key_cache: UserApiKeyCache, ) -> _BudgetCounter | None: org_id: str | None = None if valid_token.org_id is not None: diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index ff182bd98ec..43709e4e6ff 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -1,6 +1,7 @@ import os import re import secrets +from collections.abc import Mapping from datetime import datetime, timezone from datetime import datetime as dt from typing import Any, Final, Literal, cast @@ -23,7 +24,11 @@ from litellm.litellm_core_utils.core_helpers import ( reconstruct_model_name, ) from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call -from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash +from litellm.litellm_core_utils.litellm_logging import ( + coerce_model_access_groups, + is_valid_sha256_hash, + request_model_access_groups_from_litellm_params, +) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps, strip_null_bytes from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error @@ -249,6 +254,23 @@ def _extract_usage_for_ocr_call(response_obj: Any, response_obj_dict: dict) -> d return {} +def get_request_model_access_groups(kwargs: Mapping[str, object] | None) -> tuple[str, ...]: + """Model access groups that authorized this request, as stamped onto request metadata at auth time.""" + if kwargs is None: + return () + + standard_logging_payload: Final = kwargs.get("standard_logging_object") + if isinstance(standard_logging_payload, Mapping): + from_payload: Final = coerce_model_access_groups(standard_logging_payload.get("request_model_access_groups")) + if from_payload: + return from_payload + + litellm_params: Final = kwargs.get("litellm_params") + if not isinstance(litellm_params, Mapping): + return () + return request_model_access_groups_from_litellm_params(litellm_params) + + def _sl_attribution_fallback( standard_logging_payload: StandardLoggingPayload | None, field: Literal["model_id", "model_group", "api_base", "custom_llm_provider"], diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index eab56c31c39..cf56fc0b1dd 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -6031,10 +6031,42 @@ def _should_use_smtp_ssl(smtp_port: int) -> bool: return os.getenv("SMTP_USE_SSL", "False") == "True" or smtp_port == 465 -def _create_smtp_connection(smtp_host: str, smtp_port: int) -> smtplib.SMTP: +def _create_smtp_connection(smtp_host: str, smtp_port: int, timeout: float) -> smtplib.SMTP: if _should_use_smtp_ssl(smtp_port=smtp_port): - return smtplib.SMTP_SSL(host=smtp_host, port=smtp_port, context=ssl.create_default_context()) - return smtplib.SMTP(host=smtp_host, port=smtp_port) + return smtplib.SMTP_SSL(host=smtp_host, port=smtp_port, context=ssl.create_default_context(), timeout=timeout) + return smtplib.SMTP(host=smtp_host, port=smtp_port, timeout=timeout) + + +def _send_smtp_message( + email_message: MIMEMultipart, + smtp_host: str, + smtp_port: int, + smtp_username: str | None, + smtp_password: str | None, + sender_email: str, + receiver_email: str, + timeout: float, +) -> None: + using_ssl: Final = _should_use_smtp_ssl(smtp_port=smtp_port) + with _create_smtp_connection( + smtp_host=smtp_host, + smtp_port=smtp_port, + timeout=timeout, + ) as server: + if not using_ssl and os.getenv("SMTP_TLS", "True") != "False": + server.starttls(context=ssl.create_default_context()) + + if smtp_username and smtp_password: + server.login( + user=smtp_username, + password=smtp_password, + ) + + server.send_message( + msg=email_message, + from_addr=sender_email, + to_addrs=receiver_email, + ) async def send_email( @@ -6080,27 +6112,18 @@ async def send_email( email_message.attach(MIMEText(html, "html")) try: - using_ssl: Final = _should_use_smtp_ssl(smtp_port=smtp_port) - with _create_smtp_connection( + smtp_timeout: Final = float(os.getenv("SMTP_TIMEOUT", "30")) + await asyncio.to_thread( + _send_smtp_message, + email_message=email_message, smtp_host=smtp_host, smtp_port=smtp_port, - ) as server: - if not using_ssl and os.getenv("SMTP_TLS", "True") != "False": - server.starttls(context=ssl.create_default_context()) - - # Login to your email account only if smtp_username and smtp_password are provided - if smtp_username and smtp_password: - server.login( - user=smtp_username, - password=smtp_password, - ) - - # Send the email - server.send_message( - msg=email_message, - from_addr=sender_email, - to_addrs=receiver_email, - ) + smtp_username=smtp_username, + smtp_password=smtp_password, + sender_email=sender_email, + receiver_email=receiver_email, + timeout=smtp_timeout, + ) except Exception as e: verbose_proxy_logger.exception("An error occurred while sending the email:" + str(e)) diff --git a/litellm/repositories/prisma_protocols.py b/litellm/repositories/prisma_protocols.py index 2aa1b8e0e3f..d962934dfb1 100644 --- a/litellm/repositories/prisma_protocols.py +++ b/litellm/repositories/prisma_protocols.py @@ -144,4 +144,7 @@ class PrismaBatch(Protocol): @property def litellm_endusertable(self) -> BatchTable: ... + @property + def litellm_modelaccessgroupbudgettable(self) -> BatchTable: ... + async def commit(self) -> None: ... diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index 3a3593cc5f3..18cf884f267 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -104,6 +104,10 @@ 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/litellm/repositories/unit_of_work.py b/litellm/repositories/unit_of_work.py index eb11ebe3b9c..a497d0580db 100644 --- a/litellm/repositories/unit_of_work.py +++ b/litellm/repositories/unit_of_work.py @@ -118,6 +118,7 @@ class BudgetCascadeUnitOfWork: keys: LinkedSpendResetWrites organizations: LinkedSpendResetWrites tags: LinkedSpendResetWrites + model_access_groups: LinkedSpendResetWrites endusers: LinkedSpendResetWrites budgets: BudgetWindowWrites @@ -143,6 +144,7 @@ async def budget_cascade_unit_of_work( keys=LinkedSpendResetWrites(table=batch.litellm_verificationtoken), organizations=LinkedSpendResetWrites(table=batch.litellm_organizationtable), tags=LinkedSpendResetWrites(table=batch.litellm_tagtable), + model_access_groups=LinkedSpendResetWrites(table=batch.litellm_modelaccessgroupbudgettable), endusers=LinkedSpendResetWrites(table=batch.litellm_endusertable), budgets=BudgetWindowWrites(table=batch.litellm_budgettable), ) diff --git a/litellm/types/proxy/management_endpoints/model_management_endpoints.py b/litellm/types/proxy/management_endpoints/model_management_endpoints.py index 6e18787a224..9d4663631fe 100644 --- a/litellm/types/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/types/proxy/management_endpoints/model_management_endpoints.py @@ -1,6 +1,7 @@ +from datetime import datetime from typing import Any -from pydantic import BaseModel, Field +from pydantic import BaseModel, ConfigDict, Field from ...router import ModelGroupInfo @@ -53,10 +54,42 @@ class DeleteModelGroupResponse(BaseModel): message: str +class AccessGroupBudget(BaseModel): + budget_id: str + max_budget: float | None = None + soft_budget: float | None = None + budget_duration: str | None = None + budget_reset_at: datetime | None = None + + +class AccessGroupBudgetRequest(BaseModel): + budget_id: str | None = None # Link an existing budget instead of creating one + max_budget: float | None = Field(default=None, ge=0) + soft_budget: float | None = Field(default=None, ge=0) + budget_duration: str | None = None + + # rejects tpm_limit/rpm_limit/max_parallel_requests: those are not enforced per access group + model_config = ConfigDict(extra="forbid") + + +class AccessGroupBudgetResponse(BaseModel): + access_group: str + spend: float # Shared spend accrued by every key that can reach this access group + budget: AccessGroupBudget | None = None + + +class DeleteAccessGroupBudgetResponse(BaseModel): + access_group: str + budget_deleted: bool # False when the access group had no budget to begin with + message: str + + class AccessGroupInfo(BaseModel): access_group: str model_names: list[str] # List of model names in this access group deployment_count: int # Total number of deployments with this access group + spend: float | None = None # Only populated by /access_group/{access_group}/info + budget: AccessGroupBudget | None = None class ListAccessGroupsResponse(BaseModel): diff --git a/litellm/types/proxy/model_access_group_budget.py b/litellm/types/proxy/model_access_group_budget.py new file mode 100644 index 00000000000..cccbe92b5d6 --- /dev/null +++ b/litellm/types/proxy/model_access_group_budget.py @@ -0,0 +1,19 @@ +"""The model access group budget state auth and the spend reservation path share.""" + +from __future__ import annotations + +from pydantic import BaseModel + + +class ModelAccessGroupBudget(BaseModel): + """One model access group's budget, flattened out of its joined ``LiteLLM_ModelAccessGroupBudgetTable`` row. + + Both readers want only the recorded spend and the ceiling, and this sits on the per-request hot + path behind a cache, so the linked budget row is collapsed to ``max_budget`` rather than cached + whole. ``spend`` is the DB-recorded value, which lags the live counter and is only ever a + fallback for it. + """ + + access_group_name: str + spend: float = 0.0 + max_budget: float | None = None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 2d06f5e0abe..4bf8289d725 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -40,7 +40,7 @@ from pydantic import ( field_serializer, field_validator, ) -from typing_extensions import ReadOnly, Required, TypedDict +from typing_extensions import NotRequired, ReadOnly, Required, TypedDict from litellm._logging import verbose_logger from litellm._uuid import uuid @@ -3260,6 +3260,7 @@ class StandardLoggingPayload(TypedDict): cache_key: str | None saved_cache_cost: float request_tags: list + request_model_access_groups: NotRequired[ReadOnly[Sequence[str]]] end_user: str | None requester_ip_address: str | None user_agent: str | None diff --git a/schema.prisma b/schema.prisma index 7818e3acd56..60223265211 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,20 @@ model LiteLLM_EndUserTable { blocked Boolean @default(false) } +// Budget and shared spend for a model access group. The groups themselves are not rows anywhere: +// they are free-text strings in LiteLLM_ProxyModelTable.model_info.access_groups, so a row here +// exists only once someone gives that group a budget. +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 diff --git a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt index d10be89b90c..052962e078e 100644 --- a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt +++ b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt @@ -126,3 +126,6 @@ POST /customer/delete # known gap: litellm_customer GET /team/{team_id}/callback # known gap: team callback resource POST /team/{team_id}/callback # known gap: team callback resource DELETE /team/{team_id}/callback/{callback_name} # known gap: team callback resource +GET /access_group/{access_group}/budget # known gap: budget attributes on litellm_access_group +PUT /access_group/{access_group}/budget # known gap: budget attributes on litellm_access_group +DELETE /access_group/{access_group}/budget # known gap: budget attributes on litellm_access_group diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index cf912ddaa25..14b2b4e3299 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -177,15 +177,15 @@ quota_management... behavior : ratelimit | budget | spend_tracking variant : rpm | tpm | priority_generous | priority_strict key | internal_user | end_user | organization | team | team_member | tag - | model_max | soft | key_multi_window | team_multi_window - | fallback | spend_counter + | model_access_group | model_max | soft | key_multi_window + | team_multi_window | fallback | spend_counter chat_completions | stream | messages_bridge | embeddings | cache_hit | key_rollup | concurrent_burst | tags | end_user | per_model | failure | spend_calculate | pagination assertion : blocks_over_limit | resets_after_window | headers_report_remaining | picks_under_tpm | blocks_then_resets | resets_windows_independently | alerts_without_blocking - | isolates_per_model | isolates_per_member | enforced_across_keys | routes_to_fallback - | reseed_matches_db | logs_cost | zero_cost + | isolates_per_model | isolates_per_member | isolates_per_group | enforced_across_keys + | routes_to_fallback | reseed_matches_db | reports_spend | logs_cost | zero_cost | matches_sum_of_logs | loses_no_spend | attributes_spend | writes_own_rows | writes_failure_row | returns_cost | keeps_total e.g. quota_management.ratelimit.rpm.blocks_over_limit exercised_on=[chat_completions, messages] diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index 42a075681e0..d0afcaca848 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -20,6 +20,10 @@ - {id: quota_management.budget.organization.blocks_over_limit, module: quota_management, tier: P1, behavior: budget, variant: organization, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "An organization's max_budget blocks keys under its teams"} - {id: quota_management.budget.team_member.blocks_over_limit, module: quota_management, tier: P1, behavior: budget, variant: team_member, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "A member's per-team budget blocks independently of the team budget"} - {id: quota_management.budget.team_member.isolates_per_member, module: quota_management, tier: P1, behavior: budget, variant: team_member, assertions: [isolates_per_member], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "One team member's exhausted per-team budget does not block a different member on the same team"} +- {id: quota_management.budget.model_access_group.blocks_over_limit, module: quota_management, tier: P1, behavior: budget, variant: model_access_group, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "A model access group's shared max_budget blocks further calls to deployments in the group once the pool is spent"} +- {id: quota_management.budget.model_access_group.enforced_across_keys, module: quota_management, tier: P1, behavior: budget, variant: model_access_group, assertions: [enforced_across_keys], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "The pool is shared, so a key that spent nothing of its own is blocked once another key granted the same group drained it"} +- {id: quota_management.budget.model_access_group.isolates_per_group, module: quota_management, tier: P1, behavior: budget, variant: model_access_group, assertions: [isolates_per_group], exercised_on: [chat_completions], source: "proxy/db/db_spend_update_writer.py", rationale: "A request is charged only to the granted groups that serve the model it called, so an exhausted group never blocks a sibling group"} +- {id: quota_management.budget.model_access_group.reports_spend, module: quota_management, tier: P2, behavior: budget, variant: model_access_group, assertions: [reports_spend], exercised_on: [chat_completions], source: "proxy/management_endpoints/model_access_group_management_endpoints.py", rationale: "GET /access_group/{name}/budget reports the pool and the spend drawn against it, so an admin can see why calls are being refused"} - {id: quota_management.budget.tag.blocks_over_limit, module: quota_management, tier: P1, behavior: budget, variant: tag, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "router_strategy/budget_limiter.py", rationale: "Proxy-level tag budgets block tagged requests at the cap"} - {id: quota_management.budget.end_user_model_max.blocks_over_limit, module: quota_management, tier: P1, behavior: budget, variant: end_user_model_max, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "budget_management_endpoints.py", fail_before_fix: proven, rationale: "A per-model rpm_limit on an end-user budget is accepted and stored but never enforced; only key-attached budgets honour it"} - {id: quota_management.budget.model_max.isolates_per_model, module: quota_management, tier: P1, behavior: budget, variant: model_max, assertions: [isolates_per_model], exercised_on: [chat_completions], source: "proxy/hooks/model_max_budget_limiter.py", rationale: "model_max_budget caps one model without touching a sibling's budget"} diff --git a/tests/e2e/quota_management/budgets/budget_client.py b/tests/e2e/quota_management/budgets/budget_client.py index 543d5f959e0..087dc8ca522 100644 --- a/tests/e2e/quota_management/budgets/budget_client.py +++ b/tests/e2e/quota_management/budgets/budget_client.py @@ -151,6 +151,28 @@ class TagDeleteBody(BaseModel): name: str +class AccessGroupBudgetBody(BaseModel): + max_budget: float | None = None + soft_budget: float | None = None + budget_duration: str | None = None + + +class AccessGroupBudgetView(BaseModel): + budget_id: str + max_budget: float | None = None + soft_budget: float | None = None + budget_duration: str | None = None + + +class AccessGroupBudgetResponse(BaseModel): + """GET/PUT /access_group/{name}/budget: the group's shared pool and the spend + every key that can reach the group has drawn against it.""" + + access_group: str + spend: float + budget: AccessGroupBudgetView | None = None + + class BudgetNewBody(BaseModel): max_budget: float | None = None soft_budget: float | None = None @@ -514,6 +536,49 @@ class BudgetClient: response_type=NoBody, ) + # ---- model access group --------------------------------------------- + + def set_access_group_budget( + self, + access_group: str, + *, + max_budget: float | None = None, + soft_budget: float | None = None, + budget_duration: str | None = None, + ) -> AccessGroupBudgetResponse: + """Give a model access group one shared budget. Every key that can reach a + deployment in the group draws from it.""" + return unwrap( + self.proxy.transport.put( + f"/access_group/{access_group}/budget", + headers=self.proxy.transport.master, + json=AccessGroupBudgetBody( + max_budget=max_budget, + soft_budget=soft_budget, + budget_duration=budget_duration, + ), + response_type=AccessGroupBudgetResponse, + ) + ) + + def access_group_budget(self, access_group: str) -> AccessGroupBudgetResponse: + return unwrap( + self.proxy.transport.get( + f"/access_group/{access_group}/budget", + headers=self.proxy.transport.master, + params=NoBody(), + response_type=AccessGroupBudgetResponse, + ) + ) + + def delete_access_group_budget(self, access_group: str) -> None: + _ = self.proxy.transport.delete( + f"/access_group/{access_group}/budget", + headers=self.proxy.transport.master, + json=NoBody(), + response_type=NoBody, + ) + # ---- budget table --------------------------------------------------- def create_budget( diff --git a/tests/e2e/quota_management/budgets/test_model_access_group_budget_e2e.py b/tests/e2e/quota_management/budgets/test_model_access_group_budget_e2e.py new file mode 100644 index 00000000000..9c927a31216 --- /dev/null +++ b/tests/e2e/quota_management/budgets/test_model_access_group_budget_e2e.py @@ -0,0 +1,161 @@ +"""Live e2e: one shared budget across every key that can reach a model access group. + +A model access group is a free-text label on a deployment (`model_info.access_groups`), +and a key is granted the group by name. The budget hangs off the group, not the key, so +the interesting behaviors are the ones a per-key budget cannot produce: a key that has +spent nothing of its own is refused once somebody else drained the pool, and draining one +group leaves a second group untouched, because a request is only charged to the groups +the caller was granted that also serve the model being called. +""" + +from __future__ import annotations + +import os +import time +from collections.abc import Iterator +from dataclasses import dataclass +from typing import Final + +import pytest + +from budget_client import BudgetClient, is_budget_block +from e2e_config import unique_marker +from e2e_http import StreamingResponse, require_successful_call +from lifecycle import ResourceManager +from models import KeyGenerateBody, LiteLLMParamsBody, ModelInfoBody, ModelNewBody + +pytestmark = pytest.mark.e2e + +BACKEND: Final = "openai/gpt-5.4-nano" +TINY_BUDGET: Final = 5e-6 +MAX_TOKENS: Final = 16 +DRAIN_TIMEOUT_SECONDS: Final = 180 + + +@dataclass(frozen=True, slots=True) +class DrainedPool: + """A model access group whose shared budget has been spent to exhaustion, the + deployment inside it, the key that did the spending, and a second group holding + its own deployment that was never given a budget at all.""" + + access_group: str + model: str + spender_key: str + free_access_group: str + free_model: str + + +def _provider_key(env_var: str) -> str: + return os.environ.get(env_var) or f"os.environ/{env_var}" + + +def _grouped_model(model_name: str, access_group: str) -> ModelNewBody: + return ModelNewBody( + model_name=model_name, + litellm_params=LiteLLMParamsBody(model=BACKEND, api_key=_provider_key("OPENAI_API_KEY")), + model_info=ModelInfoBody(access_groups=[access_group]), + ) + + +def _call(client: BudgetClient, key: str, model: str) -> StreamingResponse: + return client.chat(key, model, f"hi {unique_marker()}", max_tokens=MAX_TOKENS) + + +def _drain(client: BudgetClient, key: str, model: str, access_group: str) -> None: + """Spend the group's pool until the proxy refuses the next request. The first call + lands under the cap and the block comes from the spend it recorded, so this needs at + least one round trip through the spend writer, not just one request.""" + deadline: Final = time.monotonic() + DRAIN_TIMEOUT_SECONDS + while time.monotonic() < deadline: + result = _call(client, key, model) + if is_budget_block(result): + return + require_successful_call(result) + time.sleep(1) + pytest.fail(f"budget on model access group {access_group!r} never blocked a request") + + +@pytest.fixture(scope="module") +def drained(client: BudgetClient) -> Iterator[DrainedPool]: + marker: Final = unique_marker() + pool: Final = DrainedPool( + access_group=f"e2e-mag-budget-{marker}", + model=f"e2e-mag-budgeted-{marker}", + spender_key=client.proxy.generate_key(KeyGenerateBody(models=[f"e2e-mag-budget-{marker}"])), + free_access_group=f"e2e-mag-free-{marker}", + free_model=f"e2e-mag-unbudgeted-{marker}", + ) + created: Final = ( + client.proxy.register_model(_grouped_model(pool.model, pool.access_group)), + client.proxy.register_model(_grouped_model(pool.free_model, pool.free_access_group)), + ) + try: + client.set_access_group_budget(pool.access_group, max_budget=TINY_BUDGET) + _drain(client, pool.spender_key, pool.model, pool.access_group) + yield pool + finally: + client.delete_access_group_budget(pool.access_group) + client.proxy.delete_key(pool.spender_key) + for model_id in created: + client.proxy.delete_model(model_id) + + +class TestModelAccessGroupBudget: + @pytest.mark.covers("quota_management.budget.model_access_group.blocks_over_limit") + def test_the_key_that_drained_the_pool_stays_blocked( + self, client: BudgetClient, drained: DrainedPool + ) -> None: + result = _call(client, drained.spender_key, drained.model) + assert is_budget_block(result), ( + f"an exhausted pool served {drained.model!r} again: {result.status_code} {result.body[:300]}" + ) + assert drained.access_group in result.body, ( + f"the block did not name the group that caused it: {result.body[:300]}" + ) + + @pytest.mark.covers("quota_management.budget.model_access_group.enforced_across_keys") + def test_a_key_that_spent_nothing_is_blocked_by_the_shared_pool( + self, client: BudgetClient, resources: ResourceManager, drained: DrainedPool + ) -> None: + newcomer = resources.key(models=[drained.access_group]) + + result = _call(client, newcomer, drained.model) + + assert is_budget_block(result), ( + "a freshly minted key with no spend of its own was served by an exhausted " + f"shared pool: {result.status_code} {result.body[:300]}" + ) + + @pytest.mark.covers("quota_management.budget.model_access_group.isolates_per_group") + def test_a_drained_group_does_not_block_a_different_group( + self, client: BudgetClient, resources: ResourceManager, drained: DrainedPool + ) -> None: + other = resources.key(models=[drained.free_access_group]) + + result = _call(client, other, drained.free_model) + + assert not is_budget_block(result), ( + f"{drained.free_access_group!r} has no budget of its own but was blocked by " + f"{drained.access_group!r}'s exhausted pool: {result.body[:300]}" + ) + require_successful_call(result) + + @pytest.mark.covers("quota_management.budget.model_access_group.reports_spend") + def test_the_budget_read_reports_the_spend_drawn_against_the_pool( + self, client: BudgetClient, drained: DrainedPool + ) -> None: + """Enforcement runs off a live counter while the group's row is written by the + batched spend writer, so the recorded spend an admin reads lands a beat after the + block. Poll for it: what matters is that it arrives and matches the pool.""" + deadline = time.monotonic() + client.proxy.poll_timeout + reported = client.access_group_budget(drained.access_group) + while reported.spend < TINY_BUDGET and time.monotonic() < deadline: + time.sleep(client.proxy.poll_interval) + reported = client.access_group_budget(drained.access_group) + + assert reported.budget is not None, "the group lost the budget that just blocked it" + assert reported.budget.max_budget == TINY_BUDGET + assert reported.spend >= TINY_BUDGET, ( + f"the pool blocked at {TINY_BUDGET} but only {reported.spend} was ever recorded " + f"against the group within {client.proxy.poll_timeout}s" + ) diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 1c717228ecc..a45a6aaa2a3 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -3735,6 +3735,60 @@ def test_get_standard_logging_object_payload_includes_litellm_call_id(logging_ob assert payload["litellm_call_id"] == call_id +def test_get_standard_logging_object_payload_carries_matched_access_groups(logging_obj): + """Access groups stamped at auth time reach the logging payload, so integrations see what a request billed.""" + from datetime import datetime + + from litellm.litellm_core_utils.litellm_logging import ( + get_standard_logging_object_payload, + ) + + now = datetime.now() + payload = get_standard_logging_object_payload( + kwargs={ + "model": "gpt-4o", + "messages": [], + "litellm_params": { + "metadata": { + "user_api_key_matched_model_access_groups": ["premium-pool", "shared-pool"] + }, + "proxy_server_request": {"body": {}}, + }, + }, + init_response_obj={}, + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + + assert payload is not None + assert payload["request_model_access_groups"] == ("premium-pool", "shared-pool") + + +def test_get_standard_logging_object_payload_has_no_access_groups_when_unstamped( + logging_obj, +): + from datetime import datetime + + from litellm.litellm_core_utils.litellm_logging import ( + get_standard_logging_object_payload, + ) + + now = datetime.now() + payload = get_standard_logging_object_payload( + kwargs={"model": "gpt-4o", "messages": []}, + init_response_obj={}, + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + + assert payload is not None + assert payload["request_model_access_groups"] == () + + def test_get_standard_logging_object_payload_preserves_absent_end_user_as_none(logging_obj): from datetime import datetime from typing import Final diff --git a/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py b/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py new file mode 100644 index 00000000000..fb82d8708fd --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_model_access_group_budgets.py @@ -0,0 +1,543 @@ +""" +Which model access groups a request is charged to. + +A group is attributed only when its name appears on an allowlist the caller was granted, so the +group is what authorized the call. Asking for a model that merely belongs to a group attributes +nothing, and every level that can name a group (key, team, team-member scope, project, org) is +unioned rather than ranked. +""" + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +import litellm +from litellm import Router +from litellm.proxy._types import ( + LiteLLM_BudgetTable, + Litellm_EntityType, + LiteLLM_OrganizationTable, + LiteLLM_ProjectTableCachedObj, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + UserAPIKeyAuth, +) +from litellm.proxy.auth.auth_checks import ( + _model_access_group_max_budget_check, + collect_matched_model_access_groups, + common_checks, + stamp_matched_model_access_groups, +) +from litellm.proxy.common_utils.reset_budget_job import _model_access_group_counter_key +from litellm.proxy.common_utils.user_api_key_cache import ( + UserApiKeyCache, + model_access_group_registry_cache_key, + model_access_group_spend_counter_key, + team_membership_reservation_cache_key, +) +from litellm.proxy.utils import ProxyLogging + +TEAM_ID = "team-1" +USER_ID = "user-1" +ORG_ID = "org-1" +BUDGETED_GROUPS = ("tier-a", "tier-b", "claude-tier") +MODEL_ACCESS_GROUP_COUNTER_KEY = model_access_group_spend_counter_key("tier-a") + +MODEL_LIST = [ + { + "model_name": "gpt-4o", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "k"}, + "model_info": {"access_groups": ["tier-a", "tier-b"]}, + }, + { + "model_name": "claude-sonnet", + "litellm_params": {"model": "anthropic/claude-sonnet", "api_key": "k"}, + "model_info": {"access_groups": ["claude-tier"]}, + }, +] + + +class _ExplodingPrismaClient: + """Every lookup in these tests is served from the injected cache; a real DB read is a bug.""" + + def __getattr__(self, name: str) -> object: + raise AssertionError(f"unexpected database access: {name}") + + +class _CountingRouter(Router): + """Counts access-group lookups, so a test can prove the registry gate skipped them.""" + + def __init__(self, **kwargs) -> None: + super().__init__(**kwargs) + self.access_group_lookups = 0 + + def get_model_access_groups(self, *args, **kwargs): + self.access_group_lookups += 1 + return super().get_model_access_groups(*args, **kwargs) + + +async def _cache( + budgeted_groups: tuple[str, ...] = BUDGETED_GROUPS, + member_allowed_models: tuple[str, ...] = (), + org_models: tuple[str, ...] = (), +) -> UserApiKeyCache: + cache = UserApiKeyCache() + await cache.async_set_cache(key=model_access_group_registry_cache_key(), value=budgeted_groups) + if member_allowed_models: + await cache.async_set_cache( + key=team_membership_reservation_cache_key(user_id=USER_ID, team_id=TEAM_ID), + value=LiteLLM_TeamMembership( + user_id=USER_ID, + team_id=TEAM_ID, + budget_id="member-budget", + litellm_budget_table=LiteLLM_BudgetTable(allowed_models=list(member_allowed_models)), + ), + model_type=LiteLLM_TeamMembership, + ) + if org_models: + await cache.async_set_cache( + key=f"org_id:{ORG_ID}", + value=LiteLLM_OrganizationTable( + organization_id=ORG_ID, + budget_id="org-budget", + models=list(org_models), + created_by=USER_ID, + updated_by=USER_ID, + ), + model_type=LiteLLM_OrganizationTable, + ) + return cache + + +async def _matched( + *, + model: str = "gpt-4o", + key_models: list[str] | None = None, + team_models: list[str] | None = None, + team_org_id: str | None = None, + project_models: list[str] | None = None, + valid_token: UserAPIKeyAuth | None = None, + cache: UserApiKeyCache | None = None, + llm_router: Router | None = None, +) -> tuple[str, ...]: + resolved_cache = cache if cache is not None else await _cache() + return await collect_matched_model_access_groups( + model=model, + valid_token=valid_token + if valid_token is not None + else UserAPIKeyAuth(api_key="hashed", models=key_models or [], team_id=TEAM_ID, user_id=USER_ID), + team_object=( + LiteLLM_TeamTable(team_id=TEAM_ID, models=team_models, organization_id=team_org_id) + if team_models is not None + else None + ), + project_object=( + LiteLLM_ProjectTableCachedObj(project_id="project-1", models=project_models) + if project_models is not None + else None + ), + llm_router=llm_router if llm_router is not None else Router(model_list=MODEL_LIST), + prisma_client=_ExplodingPrismaClient(), + user_api_key_cache=resolved_cache, + proxy_logging_obj=ProxyLogging(user_api_key_cache=resolved_cache), + ) + + +@pytest.mark.asyncio +async def test_group_named_on_the_key_is_attributed(): + assert await _matched(key_models=["tier-a"]) == ("tier-a",) + + +@pytest.mark.asyncio +async def test_model_granted_directly_on_the_key_attributes_nothing(): + assert await _matched(key_models=["gpt-4o"]) == () + + +@pytest.mark.asyncio +@pytest.mark.parametrize("key_models", [["*"], [], ["all-proxy-models"]]) +async def test_unrestricted_key_attributes_nothing(key_models: list[str]): + assert await _matched(key_models=key_models) == () + + +@pytest.mark.asyncio +async def test_group_that_does_not_serve_the_requested_model_is_not_attributed(): + assert await _matched(model="gpt-4o", key_models=["claude-tier"]) == () + + +@pytest.mark.asyncio +async def test_both_granted_groups_covering_the_model_are_attributed(): + assert await _matched(key_models=["tier-b", "tier-a"]) == ("tier-a", "tier-b") + + +@pytest.mark.asyncio +async def test_group_named_only_on_the_team_is_attributed(): + assert await _matched(key_models=[], team_models=["tier-a"]) == ("tier-a",) + + +@pytest.mark.asyncio +async def test_group_named_only_in_a_team_members_scope_is_attributed(): + assert await _matched( + model="claude-sonnet", + key_models=["*"], + team_models=["*"], + cache=await _cache(member_allowed_models=("claude-tier",)), + ) == ("claude-tier",) + + +@pytest.mark.asyncio +async def test_group_named_only_on_the_project_is_attributed(): + assert await _matched(key_models=["*"], project_models=["tier-b"]) == ("tier-b",) + + +@pytest.mark.asyncio +async def test_group_named_only_on_the_org_is_attributed(): + assert await _matched( + valid_token=UserAPIKeyAuth(api_key="hashed", models=["*"], user_id=USER_ID, org_id=ORG_ID), + cache=await _cache(org_models=("tier-a",)), + ) == ("tier-a",) + + +@pytest.mark.asyncio +async def test_group_named_on_the_teams_org_is_attributed_when_the_key_names_no_org(): + assert await _matched( + key_models=["*"], + team_models=["*"], + team_org_id=ORG_ID, + cache=await _cache(org_models=("tier-b",)), + ) == ("tier-b",) + + +@pytest.mark.asyncio +async def test_all_team_models_sentinel_on_the_key_resolves_to_the_teams_groups(): + assert await _matched( + valid_token=UserAPIKeyAuth( + api_key="hashed", + models=["all-team-models"], + team_models=["tier-a"], + team_id=TEAM_ID, + user_id=USER_ID, + ), + ) == ("tier-a",) + + +@pytest.mark.asyncio +async def test_group_without_a_budget_is_not_attributed(): + assert await _matched(key_models=["tier-a"], cache=await _cache(budgeted_groups=("tier-b",))) == () + + +@pytest.mark.asyncio +async def test_empty_registry_skips_the_access_group_matching_entirely(): + router = _CountingRouter(model_list=MODEL_LIST) + + assert await _matched(key_models=["tier-a"], cache=await _cache(budgeted_groups=()), llm_router=router) == () + assert router.access_group_lookups == 0 + + assert await _matched(key_models=["tier-a"], llm_router=router) == ("tier-a",) + assert router.access_group_lookups == 1 + + +@pytest.mark.asyncio +async def test_stamp_records_the_matched_groups_on_the_auth_object(): + cache = await _cache() + valid_token = UserAPIKeyAuth(api_key="hashed", models=["tier-a", "tier-b"], team_id=TEAM_ID, user_id=USER_ID) + + await stamp_matched_model_access_groups( + model="gpt-4o", + valid_token=valid_token, + team_object=None, + project_object=None, + llm_router=Router(model_list=MODEL_LIST), + prisma_client=_ExplodingPrismaClient(), + user_api_key_cache=cache, + proxy_logging_obj=ProxyLogging(user_api_key_cache=cache), + ) + + assert valid_token.matched_model_access_groups == ["tier-a", "tier-b"] + + +class _BrokenRouter(Router): + def get_model_access_groups(self, *args, **kwargs): + raise RuntimeError("access group store unavailable") + + +@pytest.mark.asyncio +async def test_stamp_does_not_break_auth_when_the_access_group_lookup_fails(): + cache = await _cache() + valid_token = UserAPIKeyAuth(api_key="hashed", models=["tier-a"], team_id=TEAM_ID, user_id=USER_ID) + + await stamp_matched_model_access_groups( + model="gpt-4o", + valid_token=valid_token, + team_object=None, + project_object=None, + llm_router=_BrokenRouter(model_list=MODEL_LIST), + prisma_client=_ExplodingPrismaClient(), + user_api_key_cache=cache, + proxy_logging_obj=ProxyLogging(user_api_key_cache=cache), + ) + + assert valid_token.matched_model_access_groups is None + + +@pytest.mark.asyncio +async def test_stamp_leaves_the_auth_object_untouched_when_nothing_matched(): + cache = await _cache() + valid_token = UserAPIKeyAuth(api_key="hashed", models=["gpt-4o"], team_id=TEAM_ID, user_id=USER_ID) + + await stamp_matched_model_access_groups( + model="gpt-4o", + valid_token=valid_token, + team_object=None, + project_object=None, + llm_router=Router(model_list=MODEL_LIST), + prisma_client=_ExplodingPrismaClient(), + user_api_key_cache=cache, + proxy_logging_obj=ProxyLogging(user_api_key_cache=cache), + ) + + assert valid_token.matched_model_access_groups is None + + +class _MagBudgetRow: + """One ``LiteLLM_ModelAccessGroupBudgetTable`` row as prisma hands it back.""" + + def __init__(self, access_group_name: str, spend: float = 0.0, max_budget: float | None = None) -> None: + self.access_group_name = access_group_name + self.spend = spend + self.litellm_budget_table = None if max_budget is None else SimpleNamespace(max_budget=max_budget) + + +class _RecordingPrismaClient: + """Serves budget rows and records which groups actually reached the database.""" + + def __init__(self, *rows: _MagBudgetRow) -> None: + self.rows = {row.access_group_name: row for row in rows} + self.batches: list[list[str]] = [] + self.db = SimpleNamespace( + litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=self._find_many) + ) + + async def _find_many(self, **kwargs): + requested = list(kwargs["where"]["access_group_name"]["in"]) + self.batches.append(requested) + return [self.rows[group] for group in requested if group in self.rows] + + +def _spend_reader(spend_by_counter_key: dict[str, float]): + """Stand-in for proxy_server.get_current_spend, recording every counter key it is asked for.""" + seen: list[str] = [] + + async def read(counter_key, fallback_spend, max_budget=None, **kwargs): + seen.append(counter_key) + return spend_by_counter_key.get(counter_key, fallback_spend) + + return read, seen + + +async def _enforce( + matched: tuple[str, ...], + *rows: _MagBudgetRow, + spend_by_counter_key: dict[str, float] | None = None, + prisma_client: object | None = None, + cache: UserApiKeyCache | None = None, +) -> list[str]: + read, seen = _spend_reader(spend_by_counter_key or {}) + # The check takes its client and cache as arguments, injected just below. get_current_spend is the + # one collaborator it reaches by a lazy `from litellm.proxy.proxy_server import`, with no parameter. + with patch("litellm.proxy.proxy_server.get_current_spend", read): # test-quality-ok: get_current_spend is lazily imported inside _model_access_group_max_budget_check and has no injection point + await _model_access_group_max_budget_check( + matched_model_access_groups=matched, + prisma_client=prisma_client if prisma_client is not None else _RecordingPrismaClient(*rows), + user_api_key_cache=cache if cache is not None else UserApiKeyCache(), + ) + return seen + + +@pytest.mark.asyncio +async def test_group_under_its_max_budget_passes(): + assert await _enforce( + ("tier-a",), + _MagBudgetRow("tier-a", spend=4.0, max_budget=10.0), + spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 4.0}, + ) == [MODEL_ACCESS_GROUP_COUNTER_KEY] + + +@pytest.mark.asyncio +async def test_group_exactly_at_its_max_budget_blocks_the_request(): + """A pool whose spend has reached the ceiling has nothing left, so the next request is refused. + + This is where the check departs from the tag one it otherwise mirrors, and it matches where + keys and organizations already draw the line. + """ + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _enforce( + ("tier-a",), + _MagBudgetRow("tier-a", max_budget=10.0), + spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 10.0}, + ) + + assert exc_info.value.entity_id == "tier-a" + assert exc_info.value.current_cost == 10.0 + + +@pytest.mark.asyncio +async def test_group_just_under_its_max_budget_passes(): + """Asserting the counter was read is what keeps this honest: a group that got skipped entirely, + because its row never arrived or carried no budget, would also not raise.""" + assert await _enforce( + ("tier-a",), + _MagBudgetRow("tier-a", max_budget=10.0), + spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 9.99}, + ) == [MODEL_ACCESS_GROUP_COUNTER_KEY] + + +@pytest.mark.asyncio +async def test_a_non_positive_budget_means_no_budget(): + """The reservation path treats max_budget <= 0 as unbudgeted, so the read-time check must agree. + + Without this the exclusive ceiling would turn a zero into a total freeze on one path and a + no-op on the other. + """ + assert ( + await _enforce( + ("tier-a",), + _MagBudgetRow("tier-a", max_budget=0.0), + spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 5.0}, + ) + == [] + ) + + +@pytest.mark.asyncio +async def test_group_over_its_max_budget_blocks_the_request_and_names_the_group(): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _enforce( + ("tier-a",), + _MagBudgetRow("tier-a", max_budget=10.0), + spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 10.5}, + ) + + assert exc_info.value.entity_id == "tier-a" + assert exc_info.value.entity_type == Litellm_EntityType.MODEL_ACCESS_GROUP.value + assert exc_info.value.current_cost == 10.5 + assert exc_info.value.max_budget == 10.0 + assert "tier-a" in str(exc_info.value) + + +@pytest.mark.asyncio +async def test_group_with_a_row_but_no_budget_never_blocks(): + """An admin can register a group without a ceiling; that must not become an implicit zero budget.""" + assert ( + await _enforce( + ("tier-a",), + _MagBudgetRow("tier-a", spend=9999.0), + spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 9999.0}, + ) + == [] + ) + + +@pytest.mark.asyncio +async def test_a_cold_counter_falls_back_to_the_spend_recorded_on_the_row(): + """After a counter expires the DB row is the only record of the spend, so it has to be read.""" + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _enforce(("tier-a",), _MagBudgetRow("tier-a", spend=12.0, max_budget=10.0)) + + assert exc_info.value.current_cost == 12.0 + + +@pytest.mark.asyncio +async def test_an_over_budget_group_blocks_even_when_another_matched_group_is_fine(): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _enforce( + ("tier-a", "tier-b"), + _MagBudgetRow("tier-a", max_budget=10.0), + _MagBudgetRow("tier-b", max_budget=1.0), + spend_by_counter_key={ + MODEL_ACCESS_GROUP_COUNTER_KEY: 1.0, + model_access_group_spend_counter_key("tier-b"): 5.0, + }, + ) + + assert exc_info.value.entity_id == "tier-b" + + +@pytest.mark.asyncio +async def test_request_that_matched_no_group_touches_neither_database_nor_counters(): + assert await _enforce((), prisma_client=_ExplodingPrismaClient()) == [] + + +@pytest.mark.asyncio +async def test_budget_check_reads_the_counter_key_the_reset_job_clears(): + """Reads and resets must agree, or a rollover clears a counter nobody reads.""" + reset_job_key = _model_access_group_counter_key(SimpleNamespace(access_group_name="tier-a")) + + assert await _enforce(("tier-a",), _MagBudgetRow("tier-a", max_budget=10.0)) == [reset_job_key] + + +@pytest.mark.asyncio +async def test_a_second_request_serves_the_budget_row_from_cache(): + cache = UserApiKeyCache() + prisma_client = _RecordingPrismaClient(_MagBudgetRow("tier-a", max_budget=10.0)) + + await _enforce(("tier-a",), prisma_client=prisma_client, cache=cache) + await _enforce(("tier-a",), prisma_client=prisma_client, cache=cache) + + assert prisma_client.batches == [["tier-a"]] + + +@pytest.mark.asyncio +async def test_a_database_error_does_not_block_the_request(): + class _FailingPrismaClient: + def __init__(self) -> None: + self.db = SimpleNamespace( + litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=self._boom) + ) + + async def _boom(self, **kwargs): + raise RuntimeError("database unavailable") + + assert await _enforce(("tier-a",), prisma_client=_FailingPrismaClient()) == [] + + +async def _common_checks_with_over_budget_group(*, skip_budget_checks: bool) -> bool: + cache = await _cache() + prisma_client = _RecordingPrismaClient(_MagBudgetRow("tier-a", max_budget=1.0)) + read, _ = _spend_reader({MODEL_ACCESS_GROUP_COUNTER_KEY: 99.0}) + + with ( + # common_checks resolves all three off the proxy_server module at call time; its signature + # has no client, cache or spend-reader parameter to pass them through instead. + patch("litellm.proxy.proxy_server.prisma_client", prisma_client), # test-quality-ok: common_checks lazily imports prisma_client from proxy_server and takes no client parameter + patch("litellm.proxy.proxy_server.user_api_key_cache", cache), # test-quality-ok: common_checks lazily imports user_api_key_cache from proxy_server and takes no cache parameter + patch("litellm.proxy.proxy_server.get_current_spend", read), # test-quality-ok: get_current_spend is lazily imported inside the budget check and has no injection point + ): + return await common_checks( + request_body={"model": "gpt-4o", "messages": []}, + team_object=None, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/v1/chat/completions", + llm_router=Router(model_list=MODEL_LIST), + proxy_logging_obj=ProxyLogging(user_api_key_cache=cache), + valid_token=UserAPIKeyAuth(api_key="hashed", models=["tier-a"], user_id=USER_ID), + request=SimpleNamespace(method="POST", headers={}, query_params={}, url=SimpleNamespace(path="/v1/chat/completions")), + skip_budget_checks=skip_budget_checks, + ) + + +@pytest.mark.asyncio +async def test_common_checks_blocks_a_request_whose_group_is_over_budget(): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _common_checks_with_over_budget_group(skip_budget_checks=False) + + assert exc_info.value.entity_id == "tier-a" + + +@pytest.mark.asyncio +async def test_free_model_routes_skip_the_model_access_group_budget_check(): + """skip_budget_checks is how free models stay free; it has to cover this budget too.""" + assert await _common_checks_with_over_budget_group(skip_budget_checks=True) is True diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index 6d362db3c76..03b05bd9d87 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -77,6 +77,7 @@ class MockBatcher: self.litellm_teammembership = _Table("team_membership", self) self.litellm_organizationtable = _Table("org", self) self.litellm_tagtable = _Table("tag", self) + self.litellm_modelaccessgroupbudgettable = _Table("model_access_group", self) self.litellm_endusertable = _Table("enduser", self) async def commit(self): @@ -91,6 +92,7 @@ class MockDB: self.litellm_endusertable = MockTable() self.litellm_organizationtable = MockTable() self.litellm_tagtable = MockTable() + self.litellm_modelaccessgroupbudgettable = MockTable() self.batch_calls: List[Dict[str, Any]] = [] self.batchers: List[MockBatcher] = [] @@ -521,6 +523,7 @@ _LINKED_TABLE_CASES = [ ), ("org", {"budget_id": {"in": ["7d-budget-tier"]}, "spend": {"gt": 0}}), ("tag", {"budget_id": {"in": ["7d-budget-tier"]}, "spend": {"gt": 0}}), + ("model_access_group", {"budget_id": {"in": ["7d-budget-tier"]}, "spend": {"gt": 0}}), ] @@ -1439,13 +1442,19 @@ _INVALIDATION_CASES = [ "spend:tag:tenant-42", {"tag:tenant-42"}, ), + ( + "litellm_modelaccessgroupbudgettable", + type("AccessGroup", (), {"access_group_name": "gpt-4-group"}), + "spend:model_access_group:gpt-4-group", + {"model_access_group:gpt-4-group"}, + ), ] @pytest.mark.parametrize( "table_attr, linked_row, counter_key, cache_keys", _INVALIDATION_CASES, - ids=["team_membership", "key", "org", "tag"], + ids=["team_membership", "key", "org", "tag", "model_access_group"], ) def test_budget_table_reset_invalidates_counters_and_management_cache( reset_budget_job, mock_prisma_client, monkeypatch, table_attr, linked_row, counter_key, cache_keys @@ -1499,6 +1508,102 @@ def test_budget_table_reset_commits_even_when_cache_eviction_fails(reset_budget_ assert mock_prisma_client.db.batchers[0].committed is True +# --------------------------------------------------------------------------- +# Model access group budgets ride the same cascade +# --------------------------------------------------------------------------- + + +def _model_access_group_row(name: str = "gpt-4-group", spend: float = 12.0, budget_id: str = "budget-1"): + """A LiteLLM_ModelAccessGroupBudgetTable row, shaped like prisma hands it back.""" + return type("AccessGroup", (), {"access_group_name": name, "spend": spend, "budget_id": budget_id}) + + +def test_access_group_reset_only_matches_rows_that_have_spend(reset_budget_job, mock_prisma_client, monkeypatch): + """Both the read and the write are filtered to spend > 0 on the due tiers. + + A group sitting at spend 0 has nothing to reset, and a group hanging off a + tier that is not due yet must not be swept along: both are excluded by the + filter, not by anything downstream. + """ + _make_counter_invalidation_job(monkeypatch) + mock_prisma_client.data["budget"] = [_budget_row(budget_id="budget-due", budget_duration="7d")] + mock_prisma_client.db.litellm_modelaccessgroupbudgettable.set_find_many_results( + [_model_access_group_row(budget_id="budget-due")] + ) + + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + expected_where = {"budget_id": {"in": ["budget-due"]}, "spend": {"gt": 0}} + assert mock_prisma_client.db.litellm_modelaccessgroupbudgettable.find_many_calls == [{"where": expected_where}] + writes = _batch_writes(mock_prisma_client, "model_access_group", op="update_many") + assert len(writes) == 1 + assert writes[0]["where"] == expected_where + assert writes[0]["data"] == {"spend": 0} + + +def test_access_groups_are_untouched_when_no_budget_is_due(reset_budget_job, mock_prisma_client, monkeypatch): + """No due tier means the group table is never read, written or evicted.""" + counter_cache = _make_counter_invalidation_job(monkeypatch) + mock_prisma_client.db.litellm_modelaccessgroupbudgettable.set_find_many_results([_model_access_group_row()]) + + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + assert mock_prisma_client.db.litellm_modelaccessgroupbudgettable.find_many_calls == [] + assert _batch_writes(mock_prisma_client, "model_access_group") == [] + counter_cache.in_memory_cache.set_cache.assert_not_called() + counter_cache.user_api_key_cache.async_delete_cache.assert_not_awaited() + + +def test_budget_table_reset_invalidates_every_access_group_not_just_the_first( + reset_budget_job, mock_prisma_client, monkeypatch +): + """When several groups share the expiring tier, all of them are evicted.""" + counter_cache = _make_counter_invalidation_job(monkeypatch) + mock_prisma_client.data["budget"] = [_budget_row(budget_id="budget-1")] + mock_prisma_client.db.litellm_modelaccessgroupbudgettable.set_find_many_results( + [_model_access_group_row(name=name) for name in ("group-a", "group-b", "group-c")] + ) + + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + deleted = {call.kwargs.get("key") for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list} + assert deleted == {"model_access_group:group-a", "model_access_group:group-b", "model_access_group:group-c"} + for name in ("group-a", "group-b", "group-c"): + counter_cache.in_memory_cache.set_cache.assert_any_call(key=f"spend:model_access_group:{name}", value=0.0, ttl=60) + + +def test_budget_cascade_carries_access_group_overage_when_rollover_enabled( + rollover_enabled, reset_budget_job, mock_prisma_client, monkeypatch +): + """A group 5 over the tier cap keeps a spend of 5 in the next window, the + same way a tag or a team member does: over-cap rows are decremented by the + cap, the rest are zeroed, and the counter is seeded with the carried spend.""" + counter_cache = _make_counter_invalidation_job(monkeypatch) + mock_prisma_client.data["budget"] = [_budget_row(budget_id="budget-roll", budget_duration="7d", max_budget=10.0)] + mock_prisma_client.db.litellm_modelaccessgroupbudgettable.set_find_many_results( + [_model_access_group_row(spend=15.0, budget_id="budget-roll")] + ) + + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + writes = _batch_writes(mock_prisma_client, "model_access_group") + assert { + "table": "model_access_group", + "op": "update_many", + "where": {"budget_id": "budget-roll", "spend": {"gt": 10.0}}, + "data": {"spend": {"decrement": 10.0}}, + } in writes + assert { + "table": "model_access_group", + "op": "update_many", + "where": {"budget_id": "budget-roll", "spend": {"gt": 0, "lte": 10.0}}, + "data": {"spend": 0}, + } in writes + assert _replay_spend_writes(writes, 15.0) == 5.0 + assert _replay_spend_writes(writes, 8.0) == 0 + counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:model_access_group:gpt-4-group", value=5.0, ttl=60) + + # --------------------------------------------------------------------------- # Atomicity of the budget-table cascade (LIT-5138) # --------------------------------------------------------------------------- @@ -1611,6 +1716,7 @@ def test_budget_cascade_writes_land_in_a_single_transaction(reset_budget_job, mo ("key", "update_many"), ("org", "update_many"), ("tag", "update_many"), + ("model_access_group", "update_many"), ("enduser", "update_many"), ("budget", "update_many"), } @@ -1646,7 +1752,7 @@ def test_failed_cascade_is_logged_as_a_cascade_failure(monkeypatch): assert mock_exception.call_count == 1 message = mock_exception.call_args.args[0] assert "cascade" in message - for mentioned in ("team member", "enduser", "org", "tag", "budget_reset_at"): + for mentioned in ("team member", "enduser", "org", "tag", "model access group", "budget_reset_at"): assert mentioned in message, f"failure log should mention {mentioned}: {message}" diff --git a/tests/test_litellm/proxy/db/test_model_access_group_spend.py b/tests/test_litellm/proxy/db/test_model_access_group_spend.py new file mode 100644 index 00000000000..79dffc29e7e --- /dev/null +++ b/tests/test_litellm/proxy/db/test_model_access_group_spend.py @@ -0,0 +1,508 @@ +"""Spend accumulation for model access group budgets.""" + +import asyncio +from collections.abc import Mapping, Sequence + +import pytest + +from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY +from litellm.proxy._types import DBSpendUpdateTransactions, Litellm_EntityType, SpendUpdateQueueItem +from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter, debitable_model_access_groups +from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import DailySpendUpdateQueue +from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer +from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue +from litellm.proxy.spend_tracking.spend_tracking_utils import get_request_model_access_groups + + +class _FakeRouter: + """Deployment lookup returning the access groups each deployment declares.""" + + def __init__(self, deployments: Mapping[str, Sequence[str] | None]) -> None: + self._deployments = deployments + + def get_model_info(self, id: str) -> dict | None: + if id not in self._deployments: + return None + declared = self._deployments[id] + model_info: dict = {"id": id} + if declared is not None: + model_info["access_groups"] = list(declared) + return {"model_name": "some-model", "model_info": model_info} + + +class _FakeBatchTable: + def __init__(self) -> None: + self.calls: list[tuple[dict, dict]] = [] + + def update_many(self, where: dict, data: dict) -> None: + self.calls.append((where, data)) + + +class _FakeBatcher: + def __init__(self) -> None: + self.tables: dict[str, _FakeBatchTable] = {} + + def __getattr__(self, name: str) -> _FakeBatchTable: + return self.tables.setdefault(name, _FakeBatchTable()) + + +class _FakeBatchManager: + def __init__(self, batcher: _FakeBatcher) -> None: + self._batcher = batcher + + async def __aenter__(self) -> _FakeBatcher: + return self._batcher + + async def __aexit__(self, *exc_info: object) -> bool: + return False + + +class _FakeTransaction: + def __init__(self, batcher: _FakeBatcher) -> None: + self._batcher = batcher + + def batch_(self) -> _FakeBatchManager: + return _FakeBatchManager(self._batcher) + + async def __aenter__(self) -> "_FakeTransaction": + return self + + async def __aexit__(self, *exc_info: object) -> bool: + return False + + +class _FakeDb: + def __init__(self, batcher: _FakeBatcher) -> None: + self._batcher = batcher + + def tx(self, timeout: object = None) -> _FakeTransaction: + return _FakeTransaction(self._batcher) + + +class _FakePrismaClient: + def __init__(self) -> None: + self.batcher = _FakeBatcher() + self.db = _FakeDb(self.batcher) + + +def _empty_transactions(**overrides: dict[str, float]) -> DBSpendUpdateTransactions: + return DBSpendUpdateTransactions( + user_list_transactions=overrides.get("user_list_transactions", {}), + end_user_list_transactions=overrides.get("end_user_list_transactions", {}), + key_list_transactions=overrides.get("key_list_transactions", {}), + team_list_transactions=overrides.get("team_list_transactions", {}), + team_member_list_transactions=overrides.get("team_member_list_transactions", {}), + org_list_transactions=overrides.get("org_list_transactions", {}), + tag_list_transactions=overrides.get("tag_list_transactions", {}), + agent_list_transactions=overrides.get("agent_list_transactions", {}), + model_access_group_list_transactions=overrides.get("model_access_group_list_transactions", {}), + ) + + +async def _drain(queue: SpendUpdateQueue) -> list[SpendUpdateQueueItem]: + return await queue.flush_all_updates_from_in_memory_queue() + + +# --- enqueue --------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_single_matched_group_enqueues_one_item_with_full_cost(): + writer = DBSpendUpdateWriter() + + await writer._update_model_access_group_db( + response_cost=0.42, + request_model_access_groups=["premium-pool"], + served_model_id="deployment-1", + prisma_client=object(), + router=_FakeRouter({"deployment-1": ["premium-pool"]}), + ) + + updates = await _drain(writer.spend_update_queue) + assert updates == [ + SpendUpdateQueueItem( + entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP, + entity_id="premium-pool", + response_cost=0.42, + ) + ] + + +@pytest.mark.asyncio +async def test_every_matched_group_is_charged_the_full_cost_not_a_split(): + writer = DBSpendUpdateWriter() + + await writer._update_model_access_group_db( + response_cost=0.30, + request_model_access_groups=["pool-a", "pool-b", "pool-c"], + served_model_id="deployment-1", + prisma_client=object(), + router=_FakeRouter({"deployment-1": ["pool-a", "pool-b", "pool-c"]}), + ) + + updates = await _drain(writer.spend_update_queue) + assert [update["entity_id"] for update in updates] == ["pool-a", "pool-b", "pool-c"] + assert [update["response_cost"] for update in updates] == [0.30, 0.30, 0.30] + assert {update["entity_type"] for update in updates} == {Litellm_EntityType.MODEL_ACCESS_GROUP} + + +@pytest.mark.parametrize("attributed", [None, [], ()]) +@pytest.mark.asyncio +async def test_no_attributed_groups_enqueues_nothing(attributed): + writer = DBSpendUpdateWriter() + + await writer._update_model_access_group_db( + response_cost=1.0, + request_model_access_groups=attributed, + served_model_id="deployment-1", + prisma_client=object(), + router=_FakeRouter({"deployment-1": ["premium-pool"]}), + ) + + assert await _drain(writer.spend_update_queue) == [] + + +@pytest.mark.asyncio +async def test_no_prisma_client_enqueues_nothing(): + writer = DBSpendUpdateWriter() + + await writer._update_model_access_group_db( + response_cost=1.0, + request_model_access_groups=["premium-pool"], + served_model_id="deployment-1", + prisma_client=None, + router=_FakeRouter({"deployment-1": ["premium-pool"]}), + ) + + assert await _drain(writer.spend_update_queue) == [] + + +@pytest.mark.asyncio +async def test_group_outside_the_attributed_set_is_never_debited(): + """The served deployment also sits in a pool auth never attributed; that pool stays untouched.""" + writer = DBSpendUpdateWriter() + + await writer._update_model_access_group_db( + response_cost=0.10, + request_model_access_groups=["premium-pool"], + served_model_id="deployment-1", + prisma_client=object(), + router=_FakeRouter({"deployment-1": ["premium-pool", "unattributed-pool"]}), + ) + + updates = await _drain(writer.spend_update_queue) + assert [update["entity_id"] for update in updates] == ["premium-pool"] + + +# --- fallback guard -------------------------------------------------------- + + +def test_fallback_to_a_model_in_another_pool_debits_nothing(): + assert ( + debitable_model_access_groups( + attributed=["premium-pool"], + served_model_id="fallback-deployment", + router=_FakeRouter({"fallback-deployment": ["cheap-pool"]}), + ) + == () + ) + + +def test_fallback_to_a_model_in_no_pool_debits_nothing(): + assert ( + debitable_model_access_groups( + attributed=["premium-pool"], + served_model_id="fallback-deployment", + router=_FakeRouter({"fallback-deployment": None}), + ) + == () + ) + + +def test_attributed_set_stands_when_the_served_deployment_is_unknown(): + assert debitable_model_access_groups( + attributed=["premium-pool"], + served_model_id="not-in-router", + router=_FakeRouter({"deployment-1": ["premium-pool"]}), + ) == ("premium-pool",) + + +def test_attributed_set_stands_without_a_router(): + assert debitable_model_access_groups( + attributed=["premium-pool", "premium-pool"], + served_model_id="deployment-1", + router=None, + ) == ("premium-pool",) + + +def test_partial_overlap_keeps_only_the_intersection(): + assert debitable_model_access_groups( + attributed=["pool-a", "pool-b"], + served_model_id="deployment-1", + router=_FakeRouter({"deployment-1": ["pool-b", "pool-c"]}), + ) == ("pool-b",) + + +def test_only_real_group_names_ever_become_entity_ids(): + """Whatever shape the attributed set arrives in, an empty or non-string name never reaches the queue.""" + assert debitable_model_access_groups( + attributed=["pool-a", "", "pool-a", None, 7], + served_model_id=None, + router=None, + ) == ("pool-a",) + + +# --- metadata extraction --------------------------------------------------- + + +def test_access_groups_read_from_request_metadata(): + kwargs = {"litellm_params": {"metadata": {MODEL_ACCESS_GROUP_METADATA_KEY: ["pool-a", "pool-b", "pool-a"]}}} + assert get_request_model_access_groups(kwargs) == ("pool-a", "pool-b") + + +def test_access_groups_read_from_litellm_metadata(): + kwargs = {"litellm_params": {"litellm_metadata": {MODEL_ACCESS_GROUP_METADATA_KEY: ["pool-a"]}}} + assert get_request_model_access_groups(kwargs) == ("pool-a",) + + +def test_standard_logging_payload_wins_over_metadata(): + kwargs = { + "litellm_params": {"metadata": {MODEL_ACCESS_GROUP_METADATA_KEY: ["from-metadata"]}}, + "standard_logging_object": {"request_model_access_groups": ["from-payload"]}, + } + assert get_request_model_access_groups(kwargs) == ("from-payload",) + + +def test_metadata_is_used_when_the_logging_payload_carries_no_groups(): + kwargs = { + "litellm_params": {"metadata": {MODEL_ACCESS_GROUP_METADATA_KEY: ["from-metadata"]}}, + "standard_logging_object": {"request_model_access_groups": []}, + } + assert get_request_model_access_groups(kwargs) == ("from-metadata",) + + +@pytest.mark.parametrize("stamped", ["pool-a", 7, {"pool-a": 1}]) +def test_non_list_access_group_metadata_is_ignored(stamped): + kwargs = {"litellm_params": {"metadata": {MODEL_ACCESS_GROUP_METADATA_KEY: stamped}}} + assert get_request_model_access_groups(kwargs) == () + + +def test_key_absent_from_metadata_yields_no_groups(): + """The chat path only stamps the key when something matched, so absent must mean nothing to debit.""" + kwargs = {"litellm_params": {"metadata": {"user_api_key_user_id": "u-1"}}} + assert get_request_model_access_groups(kwargs) == () + + +def test_explicit_none_yields_no_groups(): + """The pass-through path stamps the key unconditionally, so it can be present and None.""" + kwargs = {"litellm_params": {"metadata": {MODEL_ACCESS_GROUP_METADATA_KEY: None}}} + assert get_request_model_access_groups(kwargs) == () + + +@pytest.mark.parametrize( + "metadata", + [ + {"user_api_key_user_id": "u-1"}, + {MODEL_ACCESS_GROUP_METADATA_KEY: None}, + ], + ids=["key-absent", "key-present-but-none"], +) +@pytest.mark.asyncio +async def test_neither_absent_nor_none_metadata_debits_anything(metadata): + writer = DBSpendUpdateWriter() + + await writer._update_model_access_group_db( + response_cost=0.5, + request_model_access_groups=get_request_model_access_groups({"litellm_params": {"metadata": metadata}}), + served_model_id="deployment-1", + prisma_client=object(), + router=_FakeRouter({"deployment-1": ["premium-pool"]}), + ) + + assert await _drain(writer.spend_update_queue) == [] + + +def test_detached_sub_call_falls_back_to_the_auth_object(): + """Sub-calls inherit only the identity keys, so the groups come off user_api_key_auth there.""" + + class _Auth: + matched_model_access_groups = ["premium-pool"] + + kwargs = {"litellm_params": {"metadata": {"user_api_key_auth": _Auth()}}} + assert get_request_model_access_groups(kwargs) == ("premium-pool",) + + +def test_stamped_metadata_wins_over_the_auth_object(): + class _Auth: + matched_model_access_groups = ["stale-pool"] + + kwargs = { + "litellm_params": { + "metadata": { + MODEL_ACCESS_GROUP_METADATA_KEY: ["fresh-pool"], + "user_api_key_auth": _Auth(), + } + } + } + assert get_request_model_access_groups(kwargs) == ("fresh-pool",) + + +def test_auth_object_without_matched_groups_yields_no_groups(): + class _Auth: + matched_model_access_groups = None + + kwargs = {"litellm_params": {"metadata": {"user_api_key_auth": _Auth()}}} + assert get_request_model_access_groups(kwargs) == () + + +def test_non_string_entries_are_dropped(): + kwargs = {"litellm_params": {"metadata": {MODEL_ACCESS_GROUP_METADATA_KEY: ["pool-a", None, "", 3]}}} + assert get_request_model_access_groups(kwargs) == ("pool-a",) + + +def test_missing_metadata_yields_no_groups(): + assert get_request_model_access_groups(None) == () + assert get_request_model_access_groups({}) == () + assert get_request_model_access_groups({"litellm_params": {}}) == () + + +# --- queue bucketing and redis round trip ---------------------------------- + + +def test_access_group_updates_aggregate_into_their_own_bucket(): + queue = SpendUpdateQueue() + + transactions = queue.get_aggregated_db_spend_update_transactions( + [ + SpendUpdateQueueItem( + entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP, entity_id="pool-a", response_cost=0.1 + ), + SpendUpdateQueueItem( + entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP, entity_id="pool-a", response_cost=0.2 + ), + SpendUpdateQueueItem( + entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP, entity_id="pool-b", response_cost=0.5 + ), + SpendUpdateQueueItem(entity_type=Litellm_EntityType.TAG, entity_id="pool-a", response_cost=9.0), + ] + ) + + assert transactions["model_access_group_list_transactions"] == {"pool-a": pytest.approx(0.3), "pool-b": 0.5} + assert transactions["tag_list_transactions"] == {"pool-a": 9.0} + + +def test_access_group_transactions_survive_the_redis_buffer_merge(): + merged = RedisUpdateBuffer._combine_list_of_transactions( + [ + _empty_transactions(model_access_group_list_transactions={"pool-a": 0.25}), + _empty_transactions(model_access_group_list_transactions={"pool-a": 0.25, "pool-b": 1.0}), + ] + ) + + assert merged["model_access_group_list_transactions"] == {"pool-a": 0.5, "pool-b": 1.0} + + +@pytest.mark.asyncio +async def test_redis_buffer_requeues_access_group_transactions_as_queue_items(): + queue = SpendUpdateQueue() + daily_queue = DailySpendUpdateQueue() + + await RedisUpdateBuffer._restore_spend_updates_to_in_memory_queues( + db_spend_update_transactions=_empty_transactions(model_access_group_list_transactions={"pool-a": 0.75}), + daily_spend_update_transactions=None, + daily_team_spend_update_transactions=None, + daily_org_spend_update_transactions=None, + daily_end_user_spend_update_transactions=None, + daily_agent_spend_update_transactions=None, + spend_update_queue=queue, + daily_spend_update_queue=daily_queue, + daily_team_spend_update_queue=daily_queue, + daily_org_spend_update_queue=daily_queue, + daily_end_user_spend_update_queue=daily_queue, + daily_agent_spend_update_queue=daily_queue, + ) + + updates = await _drain(queue) + assert updates == [ + SpendUpdateQueueItem( + entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP, + entity_id="pool-a", + response_cost=0.75, + ) + ] + + +# --- flush to postgres ----------------------------------------------------- + + +@pytest.mark.asyncio +async def test_commit_increments_spend_on_the_model_access_group_budget_table(): + prisma_client = _FakePrismaClient() + + await DBSpendUpdateWriter()._commit_spend_updates_to_db( + prisma_client=prisma_client, + n_retry_times=0, + proxy_logging_obj=None, + db_spend_update_transactions=_empty_transactions( + model_access_group_list_transactions={"pool-b": 0.5, "pool-a": 0.25} + ), + ) + + assert prisma_client.batcher.tables["litellm_modelaccessgroupbudgettable"].calls == [ + ({"access_group_name": "pool-a"}, {"spend": {"increment": 0.25}}), + ({"access_group_name": "pool-b"}, {"spend": {"increment": 0.5}}), + ] + assert "litellm_tagtable" not in prisma_client.batcher.tables + + +# --- end-to-end through the batched fan-out -------------------------------- + + +@pytest.mark.asyncio +async def test_batch_database_updates_enqueues_access_group_spend(): + writer = DBSpendUpdateWriter() + + await writer._batch_database_updates( + response_cost=0.15, + user_id=None, + hashed_token=None, + team_id=None, + org_id=None, + end_user_id=None, + prisma_client=object(), + litellm_proxy_budget_name=None, + payload={"model_id": "deployment-1", "spend": 0.15}, + request_model_access_groups=("pool-a", "pool-b"), + ) + await asyncio.sleep(0) + + access_group_updates = [ + update + for update in await _drain(writer.spend_update_queue) + if update["entity_type"] is Litellm_EntityType.MODEL_ACCESS_GROUP + ] + assert [(update["entity_id"], update["response_cost"]) for update in access_group_updates] == [ + ("pool-a", 0.15), + ("pool-b", 0.15), + ] + + +@pytest.mark.asyncio +async def test_batch_database_updates_enqueues_nothing_without_access_groups(): + writer = DBSpendUpdateWriter() + + await writer._batch_database_updates( + response_cost=0.15, + user_id=None, + hashed_token=None, + team_id=None, + org_id=None, + end_user_id=None, + prisma_client=object(), + litellm_proxy_budget_name=None, + payload={"model_id": "deployment-1", "spend": 0.15}, + ) + await asyncio.sleep(0) + + updates = await _drain(writer.spend_update_queue) + assert [update for update in updates if update["entity_type"] is Litellm_EntityType.MODEL_ACCESS_GROUP] == [] diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 8fd8730d409..2a540f4f522 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -1,14 +1,14 @@ -import pytest - - from datetime import datetime from unittest.mock import AsyncMock, MagicMock, patch +import pytest + +from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.proxy_track_cost_callback import ( - _ProxyDBLogger, _get_budget_reservation_from_metadata, + _ProxyDBLogger, _should_track_cost_callback, _update_database_and_spend_counters, ) @@ -589,6 +589,7 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda response_cost=0.2, budget_reservation=budget_reservation, request_tags=["tag-a"], + model_access_groups=("premium",), ) proxy_logging_obj.db_spend_update_writer.update_database.assert_awaited_once() @@ -603,6 +604,7 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda tags=["tag-a"], request_id="chatcmpl-abc123", request_started_at=start_time, + model_access_groups=("premium",), ) @@ -1935,3 +1937,113 @@ async def test_update_database_and_spend_counters_forwards_a_missing_request_id_ ) assert increment_spend_counters.await_args.kwargs["request_id"] is None + + +class _FakeDeploymentLookup: + """Deployment lookup returning the access groups each deployment declares.""" + + def __init__(self, deployments): + self._deployments = deployments + + def get_model_info(self, id): + if id not in self._deployments: + return None + return {"model_name": "premium-haiku", "model_info": {"id": id, "access_groups": list(self._deployments[id])}} + + +def _model_access_group_kwargs(granted, served_model_id=None): + metadata = {"user_api_key": "hashed-key", "user_api_key_user_id": "user-1"} + if granted is not None: + metadata[MODEL_ACCESS_GROUP_METADATA_KEY] = list(granted) + return { + "call_type": "acompletion", + "model": "premium-haiku", + "litellm_call_id": "test-call-id", + "litellm_params": {"metadata": metadata}, + "stream": False, + "standard_logging_object": {"response_cost": 0.25, "request_tags": None, "model_id": served_model_id}, + } + + +async def _groups_charged_by_the_callback(kwargs, deployments=None): + """The groups the callback hands the spend counters for one request. + + The callback resolves ``proxy_logging_obj`` and the router by importing them off + ``proxy_server`` inside its own body, so there is no seam to inject either through. + """ + logger = _ProxyDBLogger() + with ( + patch( # test-quality-ok: callback imports proxy_logging_obj off proxy_server in its body, no seam + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as mock_proxy_logging, + patch( # test-quality-ok: the arguments to this call are the boundary under test + "litellm.proxy.hooks.proxy_track_cost_callback._update_database_and_spend_counters", + new=AsyncMock(), + ) as mock_update, + patch( # test-quality-ok: llm_router is a proxy_server global the callback reads lazily, no seam + "litellm.proxy.proxy_server.llm_router", new=_FakeDeploymentLookup(deployments or {}) + ), + ): + mock_proxy_logging.failed_tracking_alert = AsyncMock() + mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() + mock_proxy_logging.db_spend_update_writer = MagicMock() + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + + await logger._PROXY_track_cost_callback( + kwargs=kwargs, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + return mock_update.await_args.kwargs["model_access_groups"] + + +@pytest.mark.asyncio +async def test_track_cost_callback_charges_the_model_access_groups_auth_stamped(): + """Auth stamps the matched groups onto request metadata; the callback has to carry them through. + + Without this hop nothing writes ``spend:model_access_group:*`` on the normal path, so with + reservations disabled the budget check reads a counter no one maintains. + """ + charged = await _groups_charged_by_the_callback( + kwargs=_model_access_group_kwargs(granted=["premium", "starter"]), + ) + + assert charged == ("premium", "starter") + + +@pytest.mark.asyncio +async def test_track_cost_callback_charges_no_model_access_group_when_none_were_stamped(): + """A request no budgeted group authorized must not debit anything.""" + charged = await _groups_charged_by_the_callback( + kwargs=_model_access_group_kwargs(granted=None), + ) + + assert charged == () + + +@pytest.mark.asyncio +async def test_spend_counters_only_debit_the_group_the_served_deployment_belongs_to(): + """A caller granted two pools that both cover the model group only draws down the pool that served. + + The database writer already narrows by served deployment, so passing the unnarrowed set to the + live counters let one request block a pool the persisted spend never debited. + """ + charged = await _groups_charged_by_the_callback( + kwargs=_model_access_group_kwargs(granted=["premium", "tier0"], served_model_id="deployment-premium"), + deployments={"deployment-premium": ["premium"], "deployment-tier0": ["tier0"]}, + ) + + assert charged == ("premium",) + + +@pytest.mark.asyncio +async def test_spend_counters_keep_every_granted_group_when_the_deployment_is_unknown(): + """An unidentifiable deployment leaves the auth-time set standing, so nothing silently stops billing.""" + charged = await _groups_charged_by_the_callback( + kwargs=_model_access_group_kwargs(granted=["premium", "tier0"], served_model_id="deployment-gone"), + deployments={"deployment-premium": ["premium"]}, + ) + + assert charged == ("premium", "tier0") diff --git a/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py b/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py index db0557cfbf0..de5fc96c7c3 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py +++ b/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py @@ -2,6 +2,10 @@ Test access group management endpoints """ +from contextlib import contextmanager +from dataclasses import dataclass +from datetime import datetime +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -449,6 +453,7 @@ async def test_delete_access_group_ignores_models_that_were_already_dead(): mock_prisma = MagicMock() mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[deploy_broken]) mock_prisma.db.litellm_proxymodeltable.update = AsyncMock() + mock_prisma.db.litellm_modelaccessgroupbudgettable.delete = AsyncMock(return_value=None) from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( @@ -468,6 +473,7 @@ async def test_delete_access_group_ignores_models_that_were_already_dead(): response = await delete_access_group( access_group="doomed-group", user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + auth_cache=_FakeAuthCache(), ) assert response.models_updated == 1 @@ -568,3 +574,639 @@ async def test_create_access_group_model_missing_everywhere_still_400s(): assert exc_info.value.status_code == 400 assert model_name in str(exc_info.value.detail) + +@dataclass +class _FakeBudgetRow: + budget_id: str + max_budget: float | None = None + soft_budget: float | None = None + budget_duration: str | None = None + budget_reset_at: datetime | None = None + + +@dataclass +class _FakeAccessGroupBudgetRow: + access_group_name: str + budget_id: str | None = None + spend: float = 0.0 + litellm_budget_table: _FakeBudgetRow | None = None + + +@dataclass +class _FakeDeployment: + model_id: str + model_name: str + model_info: dict + + +class _FakeBudgetTable: + """Stands in for litellm_budgettable so a test can see whether a budget row was created, + updated in place, or left orphaned.""" + + def __init__(self, journal: list[str]) -> None: + self.journal = journal + self.rows: dict[str, _FakeBudgetRow] = {} + self.create_calls: list[dict] = [] + self.update_calls: list[tuple[str, dict]] = [] + self.deleted_ids: list[str] = [] + self._sequence = 0 + + async def create(self, data, include=None): + self._sequence += 1 + budget_id = str(data.get("budget_id") or f"budget-{self._sequence}") + row = _FakeBudgetRow( + budget_id=budget_id, + max_budget=data.get("max_budget"), + soft_budget=data.get("soft_budget"), + budget_duration=data.get("budget_duration"), + ) + self.rows[budget_id] = row + self.create_calls.append(dict(data)) + self.journal.append(f"budget_table.create:{budget_id}") + return row + + async def update(self, where, data, include=None): + budget_id = where["budget_id"] + self.update_calls.append((budget_id, dict(data))) + self.journal.append(f"budget_table.update:{budget_id}") + row = self.rows.get(budget_id) + if row is None: + return None + for field_name in ("max_budget", "soft_budget", "budget_duration"): + if data.get(field_name) is not None: + setattr(row, field_name, data[field_name]) + return row + + async def delete(self, where, include=None): + budget_id = where["budget_id"] + self.journal.append(f"budget_table.delete:{budget_id}") + self.deleted_ids.append(budget_id) + return self.rows.pop(budget_id, None) + + +class _FakeAccessGroupBudgetTable: + """Stands in for litellm_modelaccessgroupbudgettable, resolving `include` against the fake + budget table the way prisma resolves the relation.""" + + def __init__(self, journal: list[str], budget_table: _FakeBudgetTable) -> None: + self.journal = journal + self.budget_table = budget_table + self.rows: dict[str, _FakeAccessGroupBudgetRow] = {} + self.upsert_calls: list[dict] = [] + + def _resolve(self, row, include): + if row is None: + return None + row.litellm_budget_table = ( + self.budget_table.rows.get(row.budget_id) if include and row.budget_id is not None else None + ) + return row + + async def find_unique(self, where, include=None): + return self._resolve(self.rows.get(where["access_group_name"]), include) + + async def upsert(self, where, data, include=None): + access_group_name = where["access_group_name"] + self.upsert_calls.append(dict(data)) + self.journal.append(f"access_group_budget.upsert:{access_group_name}") + existing = self.rows.get(access_group_name) + payload = data["update"] if existing is not None else data["create"] + row = existing or _FakeAccessGroupBudgetRow(access_group_name=access_group_name) + row.budget_id = payload.get("budget_id") + self.rows[access_group_name] = row + return self._resolve(row, include) + + async def delete(self, where, include=None): + access_group_name = where["access_group_name"] + self.journal.append(f"access_group_budget.delete:{access_group_name}") + return self.rows.pop(access_group_name, None) + + +class _FakeModelTable: + def __init__(self, journal: list[str], deployments) -> None: + self.journal = journal + self.deployments = list(deployments) + self.updates: list[tuple[dict, dict]] = [] + + async def find_many(self, where=None, **kwargs): + return list(self.deployments) + + async def find_unique(self, where, include=None): + return next((d for d in self.deployments if d.model_id == where["model_id"]), None) + + async def update(self, where, data, include=None): + self.journal.append(f"model_table.update:{where['model_id']}") + self.updates.append((dict(where), dict(data))) + return None + + +class _FakePrismaClient: + def __init__(self, journal: list[str], deployments=()) -> None: + self.budget_table = _FakeBudgetTable(journal) + self.access_group_budget_table = _FakeAccessGroupBudgetTable(journal, self.budget_table) + self.model_table = _FakeModelTable(journal, deployments) + self.db = SimpleNamespace( + litellm_budgettable=self.budget_table, + litellm_modelaccessgroupbudgettable=self.access_group_budget_table, + litellm_proxymodeltable=self.model_table, + ) + + def jsonify_object(self, data): + return dict(data) + + +class _FakeAuthCache: + """Spy for the auth cache the endpoints evict through. Injected into the endpoint rather than + patched over the proxy_server global, so dropping the eviction call fails a test.""" + + def __init__(self, journal: list[str] | None = None) -> None: + self.journal = journal if journal is not None else [] + self.deleted_keys: list[str] = [] + + async def async_delete_cache(self, key): + self.deleted_keys.append(key) + self.journal.append(f"auth_cache.delete:{key}") + + +def _admin(): + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + + return UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + + +def _deployment(model_id="deploy-1", model_name="gpt-4o", access_groups=("prod-models",)): + return _FakeDeployment( + model_id=model_id, + model_name=model_name, + model_info={"access_groups": list(access_groups)}, + ) + + +def _seed_budget(prisma, access_group, spend=0.0, budget_id="budget-seed", **budget_fields): + prisma.budget_table.rows[budget_id] = _FakeBudgetRow(budget_id=budget_id, **budget_fields) + prisma.access_group_budget_table.rows[access_group] = _FakeAccessGroupBudgetRow( + access_group_name=access_group, + budget_id=budget_id, + spend=spend, + ) + + +@contextmanager +def _proxy(prisma): + with patch( # test-quality-ok: the endpoints import proxy_server.prisma_client themselves; no parameter to inject + "litellm.proxy.proxy_server.prisma_client", prisma + ): + yield + + +@contextmanager +def _proxy_with_stubbed_reload(prisma): + """delete_access_group finishes by reloading the router and judging what it serves afterwards. + Both collaborators it reaches for there are module globals it imports itself, so a fake can only + get in by patching them; auth_cache and prisma are the ones with a real seam.""" + never_served_router = MagicMock() + never_served_router.get_model_ids.return_value = [] + with ( + _proxy(prisma), + patch( # test-quality-ok: live_model_ids_snapshot() reads the llm_router global; the endpoint takes no router + "litellm.proxy.proxy_server.llm_router", never_served_router + ), + patch( # test-quality-ok: the endpoint calls its module-level clear_cache import; there is no parameter for it + "litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache", + new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), + ), + ): + yield + + +def _eviction_journal(access_group): + """Both auth cache keys, in the order a write path has to evict them.""" + from litellm.proxy.common_utils.user_api_key_cache import ( + model_access_group_cache_key, + model_access_group_registry_cache_key, + ) + + return [ + f"auth_cache.delete:{model_access_group_cache_key(access_group)}", + f"auth_cache.delete:{model_access_group_registry_cache_key()}", + ] + + +def _assert_evicted_after_write(journal, access_group, write_entry): + """Exactly the two keys, in order, after the DB write. Deliberately not a tail slice: what + has to hold is that the eviction follows the write, not that nothing follows the eviction.""" + evictions = [entry for entry in journal if entry.startswith("auth_cache.delete:")] + assert evictions == _eviction_journal(access_group) + assert journal.index(write_entry) < journal.index(evictions[0]) + + +@pytest.mark.asyncio +async def test_put_access_group_budget_creates_the_row_and_its_budget(): + """First PUT has to create both halves: the budget row it links, and the access group row + that carries the link and the shared spend.""" + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + set_access_group_budget, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + AccessGroupBudgetRequest, + ) + + prisma = _FakePrismaClient([], deployments=[_deployment()]) + cache = _FakeAuthCache() + + with _proxy(prisma): + response = await set_access_group_budget( + access_group="prod-models", + data=AccessGroupBudgetRequest(max_budget=100.0, soft_budget=80.0, budget_duration="30d"), + user_api_key_dict=_admin(), + auth_cache=cache, + ) + + assert response.access_group == "prod-models" + assert response.spend == 0.0 + assert response.budget is not None + assert response.budget.max_budget == 100.0 + assert response.budget.soft_budget == 80.0 + assert response.budget.budget_duration == "30d" + assert len(prisma.budget_table.create_calls) == 1 + assert prisma.access_group_budget_table.rows["prod-models"].budget_id == response.budget.budget_id + + +@pytest.mark.asyncio +async def test_second_put_replaces_the_budget_instead_of_creating_another(): + """PUT is idempotent: a second call must update the budget already linked to the group, + not leave a second budget row (and a second group row) behind.""" + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + set_access_group_budget, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + AccessGroupBudgetRequest, + ) + + prisma = _FakePrismaClient([], deployments=[_deployment()]) + cache = _FakeAuthCache() + + with _proxy(prisma): + first = await set_access_group_budget( + access_group="prod-models", + data=AccessGroupBudgetRequest(max_budget=100.0), + user_api_key_dict=_admin(), + auth_cache=cache, + ) + second = await set_access_group_budget( + access_group="prod-models", + data=AccessGroupBudgetRequest(max_budget=250.0), + user_api_key_dict=_admin(), + auth_cache=cache, + ) + + assert first.budget is not None and second.budget is not None + assert second.budget.budget_id == first.budget.budget_id + assert second.budget.max_budget == 250.0 + assert len(prisma.budget_table.create_calls) == 1 + assert len(prisma.budget_table.rows) == 1 + assert len(prisma.access_group_budget_table.rows) == 1 + assert prisma.budget_table.update_calls[-1][0] == first.budget.budget_id + + +@pytest.mark.asyncio +async def test_put_access_group_budget_links_an_existing_budget_without_creating_one(): + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + set_access_group_budget, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + AccessGroupBudgetRequest, + ) + + prisma = _FakePrismaClient([], deployments=[_deployment()]) + prisma.budget_table.rows["shared-budget"] = _FakeBudgetRow(budget_id="shared-budget", max_budget=7.0) + cache = _FakeAuthCache() + + with _proxy(prisma): + response = await set_access_group_budget( + access_group="prod-models", + data=AccessGroupBudgetRequest(budget_id="shared-budget"), + user_api_key_dict=_admin(), + auth_cache=cache, + ) + + assert prisma.budget_table.create_calls == [] + assert response.budget is not None + assert response.budget.budget_id == "shared-budget" + assert response.budget.max_budget == 7.0 + assert prisma.access_group_budget_table.rows["prod-models"].budget_id == "shared-budget" + + +@pytest.mark.asyncio +async def test_put_access_group_budget_rejects_an_empty_body(): + """An empty PUT would register the group as budgeted while enforcing nothing.""" + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + set_access_group_budget, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + AccessGroupBudgetRequest, + ) + + prisma = _FakePrismaClient([], deployments=[_deployment()]) + cache = _FakeAuthCache() + + with _proxy(prisma), pytest.raises(HTTPException) as exc_info: + await set_access_group_budget( + access_group="prod-models", + data=AccessGroupBudgetRequest(), + user_api_key_dict=_admin(), + auth_cache=cache, + ) + + assert exc_info.value.status_code == 400 + assert prisma.access_group_budget_table.rows == {} + assert cache.deleted_keys == [] + + +@pytest.mark.asyncio +async def test_put_access_group_budget_rejects_an_unparseable_duration(): + """An unparseable duration can only be discovered by the reset job, long after the write.""" + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + set_access_group_budget, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + AccessGroupBudgetRequest, + ) + + prisma = _FakePrismaClient([], deployments=[_deployment()]) + cache = _FakeAuthCache() + + with _proxy(prisma), pytest.raises(HTTPException) as exc_info: + await set_access_group_budget( + access_group="prod-models", + data=AccessGroupBudgetRequest(max_budget=10.0, budget_duration="every other tuesday"), + user_api_key_dict=_admin(), + auth_cache=cache, + ) + + assert exc_info.value.status_code == 400 + assert prisma.budget_table.create_calls == [] + assert prisma.access_group_budget_table.rows == {} + + +def test_access_group_budget_request_rejects_rate_limit_fields(): + """tpm/rpm/max_parallel_requests are not enforced per access group, so accepting them would + promise rate limiting that never happens.""" + from pydantic import ValidationError + + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + AccessGroupBudgetRequest, + ) + + for unsupported in ({"tpm_limit": 10}, {"rpm_limit": 10}, {"max_parallel_requests": 10}): + with pytest.raises(ValidationError): + AccessGroupBudgetRequest(max_budget=1.0, **unsupported) + + +@pytest.mark.asyncio +async def test_get_access_group_budget_returns_the_budget_and_the_shared_spend(): + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + get_access_group_budget, + ) + + prisma = _FakePrismaClient([], deployments=[_deployment()]) + _seed_budget(prisma, "prod-models", spend=42.5, max_budget=100.0, budget_duration="30d") + + with _proxy(prisma): + response = await get_access_group_budget(access_group="prod-models") + + assert response.access_group == "prod-models" + assert response.spend == 42.5 + assert response.budget is not None + assert response.budget.max_budget == 100.0 + assert response.budget.budget_duration == "30d" + + +@pytest.mark.asyncio +async def test_get_access_group_budget_on_a_budgetless_group_is_200_not_404(): + """A real group that simply has no budget is not an error; only an unknown group is.""" + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + get_access_group_budget, + ) + + prisma = _FakePrismaClient([], deployments=[_deployment()]) + + with _proxy(prisma): + response = await get_access_group_budget(access_group="prod-models") + + assert response.spend == 0.0 + assert response.budget is None + + +@pytest.mark.asyncio +async def test_access_group_budget_routes_404_on_an_unknown_group(): + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + delete_access_group_budget, + get_access_group_budget, + set_access_group_budget, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + AccessGroupBudgetRequest, + ) + + prisma = _FakePrismaClient([], deployments=[_deployment()]) + cache = _FakeAuthCache() + admin = _admin() + + calls = ( + lambda: get_access_group_budget(access_group="ghost-group"), + lambda: set_access_group_budget( + access_group="ghost-group", + data=AccessGroupBudgetRequest(max_budget=1.0), + user_api_key_dict=admin, + auth_cache=cache, + ), + lambda: delete_access_group_budget(access_group="ghost-group", auth_cache=cache), + ) + + with _proxy(prisma): + for make_call in calls: + with pytest.raises(HTTPException) as exc_info: + await make_call() + assert exc_info.value.status_code == 404 + + assert prisma.budget_table.create_calls == [] + assert prisma.access_group_budget_table.rows == {} + + +@pytest.mark.asyncio +async def test_delete_access_group_budget_drops_the_row_and_spares_the_shared_budget(): + """The group row goes; the LiteLLM_BudgetTable row it linked survives, as /tag/delete leaves a + tag's. That row can be shared, so deleting it would be data loss for whatever else points at it.""" + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + delete_access_group_budget, + ) + + prisma = _FakePrismaClient([], deployments=[_deployment()]) + _seed_budget(prisma, "prod-models", spend=12.0, max_budget=100.0) + cache = _FakeAuthCache() + + with _proxy(prisma): + response = await delete_access_group_budget(access_group="prod-models", auth_cache=cache) + + assert response.budget_deleted is True + assert prisma.access_group_budget_table.rows == {} + assert prisma.budget_table.deleted_ids == [] + assert prisma.budget_table.rows["budget-seed"].max_budget == 100.0 + + +@pytest.mark.asyncio +async def test_delete_access_group_budget_on_a_budgetless_group_still_evicts(): + """budget_deleted is False, but the group can still be sitting in the cached registry of + budgeted groups, so the eviction has to run whether or not a row was there to drop.""" + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + delete_access_group_budget, + ) + + journal: list[str] = [] + prisma = _FakePrismaClient(journal, deployments=[_deployment()]) + cache = _FakeAuthCache(journal) + + with _proxy(prisma): + response = await delete_access_group_budget(access_group="prod-models", auth_cache=cache) + + assert response.budget_deleted is False + assert prisma.budget_table.deleted_ids == [] + _assert_evicted_after_write(journal, "prod-models", "access_group_budget.delete:prod-models") + + +@pytest.mark.asyncio +async def test_deleting_the_access_group_strips_deployments_before_dropping_the_budget(): + """Ordering is the point: stripping first means a failure leaves an unreachable budget row, + while the reverse leaves a live group whose enforcement silently vanished. The shared + LiteLLM_BudgetTable row survives here too.""" + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + delete_access_group, + ) + + journal: list[str] = [] + prisma = _FakePrismaClient(journal, deployments=[_deployment()]) + _seed_budget(prisma, "prod-models", spend=3.0, max_budget=100.0) + cache = _FakeAuthCache() + + with _proxy_with_stubbed_reload(prisma): + response = await delete_access_group( + access_group="prod-models", user_api_key_dict=_admin(), auth_cache=cache + ) + + assert response.models_updated == 1 + assert prisma.access_group_budget_table.rows == {} + assert prisma.budget_table.deleted_ids == [] + assert prisma.budget_table.rows["budget-seed"].max_budget == 100.0 + assert journal.index("model_table.update:deploy-1") < journal.index("access_group_budget.delete:prod-models") + + +@pytest.mark.asyncio +async def test_access_group_info_surfaces_the_budget_and_spend(): + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + get_access_group_info, + ) + + prisma = _FakePrismaClient([], deployments=[_deployment()]) + _seed_budget(prisma, "prod-models", spend=9.5, max_budget=100.0, soft_budget=50.0) + + with _proxy(prisma): + info = await get_access_group_info(access_group="prod-models", user_api_key_dict=_admin()) + + assert info.model_names == ["gpt-4o"] + assert info.spend == 9.5 + assert info.budget is not None + assert info.budget.max_budget == 100.0 + assert info.budget.soft_budget == 50.0 + + +@pytest.mark.asyncio +async def test_put_access_group_budget_evicts_both_auth_cache_keys(): + """Auth reads the per-group row and the registry of budgeted groups cache-first with no + freshness check, so a PUT that skips either eviction returns 200 and enforces nothing until + the TTL expires. Both keys, after the write.""" + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + set_access_group_budget, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import ( + AccessGroupBudgetRequest, + ) + + journal: list[str] = [] + prisma = _FakePrismaClient(journal, deployments=[_deployment()]) + cache = _FakeAuthCache(journal) + + with _proxy(prisma): + await set_access_group_budget( + access_group="prod-models", + data=AccessGroupBudgetRequest(max_budget=100.0), + user_api_key_dict=_admin(), + auth_cache=cache, + ) + + _assert_evicted_after_write(journal, "prod-models", "access_group_budget.upsert:prod-models") + + +@pytest.mark.asyncio +async def test_delete_access_group_budget_evicts_both_auth_cache_keys(): + """Clearing a budget has the same window as setting one: until both keys are dropped, auth + keeps enforcing the budget that is already gone.""" + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + delete_access_group_budget, + ) + + journal: list[str] = [] + prisma = _FakePrismaClient(journal, deployments=[_deployment()]) + _seed_budget(prisma, "prod-models", spend=12.0, max_budget=100.0) + cache = _FakeAuthCache(journal) + + with _proxy(prisma): + await delete_access_group_budget(access_group="prod-models", auth_cache=cache) + + _assert_evicted_after_write(journal, "prod-models", "access_group_budget.delete:prod-models") + + +@pytest.mark.asyncio +async def test_deleting_the_access_group_evicts_both_auth_cache_keys(): + """The group-delete cascade drops the budget row too, so it owes the same two evictions.""" + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + delete_access_group, + ) + + journal: list[str] = [] + prisma = _FakePrismaClient(journal, deployments=[_deployment()]) + _seed_budget(prisma, "prod-models", spend=3.0, max_budget=100.0) + cache = _FakeAuthCache(journal) + + with _proxy_with_stubbed_reload(prisma): + await delete_access_group(access_group="prod-models", user_api_key_dict=_admin(), auth_cache=cache) + + _assert_evicted_after_write(journal, "prod-models", "access_group_budget.delete:prod-models") + + +@pytest.mark.asyncio +async def test_deleting_an_access_group_that_never_had_a_budget_still_evicts(): + """The cascade's delete finds no row and reports nothing dropped, but the group can still be + sitting in the cached registry of budgeted groups, so both keys have to go regardless.""" + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + delete_access_group, + ) + + journal: list[str] = [] + prisma = _FakePrismaClient(journal, deployments=[_deployment()]) + cache = _FakeAuthCache(journal) + + with _proxy_with_stubbed_reload(prisma): + response = await delete_access_group( + access_group="prod-models", user_api_key_dict=_admin(), auth_cache=cache + ) + + assert response.models_updated == 1 + assert prisma.access_group_budget_table.rows == {} + _assert_evicted_after_write(journal, "prod-models", "access_group_budget.delete:prod-models") diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 38a346e7fb7..95067929ac1 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -2,6 +2,7 @@ import asyncio import threading from collections.abc import Mapping from datetime import datetime, timedelta, timezone +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -10,15 +11,16 @@ from fastapi import HTTPException import litellm from litellm.caching.dual_cache import DualCache from litellm.constants import STREAM_SSE_KEEPALIVE_PING_BYTES -from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( - AnthropicMessagesStreamingResponse, -) from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import ( AgenticAnthropicStreamingIterator, ) +from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( + AnthropicMessagesStreamingResponse, +) from litellm.proxy._types import ( LiteLLM_BudgetTable, LiteLLM_EndUserTable, + Litellm_EntityType, LiteLLM_OrganizationTable, LiteLLM_TagTable, LiteLLM_TeamMembership, @@ -27,9 +29,16 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.common_utils.reset_budget_job import _model_access_group_counter_key +from litellm.proxy.common_utils.user_api_key_cache import ( + UserApiKeyCache, + model_access_group_cache_key, + model_access_group_spend_counter_key, +) from litellm.proxy.spend_tracking.budget_reservation import ( TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS, _approximate_input_size, + _get_model_access_group_budget_counters, estimate_request_max_cost, get_budget_window_start, invalidate_budget_reservation_counters, @@ -39,6 +48,7 @@ from litellm.proxy.spend_tracking.budget_reservation import ( ) from litellm.proxy.utils import ProxyLogging from litellm.router import Router +from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget @pytest.fixture() @@ -2962,3 +2972,214 @@ async def test_small_prompt_is_tokenized_inline(spend_counter_state): assert reservation is not None assert threads == [threading.main_thread()] + + +class _ModelAccessGroupBudgetPrisma: + """Serves ``LiteLLM_ModelAccessGroupBudgetTable`` rows, recording what reached the database.""" + + def __init__(self, **max_budget_by_group) -> None: + self.rows = { + group: SimpleNamespace( + access_group_name=group, + spend=7.0, + litellm_budget_table=None if max_budget is None else SimpleNamespace(max_budget=max_budget), + ) + for group, max_budget in max_budget_by_group.items() + } + self.batches = [] + self.db = SimpleNamespace( + litellm_modelaccessgroupbudgettable=SimpleNamespace(find_many=self._find_many) + ) + + async def _find_many(self, **kwargs): + requested = list(kwargs["where"]["access_group_name"]["in"]) + self.batches.append(requested) + return [self.rows[group] for group in requested if group in self.rows] + + +async def _model_access_group_counters(matched, **max_budget_by_group): + return await _get_model_access_group_budget_counters( + valid_token=UserAPIKeyAuth(api_key="hashed", matched_model_access_groups=matched), + prisma_client=_ModelAccessGroupBudgetPrisma(**max_budget_by_group), + user_api_key_cache=UserApiKeyCache(), + ) + + +@pytest.mark.asyncio +async def test_model_access_group_with_a_budget_reserves_against_the_reset_jobs_counter_key(): + counters = await _model_access_group_counters(["premium"], premium=25.0) + + assert len(counters) == 1 + counter = counters[0] + assert counter.counter_key == _model_access_group_counter_key(SimpleNamespace(access_group_name="premium")) + assert counter.source_cache_key == model_access_group_cache_key("premium") + assert counter.max_budget == 25.0 + assert counter.fallback_spend == 7.0 + assert counter.entity_type == "Model access group" + assert counter.entity_id == "premium" + + +@pytest.mark.asyncio +async def test_model_access_group_without_a_budget_reserves_nothing(): + assert await _model_access_group_counters(["premium"], premium=None) == [] + + +@pytest.mark.asyncio +async def test_model_access_group_with_a_zero_budget_reserves_nothing(): + """Zero is how a budget is cleared, not a ceiling that blocks every request.""" + assert await _model_access_group_counters(["premium"], premium=0.0) == [] + + +@pytest.mark.asyncio +async def test_model_access_group_counters_come_from_the_auth_object(): + """Auth already resolved which granted groups serve the model; re-deriving it here would drift.""" + assert await _model_access_group_counters(None, premium=25.0) == [] + + +@pytest.mark.asyncio +async def test_repeated_model_access_group_reserves_once(): + counters = await _model_access_group_counters(["premium", "premium"], premium=25.0) + + assert [counter.entity_id for counter in counters] == ["premium"] + + +@pytest.mark.asyncio +async def test_model_access_group_counter_blocks_a_request_over_the_group_budget(spend_counter_state): + """End to end through the reservation path, which is what runs when reservations are enabled.""" + counter_cache, key_cache = spend_counter_state + prisma_client = _ModelAccessGroupBudgetPrisma(premium=1.0) + valid_token = UserAPIKeyAuth(api_key="hashed", token="tok", matched_model_access_groups=["premium"]) + + with patch( # test-quality-ok: reserve_budget_for_request takes no estimator, so pinning the estimate needs this attribute + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.5, + ): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=prisma_client, + user_api_key_cache=key_cache, + proxy_logging_obj=ProxyLogging(user_api_key_cache=key_cache), + ) + + assert exc_info.value.entity_id == "premium" + assert exc_info.value.entity_type == Litellm_EntityType.MODEL_ACCESS_GROUP.value + + +async def _cache_model_access_group_budget(key_cache, group, spend, max_budget=None): + await key_cache.async_set_cache( + key=model_access_group_cache_key(group), + value=ModelAccessGroupBudget(access_group_name=group, spend=spend, max_budget=max_budget), + model_type=ModelAccessGroupBudget, + ) + + +async def _reserve_for_model_access_groups(key_cache, groups, estimate): + """Reserve against the given groups, whose rows are already cached, so nothing hits the DB.""" + with patch( # test-quality-ok: reserve_budget_for_request takes no estimator, so pinning the estimate needs this attribute + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=estimate, + ): + return await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=UserAPIKeyAuth( + api_key="hashed", token="tok-mag-counter", matched_model_access_groups=list(groups) + ), + team_object=None, + user_object=None, + prisma_client=_ModelAccessGroupBudgetPrisma(), + user_api_key_cache=key_cache, + proxy_logging_obj=ProxyLogging(user_api_key_cache=key_cache), + ) + + +@pytest.mark.asyncio +async def test_model_access_group_counter_accumulates_across_calls_without_a_reservation(spend_counter_state): + """With reservations disabled nothing writes the counter up front, so the cost callback must. + + Otherwise the read-time budget check enforces against the DB row's spend, which the cache + holds for the full TTL, and a caller runs past the ceiling for that whole window. + """ + counter_cache, key_cache = spend_counter_state + await _cache_model_access_group_budget(key_cache, "premium", spend=1.0, max_budget=25.0) + + from litellm.proxy.proxy_server import increment_spend_counters + + counter_key = model_access_group_spend_counter_key("premium") + + await increment_spend_counters( + token=None, team_id=None, user_id=None, response_cost=0.25, model_access_groups=["premium"] + ) + assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(1.25) + + await increment_spend_counters( + token=None, team_id=None, user_id=None, response_cost=0.75, model_access_groups=["premium", "premium", ""] + ) + assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(2.0) + assert counter_cache.in_memory_cache.get_cache(key=model_access_group_spend_counter_key("")) is None + + +@pytest.mark.asyncio +async def test_reserved_model_access_group_is_not_charged_twice(spend_counter_state): + """The reservation already wrote this counter, so the post-call pass has to skip it.""" + counter_cache, key_cache = spend_counter_state + await _cache_model_access_group_budget(key_cache, "premium", spend=1.0, max_budget=25.0) + + reservation = await _reserve_for_model_access_groups(key_cache, ["premium"], estimate=0.6) + counter_key = model_access_group_spend_counter_key("premium") + assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(1.6) + + from litellm.proxy.proxy_server import increment_spend_counters + + await increment_spend_counters( + token=None, + team_id=None, + user_id=None, + response_cost=0.2, + budget_reservation=reservation, + model_access_groups=["premium"], + ) + + # 1.0 recorded + the reservation reconciled down to the 0.2 actually spent. A second + # increment would land at 1.4. + assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(1.2) + + +@pytest.mark.asyncio +async def test_unreserved_model_access_group_is_charged_alongside_a_reserved_one(spend_counter_state): + """A budgetless group reserves nothing, so only the post-call pass can charge it. + + Both groups authorized the request and both get debited, each exactly once, whether or not + the reservation path happened to hold a counter for them. + """ + counter_cache, key_cache = spend_counter_state + await _cache_model_access_group_budget(key_cache, "premium", spend=1.0, max_budget=25.0) + await _cache_model_access_group_budget(key_cache, "starter", spend=4.0) + + reservation = await _reserve_for_model_access_groups(key_cache, ["premium", "starter"], estimate=0.6) + assert [entry["entity_id"] for entry in reservation["entries"]] == ["premium"] + + from litellm.proxy.proxy_server import increment_spend_counters + + await increment_spend_counters( + token=None, + team_id=None, + user_id=None, + response_cost=0.2, + budget_reservation=reservation, + model_access_groups=["premium", "starter", "starter", "premium"], + ) + + assert counter_cache.in_memory_cache.get_cache( + key=model_access_group_spend_counter_key("premium") + ) == pytest.approx(1.2) + assert counter_cache.in_memory_cache.get_cache( + key=model_access_group_spend_counter_key("starter") + ) == pytest.approx(4.2) diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 649bcbb6511..ee0e2014951 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -33,6 +33,8 @@ from litellm.proxy.litellm_pre_call_utils import ( check_if_token_is_service_account, clean_headers, ) +from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs +from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY from litellm.litellm_core_utils.get_provider_specific_headers import ( ProviderSpecificHeaderUtils, ) @@ -7683,3 +7685,37 @@ async def test_add_litellm_data_to_request_keeps_litellm_metadata_on_litellm_met ) assert updated["litellm_metadata"]["trace_id"] == "abc" + + +def _stamp_model_access_groups(matched_model_access_groups, metadata_variable_name="metadata"): + user_api_key_dict = UserAPIKeyAuth(api_key="hashed-key") + user_api_key_dict.matched_model_access_groups = matched_model_access_groups + return LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata( + data={metadata_variable_name: {}}, + user_api_key_dict=user_api_key_dict, + _metadata_variable_name=metadata_variable_name, + )[metadata_variable_name] + + +def test_matched_model_access_groups_are_stamped_into_request_metadata(): + """The post-call spend writer reads the groups off request metadata, not off UserAPIKeyAuth.""" + stamped = _stamp_model_access_groups(["tier-a", "tier-b"]) + + assert stamped[MODEL_ACCESS_GROUP_METADATA_KEY] == ["tier-a", "tier-b"] + assert MODEL_ACCESS_GROUP_METADATA_KEY not in _stamp_model_access_groups(None) + + +def test_stamped_model_access_groups_survive_the_litellm_metadata_merge(): + """ + The key must keep its ``user_api_key`` prefix: when a request carries both metadata dicts, + get_litellm_metadata_from_kwargs returns litellm_metadata and copies a key over from metadata + only when that substring is in its name, so an unprefixed key is silently dropped. + """ + kwargs = { + "litellm_params": { + "metadata": _stamp_model_access_groups(["tier-a"]), + "litellm_metadata": {"trace_id": "abc"}, + } + } + + assert get_litellm_metadata_from_kwargs(kwargs)[MODEL_ACCESS_GROUP_METADATA_KEY] == ["tier-a"] diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index 6920cc0dae3..dcaad968663 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -1266,13 +1266,14 @@ class TestCreateSmtpConnection: patch("smtplib.SMTP_SSL") as mock_smtp_ssl, patch("smtplib.SMTP") as mock_smtp, ): - result = _create_smtp_connection(smtp_host="mail.example.com", smtp_port=465) + result = _create_smtp_connection(smtp_host="mail.example.com", smtp_port=465, timeout=30.0) mock_smtp.assert_not_called() assert result is mock_smtp_ssl.return_value _, kwargs = mock_smtp_ssl.call_args assert kwargs["host"] == "mail.example.com" assert kwargs["port"] == 465 + assert kwargs["timeout"] == 30.0 context = kwargs["context"] assert isinstance(context, ssl.SSLContext) assert context.verify_mode == ssl.CERT_REQUIRED @@ -1286,11 +1287,11 @@ class TestCreateSmtpConnection: patch("smtplib.SMTP_SSL") as mock_smtp_ssl, patch("smtplib.SMTP") as mock_smtp, ): - result = _create_smtp_connection(smtp_host="mail.example.com", smtp_port=587) + result = _create_smtp_connection(smtp_host="mail.example.com", smtp_port=587, timeout=30.0) mock_smtp_ssl.assert_not_called() assert result is mock_smtp.return_value - mock_smtp.assert_called_once_with(host="mail.example.com", port=587) + mock_smtp.assert_called_once_with(host="mail.example.com", port=587, timeout=30.0) class TestSendEmailStartTls: diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py b/tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py index 19abcb5d66d..fce51c9296c 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py @@ -14,6 +14,7 @@ from __future__ import annotations import asyncio import sys +import threading from dataclasses import dataclass, field from email.message import EmailMessage from pathlib import Path @@ -320,6 +321,7 @@ class _SentMessage: body: Optional[str] starttls_called: bool login_args: Optional[tuple] + thread_ident: int @dataclass @@ -328,6 +330,7 @@ class InMemorySMTP: sent: List[_SentMessage] = field(default_factory=list) raise_on_send: Optional[Exception] = None + connection_kwargs: List[Dict[str, Any]] = field(default_factory=list) def server_factory(self) -> Callable[..., Any]: outer = self @@ -370,10 +373,12 @@ class InMemorySMTP: body=body, starttls_called=self._starttls_called, login_args=self._login_args, + thread_ident=threading.get_ident(), ) ) def _factory(*args: Any, **kwargs: Any) -> _Conn: + outer.connection_kwargs.append(dict(kwargs)) return _Conn() return _factory diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_send_email.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_send_email.py index 739e942de52..0e8aba0a03b 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_send_email.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_send_email.py @@ -6,6 +6,7 @@ Symbols pinned here: from __future__ import annotations +import threading from typing import Any import pytest @@ -51,9 +52,7 @@ async def test_send_email_dispatches_via_smtp(in_memory_smtp: Any) -> None: @pytest.mark.asyncio -async def test_send_email_starttls_uses_ssl( - in_memory_smtp: Any, monkeypatch: pytest.MonkeyPatch -) -> None: +async def test_send_email_starttls_uses_ssl(in_memory_smtp: Any, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("SMTP_USE_SSL", "True") await send_email( receiver_email="to@invalid", @@ -82,9 +81,7 @@ async def test_send_email_error_missing_sender_email( ) -> None: monkeypatch.delenv("SMTP_SENDER_EMAIL", raising=False) with pytest.raises(ValueError, match="SMTP_SENDER_EMAIL"): - await send_email( - receiver_email="x@y", subject="s", html="

h

" - ) + await send_email(receiver_email="x@y", subject="s", html="

h

") @pytest.mark.asyncio @@ -105,6 +102,49 @@ async def test_send_email_error_missing_html() -> None: await send_email(receiver_email="x@y", subject="s", html=None) +@pytest.mark.asyncio +async def test_send_email_sets_connection_timeout(in_memory_smtp: Any) -> None: + await send_email( + receiver_email="to@invalid", + subject="Hi", + html="

x

", + ) + assert in_memory_smtp.connection_kwargs[0].get("timeout") == 30.0 + + +@pytest.mark.asyncio +async def test_send_email_timeout_env_override(in_memory_smtp: Any, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("SMTP_TIMEOUT", "5") + monkeypatch.setenv("SMTP_USE_SSL", "True") + await send_email( + receiver_email="to@invalid", + subject="Hi", + html="

x

", + ) + assert in_memory_smtp.connection_kwargs[0].get("timeout") == 5.0 + + +@pytest.mark.asyncio +async def test_send_email_malformed_timeout_is_swallowed(in_memory_smtp: Any, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("SMTP_TIMEOUT", "30s") + await send_email( + receiver_email="to@invalid", + subject="Hi", + html="

x

", + ) + assert in_memory_smtp.sent == [] + + +@pytest.mark.asyncio +async def test_send_email_runs_off_event_loop_thread(in_memory_smtp: Any) -> None: + await send_email( + receiver_email="to@invalid", + subject="Hi", + html="

x

", + ) + assert in_memory_smtp.sent[0].thread_ident != threading.get_ident() + + @pytest.mark.asyncio async def test_send_email_smtp_failure_is_swallowed( in_memory_smtp: Any, @@ -113,7 +153,5 @@ async def test_send_email_smtp_failure_is_swallowed( does not raise so a failing email never blocks the proxy. """ in_memory_smtp.raise_on_send = RuntimeError("smtp boom") - await send_email( - receiver_email="to@invalid", subject="Hi", html="

x

" - ) + await send_email(receiver_email="to@invalid", subject="Hi", html="

x

") assert in_memory_smtp.sent == [] diff --git a/tests/test_litellm/repositories/test_unit_of_work.py b/tests/test_litellm/repositories/test_unit_of_work.py index 1ebfd917e36..1a76b537e95 100644 --- a/tests/test_litellm/repositories/test_unit_of_work.py +++ b/tests/test_litellm/repositories/test_unit_of_work.py @@ -1,9 +1,11 @@ +from dataclasses import fields from datetime import datetime, timezone from typing import Any, Dict, List, Mapping, Tuple import pytest from litellm.repositories.unit_of_work import ( + LinkedSpendResetWrites, budget_cascade_unit_of_work, spend_reset_unit_of_work, ) @@ -32,6 +34,7 @@ class FakeBatch: self.litellm_teammembership = FakeBatchTable("litellm_teammembership", self.calls) self.litellm_organizationtable = FakeBatchTable("litellm_organizationtable", self.calls) self.litellm_tagtable = FakeBatchTable("litellm_tagtable", self.calls) + self.litellm_modelaccessgroupbudgettable = FakeBatchTable("litellm_modelaccessgroupbudgettable", self.calls) self.litellm_endusertable = FakeBatchTable("litellm_endusertable", self.calls) async def commit(self) -> None: @@ -90,6 +93,7 @@ async def test_budget_cascade_dependents_and_window_advance_share_one_batch(): uow.keys.queue_spend_zero(where=linked) uow.organizations.queue_spend_zero(where=linked) uow.tags.queue_spend_zero(where=linked) + uow.model_access_groups.queue_spend_zero(where=linked) uow.endusers.queue_spend_zero(where={"user_id": {"in": ["enduser-1"]}}) uow.budgets.queue_window_advance(budget_id="budget-1", budget_reset_at=reset_at) assert batch.commit_count == 0 @@ -100,6 +104,7 @@ async def test_budget_cascade_dependents_and_window_advance_share_one_batch(): ("litellm_verificationtoken.update_many", linked, {"spend": 0}), ("litellm_organizationtable.update_many", linked, {"spend": 0}), ("litellm_tagtable.update_many", linked, {"spend": 0}), + ("litellm_modelaccessgroupbudgettable.update_many", linked, {"spend": 0}), ("litellm_endusertable.update_many", {"user_id": {"in": ["enduser-1"]}}, {"spend": 0}), ("litellm_budgettable.update_many", {"budget_id": "budget-1"}, {"budget_reset_at": reset_at}), ] @@ -117,6 +122,39 @@ async def test_budget_window_advance_tolerates_a_tier_deleted_mid_chunk(): assert [call[0] for call in batch.calls] == ["litellm_budgettable.update_many"] +async def test_every_cascade_dependent_writes_to_its_own_table_on_the_one_batch(): + """Walks the dataclass instead of naming tables, so a dependent added to + BudgetCascadeUnitOfWork later cannot go uncovered. + + The named test above only proves the tables it lists, and an unbound + dependent surfaces as an AttributeError from whichever tests happen to + open a cascade. This pins the real contract: every field writes, each to a + distinct table, all on the same batch. + """ + reset_at = datetime(2026, 8, 3, 12, 0, tzinfo=timezone.utc) + batches: List[FakeBatch] = [] + + def _new_batch() -> FakeBatch: + # Fresh per call like db.batch_(), unlike the `lambda: batch` above: a + # second transaction would otherwise alias onto the first and hide. + batches.append(FakeBatch()) + return batches[-1] + + async with budget_cascade_unit_of_work(_new_batch) as uow: + writes = [getattr(uow, field.name) for field in fields(uow)] + for write in writes: + if isinstance(write, LinkedSpendResetWrites): + write.queue_spend_zero(where={"budget_id": "budget-1"}) + else: + write.queue_window_advance(budget_id="budget-1", budget_reset_at=reset_at) + + assert len(batches) == 1, "the cascade must open exactly one transaction" + batch = batches[0] + assert len(batch.calls) == len(writes), "a dependent bound to a batch of its own would not land here" + assert len({call[0] for call in batch.calls}) == len(writes), "two dependents share one table" + assert batch.commit_count == 1 + + async def test_budget_cascade_raising_inside_block_skips_commit(): """A failure part-way through must leave budget_reset_at where it was, so the tier is still due on the next tick.""" diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 6c5319f0361..c762183dec4 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -497,6 +497,91 @@ export interface paths { patch?: never; trace?: never; }; + "/access_group/{access_group}/budget": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Access Group Budget + * @description Get the shared budget of an access group, and the spend drawn against it. + * + * Example: + * ```bash + * curl -X GET 'http://localhost:4000/access_group/production-models/budget' \ + * -H 'Authorization: Bearer sk-1234' + * ``` + * + * Parameters: + * - access_group: str - The access group name (URL path parameter) + * + * Returns: + * - AccessGroupBudgetResponse; budget is null when the group has no budget set + * + * Raises: + * - HTTPException 404: If access group not found + */ + get: operations["get_access_group_budget_access_group__access_group__budget_get"]; + /** + * Set Access Group Budget + * @description Set or replace the shared budget of an access group. Idempotent. + * + * Every key that can reach a model in the group draws from this one budget. + * + * Example: + * ```bash + * curl -X PUT 'http://localhost:4000/access_group/production-models/budget' \ + * -H 'Authorization: Bearer sk-1234' \ + * -H 'Content-Type: application/json' \ + * -d '{ + * "max_budget": 100.0, + * "budget_duration": "30d" + * }' + * ``` + * + * Parameters: + * - access_group: str - The access group name (URL path parameter) + * - max_budget: Optional[float] - Requests fail once the group's shared spend exceeds this + * - soft_budget: Optional[float] - Fires an alert when reached; requests still succeed + * - budget_duration: Optional[str] - Frequency of resetting the group's spend (e.g. '30d') + * - budget_id: Optional[str] - Link an existing budget instead of creating one + * + * Returns: + * - AccessGroupBudgetResponse with the stored budget and current spend + * + * Raises: + * - HTTPException 400: If no budget field is given, or budget_duration cannot be parsed + * - HTTPException 404: If access group not found + */ + put: operations["set_access_group_budget_access_group__access_group__budget_put"]; + post?: never; + /** + * Delete Access Group Budget + * @description Clear the shared budget of an access group, leaving the group itself in place. + * + * Example: + * ```bash + * curl -X DELETE 'http://localhost:4000/access_group/production-models/budget' \ + * -H 'Authorization: Bearer sk-1234' + * ``` + * + * Parameters: + * - access_group: str - The access group name (URL path parameter) + * + * Returns: + * - DeleteAccessGroupBudgetResponse; budget_deleted is false when there was nothing to clear + * + * Raises: + * - HTTPException 404: If access group not found + */ + delete: operations["delete_access_group_budget_access_group__access_group__budget_delete"]; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/access_group/{access_group}/delete": { parameters: { query?: never; @@ -555,7 +640,7 @@ export interface paths { * - access_group: str - The access group name (URL path parameter) * * Returns: - * - AccessGroupInfo with the access group details + * - AccessGroupInfo with the access group details, its shared budget and its spend * * Raises: * - HTTPException 404: If access group not found @@ -22259,6 +22344,38 @@ export interface components { */ type: "restricted_sso_group"; }; + /** AccessGroupBudget */ + AccessGroupBudget: { + /** Budget Duration */ + budget_duration?: string | null; + /** Budget Id */ + budget_id: string; + /** Budget Reset At */ + budget_reset_at?: string | null; + /** Max Budget */ + max_budget?: number | null; + /** Soft Budget */ + soft_budget?: number | null; + }; + /** AccessGroupBudgetRequest */ + AccessGroupBudgetRequest: { + /** Budget Duration */ + budget_duration?: string | null; + /** Budget Id */ + budget_id?: string | null; + /** Max Budget */ + max_budget?: number | null; + /** Soft Budget */ + soft_budget?: number | null; + }; + /** AccessGroupBudgetResponse */ + AccessGroupBudgetResponse: { + /** Access Group */ + access_group: string; + budget?: components["schemas"]["AccessGroupBudget"] | null; + /** Spend */ + spend: number; + }; /** AccessGroupCreateRequest */ AccessGroupCreateRequest: { /** Access Agent Ids */ @@ -22280,10 +22397,13 @@ export interface components { AccessGroupInfo: { /** Access Group */ access_group: string; + budget?: components["schemas"]["AccessGroupBudget"] | null; /** Deployment Count */ deployment_count: number; /** Model Names */ model_names: string[]; + /** Spend */ + spend?: number | null; }; /** AccessGroupResponse */ AccessGroupResponse: { @@ -26132,6 +26252,15 @@ export interface components { [key: string]: unknown; }; }; + /** DeleteAccessGroupBudgetResponse */ + DeleteAccessGroupBudgetResponse: { + /** Access Group */ + access_group: string; + /** Budget Deleted */ + budget_deleted: boolean; + /** Message */ + message: string; + }; /** * DeleteCustomerRequest * @description Delete multiple Customers @@ -39152,6 +39281,103 @@ export interface operations { }; }; }; + get_access_group_budget_access_group__access_group__budget_get: { + parameters: { + query?: never; + header?: never; + path: { + access_group: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["AccessGroupBudgetResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + set_access_group_budget_access_group__access_group__budget_put: { + parameters: { + query?: never; + header?: never; + path: { + access_group: string; + }; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["AccessGroupBudgetRequest"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["AccessGroupBudgetResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + delete_access_group_budget_access_group__access_group__budget_delete: { + parameters: { + query?: never; + header?: never; + path: { + access_group: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["DeleteAccessGroupBudgetResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; delete_access_group_access_group__access_group__delete_delete: { parameters: { query?: never;