Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py: 57%
372 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1"""
2Allow proxy admin to manage model access groups
4Endpoints here:
5- POST /access_group/new - Create a new access group with multiple model names
6- GET /access_group/list - List every access group
7- GET /access_group/{access_group}/info - Read one access group, including its budget
8- PUT /access_group/{access_group}/update - Replace an access group's deployments
9- DELETE /access_group/{access_group}/delete - Delete an access group and its budget
10- GET /access_group/{access_group}/budget - Read an access group's shared budget and spend
11- PUT /access_group/{access_group}/budget - Set or replace an access group's shared budget
12- DELETE /access_group/{access_group}/budget - Clear an access group's shared budget
13"""
15import json
16from collections.abc import Mapping, Sequence
17from datetime import datetime
18from types import MappingProxyType
19from typing import TYPE_CHECKING, Annotated, Any, Final, Protocol
21from fastapi import APIRouter, Depends, HTTPException
22from typing_extensions import ReadOnly, TypedDict
24from litellm._logging import verbose_proxy_logger
25from litellm.proxy._types import UserAPIKeyAuth
26from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
27from litellm.proxy.common_utils.user_api_key_cache import (
28 UserApiKeyCache,
29 model_access_group_cache_key,
30 model_access_group_registry_cache_key,
31)
32from litellm.proxy.management_endpoints.common_utils import validate_budget_duration
34# Clear cache and reload models to pick up the access group changes
35from litellm.proxy.management_endpoints.model_management_endpoints import (
36 clear_cache,
37 live_model_ids_snapshot,
38 model_info_as_mapping,
39 reload_serving_verdict,
40)
41from litellm.proxy.management_helpers.utils import handle_budget_for_entity
42from litellm.proxy.utils import PrismaClient
43from litellm.repositories.model_repository import ModelRepository
44from litellm.repositories.table_repositories import ModelAccessGroupBudgetRepository
45from litellm.types.proxy.management_endpoints.model_management_endpoints import (
46 AccessGroupBudget,
47 AccessGroupBudgetRequest,
48 AccessGroupBudgetResponse,
49 AccessGroupInfo,
50 DeleteAccessGroupBudgetResponse,
51 DeleteModelGroupResponse,
52 ListAccessGroupsResponse,
53 NewModelGroupRequest,
54 NewModelGroupResponse,
55 UpdateModelGroupRequest,
56)
58if TYPE_CHECKING: 58 ↛ 59line 58 didn't jump to line 59 because the condition on line 58 was never true
59 from litellm import Router
61router: Final = APIRouter(tags=["model management"])
63_AUTH_DEPENDENCIES: Final = (Depends(user_api_key_auth),)
66class _ErrorDetail(TypedDict):
67 error: ReadOnly[str]
70class _ModelAccessGroupWhere(TypedDict):
71 access_group_name: ReadOnly[str]
74class _BudgetInclude(TypedDict):
75 litellm_budget_table: ReadOnly[bool]
78class _ModelAccessGroupBudgetCreate(TypedDict):
79 access_group_name: ReadOnly[str]
80 budget_id: ReadOnly[str | None]
81 created_by: ReadOnly[str]
82 updated_by: ReadOnly[str]
85class _ModelAccessGroupBudgetUpdate(TypedDict):
86 budget_id: ReadOnly[str | None]
87 updated_by: ReadOnly[str]
90class _ModelAccessGroupBudgetUpsert(TypedDict):
91 create: ReadOnly[_ModelAccessGroupBudgetCreate]
92 update: ReadOnly[_ModelAccessGroupBudgetUpdate]
95def _http_error(status_code: int, message: str) -> HTTPException:
96 detail: Final[_ErrorDetail] = {"error": message}
97 return HTTPException(status_code=status_code, detail=detail)
100class _DeploymentRow(Protocol):
101 @property
102 def model_id(self) -> str: ... 102 ↛ exitline 102 didn't return from function 'model_id' because
104 @property
105 def model_name(self) -> str: ... 105 ↛ exitline 105 didn't return from function 'model_name' because
107 @property
108 def model_info(self) -> object: ... 108 ↛ exitline 108 didn't return from function 'model_info' because
111class _ModelTableClient(Protocol):
112 async def find_many(self, *, where: Mapping[str, object] | None = None) -> Sequence[_DeploymentRow]: ... 112 ↛ exitline 112 didn't return from function 'find_many' because
114 async def find_unique(self, *, where: Mapping[str, object]) -> _DeploymentRow | None: ... 114 ↛ exitline 114 didn't return from function 'find_unique' because
116 async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> object: ... 116 ↛ exitline 116 didn't return from function 'update' because
119class _BudgetRow(Protocol):
120 @property
121 def budget_id(self) -> str: ... 121 ↛ exitline 121 didn't return from function 'budget_id' because
123 @property
124 def max_budget(self) -> float | None: ... 124 ↛ exitline 124 didn't return from function 'max_budget' because
126 @property
127 def soft_budget(self) -> float | None: ... 127 ↛ exitline 127 didn't return from function 'soft_budget' because
129 @property
130 def budget_duration(self) -> str | None: ... 130 ↛ exitline 130 didn't return from function 'budget_duration' because
132 @property
133 def budget_reset_at(self) -> datetime | None: ... 133 ↛ exitline 133 didn't return from function 'budget_reset_at' because
136class _ModelAccessGroupBudgetRow(Protocol):
137 @property
138 def access_group_name(self) -> str: ... 138 ↛ exitline 138 didn't return from function 'access_group_name' because
140 @property
141 def spend(self) -> float: ... 141 ↛ exitline 141 didn't return from function 'spend' because
143 @property
144 def budget_id(self) -> str | None: ... 144 ↛ exitline 144 didn't return from function 'budget_id' because
146 @property
147 def litellm_budget_table(self) -> _BudgetRow | None: ... 147 ↛ exitline 147 didn't return from function 'litellm_budget_table' because
150class _ModelAccessGroupBudgetTableClient(Protocol):
151 async def find_unique( 151 ↛ exitline 151 didn't return from function 'find_unique' because
152 self, *, where: Mapping[str, object], include: Mapping[str, object] | None = None
153 ) -> _ModelAccessGroupBudgetRow | None: ...
155 async def upsert( 155 ↛ exitline 155 didn't return from function 'upsert' because
156 self,
157 *,
158 where: Mapping[str, object],
159 data: Mapping[str, object],
160 include: Mapping[str, object] | None = None,
161 ) -> _ModelAccessGroupBudgetRow: ...
163 async def find_many( 163 ↛ exitline 163 didn't return from function 'find_many' because
164 self, *, include: Mapping[str, object] | None = None
165 ) -> Sequence[_ModelAccessGroupBudgetRow]: ...
167 async def delete(self, *, where: Mapping[str, object]) -> _ModelAccessGroupBudgetRow | None: ... 167 ↛ exitline 167 didn't return from function 'delete' because
170def _model_table(prisma_client: PrismaClient) -> _ModelTableClient:
171 return ModelRepository(prisma_client).table
174def _model_access_group_budget_table(prisma_client: PrismaClient) -> _ModelAccessGroupBudgetTableClient:
175 return ModelAccessGroupBudgetRepository(prisma_client).table
178def _prisma_client_or_500() -> PrismaClient:
179 from litellm.proxy.proxy_server import prisma_client
181 if prisma_client is None: 181 ↛ 182line 181 didn't jump to line 182 because the condition on line 181 was never true
182 raise _http_error(500, "Database not connected.")
183 return prisma_client
186def _auth_cache() -> UserApiKeyCache:
187 from litellm.proxy.proxy_server import user_api_key_cache
189 return user_api_key_cache
192async def _evict_model_access_group_cache_keys(access_group: str, auth_cache: UserApiKeyCache) -> None:
193 """
194 Every endpoint that writes an access group budget row must call this, or the budget stays
195 unenforced until the TTL expires: auth gates the feature on a cached registry of the groups
196 that have a budget row, read cache-first with no freshness check.
197 """
198 from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
199 evict_and_broadcast,
200 )
202 await evict_and_broadcast(
203 cache_keys=(model_access_group_cache_key(access_group), model_access_group_registry_cache_key()),
204 user_api_key_cache=auth_cache,
205 )
208async def _model_access_group_budget_row(
209 access_group: str, prisma_client: PrismaClient
210) -> _ModelAccessGroupBudgetRow | None:
211 where: Final[_ModelAccessGroupWhere] = {"access_group_name": access_group}
212 include: Final[_BudgetInclude] = {"litellm_budget_table": True}
213 return await _model_access_group_budget_table(prisma_client).find_unique(where=where, include=include)
216async def _model_access_group_budget_rows(
217 prisma_client: PrismaClient,
218) -> Mapping[str, _ModelAccessGroupBudgetRow]:
219 """Every group's budget row in one read, so listing groups does not fan out into one query
220 per group."""
221 include: Final[_BudgetInclude] = {"litellm_budget_table": True}
222 rows: Final = await _model_access_group_budget_table(prisma_client).find_many(include=include)
223 return MappingProxyType({row.access_group_name: row for row in rows})
226def _with_budget(info: AccessGroupInfo, row: _ModelAccessGroupBudgetRow | None) -> AccessGroupInfo:
227 """The group as listed, plus whatever budget hangs off it. A group with no row has spent
228 nothing, because clearing a budget drops the row that recorded the spend."""
229 return AccessGroupInfo(
230 access_group=info.access_group,
231 model_names=info.model_names,
232 deployment_count=info.deployment_count,
233 spend=row.spend if row is not None else 0.0,
234 budget=_budget_or_none(row),
235 )
238def _budget_or_none(row: _ModelAccessGroupBudgetRow | None) -> AccessGroupBudget | None:
239 budget: Final = row.litellm_budget_table if row is not None else None
240 if budget is None:
241 return None
242 return AccessGroupBudget(
243 budget_id=budget.budget_id,
244 max_budget=budget.max_budget,
245 soft_budget=budget.soft_budget,
246 budget_duration=budget.budget_duration,
247 budget_reset_at=budget.budget_reset_at,
248 )
251def _budget_response(access_group: str, row: _ModelAccessGroupBudgetRow | None) -> AccessGroupBudgetResponse:
252 return AccessGroupBudgetResponse(
253 access_group=access_group,
254 spend=row.spend if row is not None else 0.0,
255 budget=_budget_or_none(row),
256 )
259async def _delete_model_access_group_budget_row(
260 access_group: str, prisma_client: PrismaClient, auth_cache: UserApiKeyCache
261) -> bool:
262 """
263 Drop the group's budget row only, matching /tag/delete: the LiteLLM_BudgetTable row survives
264 because the link is ON DELETE SET NULL and a budget_id an admin passed in may be shared with
265 other entities.
267 Evicts unconditionally: a group with no row of its own can still be sitting in the cached
268 registry, so skipping the eviction when nothing was deleted would leave that stale.
269 """
270 where: Final[_ModelAccessGroupWhere] = {"access_group_name": access_group}
271 row: Final = await _model_access_group_budget_table(prisma_client).delete(where=where)
272 await _evict_model_access_group_cache_keys(access_group, auth_cache)
273 return row is not None
276async def _raise_404_if_model_access_group_missing(access_group: str, prisma_client: PrismaClient) -> None:
277 access_groups_map: Final = await get_all_access_groups_from_db(prisma_client=prisma_client)
278 if access_group not in access_groups_map: 278 ↛ exitline 278 didn't return from function '_raise_404_if_model_access_group_missing' because the condition on line 278 was always true
279 raise _http_error(404, f"Access group '{access_group}' not found")
282def validate_models_exist(model_names: Sequence[str], llm_router: "Router | None") -> tuple[bool, Sequence[str]]:
283 """
284 Validate that all requested model names exist in the router.
285 Checks only exact model name matches.
287 Returns:
288 (all_valid, missing_models)
289 """
290 if llm_router is None: 290 ↛ 291line 290 didn't jump to line 291 because the condition on line 290 was never true
291 return False, model_names
293 router_model_names: Final = frozenset(llm_router.get_model_names())
294 missing: Final = tuple(m for m in model_names if m not in router_model_names)
295 return (not missing, missing)
298async def _missing_models_after_read_through(
299 model_names: Sequence[str], llm_router: "Router | None"
300) -> tuple[str, ...]:
301 from litellm.proxy import proxy_server
302 from litellm.proxy.common_utils.registry_read_through import (
303 model_registry_read_through,
304 )
306 _, missing = validate_models_exist(model_names=model_names, llm_router=llm_router)
307 if not missing: 307 ↛ 308line 307 didn't jump to line 308 because the condition on line 307 was never true
308 return ()
309 for name in missing:
310 await model_registry_read_through.attempt(name)
311 _, still_missing = validate_models_exist(model_names=model_names, llm_router=proxy_server.llm_router)
312 return tuple(still_missing)
315def add_access_group_to_deployment(model_info: dict[str, Any], access_group: str) -> tuple[dict[str, Any], bool]:
316 """
317 Add an access group to a deployment's model_info.
319 Args:
320 model_info: The model_info dictionary from the deployment
321 access_group: The access group name to add
323 Returns:
324 Tuple[Dict[str, Any], bool]: (updated_model_info, was_modified)
325 """
326 access_groups: Final = model_info.get("access_groups", [])
328 # Check if access group already exists
329 if access_group in access_groups:
330 return model_info, False
332 # Add the access group
333 access_groups.append(access_group)
334 model_info["access_groups"] = access_groups
336 return model_info, True
339def _raise_http_if_reload_degraded_serving(
340 before: frozenset[str],
341 written_models: Sequence[tuple[str, object]],
342 access_group: str,
343 still_desired: frozenset[str] | None,
344 live_after: frozenset[str] | None,
345) -> None:
346 """Same verdict as the model-write endpoints, expressed through this file's
347 HTTPException error convention, with the metadata-only obligation: these writes
348 change group membership, not the models themselves, so a row that was already not
349 serving before the reload is never blamed here; only a model this reload stopped
350 serving is reported."""
351 missing, collateral = reload_serving_verdict(
352 before=before,
353 written_models=written_models,
354 written_must_serve=False,
355 still_desired=still_desired,
356 live_after=live_after,
357 )
358 gone: Final = tuple(dict.fromkeys((*missing, *collateral)))
359 if not gone:
360 return
361 raise HTTPException(
362 status_code=500,
363 detail={
364 "error": (
365 f"Access group '{access_group}' was saved to the database, but model id(s) {list(gone)} that "
366 "this pod was serving are no longer live after the reload it triggered. Other pods reload on "
367 "their own interval. Check server logs for 'Error upserting deployment' for the cause."
368 )
369 },
370 )
373async def _tag_deployment_with_access_group(
374 model_id: str,
375 model_info: object,
376 access_group: str,
377 prisma_client: PrismaClient,
378) -> tuple[str, Mapping[str, object]] | None:
379 """Write `access_group` into one deployment's model_info; returns the
380 (model_id, updated model_info) pair when a write happened, None when the
381 deployment already carried the group."""
382 updated_model_info, was_modified = add_access_group_to_deployment(
383 model_info=dict(_readable_model_info_or_raise(model_id=model_id, model_info=model_info)),
384 access_group=access_group,
385 )
386 if not was_modified:
387 return None
388 await _model_table(prisma_client).update(
389 where={"model_id": model_id},
390 data={"model_info": json.dumps(updated_model_info)},
391 )
392 verbose_proxy_logger.debug("Updated deployment %s with access group: %s", model_id, access_group)
393 return (model_id, updated_model_info)
396def _readable_model_info_or_raise(model_id: str, model_info: object) -> Mapping[str, object]:
397 """These helpers rewrite the model_info column wholesale, so a present-but-unreadable
398 value must refuse loudly rather than be silently replaced with a fresh object; an
399 absent value stays a legitimate empty start."""
400 parsed: Final = model_info_as_mapping(model_info)
401 if parsed is None and model_info is not None:
402 raise ValueError(f"model_info for deployment {model_id} is not a readable JSON object; refusing to rewrite it")
403 return parsed or {}
406async def _strip_access_group_from_deployment(
407 model_id: str,
408 model_info: object,
409 access_group: str,
410 prisma_client: PrismaClient,
411) -> tuple[str, Mapping[str, object]] | None:
412 """Remove `access_group` from one deployment's model_info; returns the
413 (model_id, updated model_info) pair when a write happened, None when the
414 deployment did not carry the group."""
415 updated_model_info, was_modified = remove_access_group_from_deployment(
416 model_info=dict(_readable_model_info_or_raise(model_id=model_id, model_info=model_info)),
417 access_group=access_group,
418 )
419 if not was_modified:
420 return None
421 await _model_table(prisma_client).update(
422 where={"model_id": model_id},
423 data={"model_info": json.dumps(updated_model_info)},
424 )
425 return (model_id, updated_model_info)
428async def update_deployments_with_access_group(
429 model_names: list[str],
430 access_group: str,
431 prisma_client: PrismaClient,
432) -> tuple[tuple[str, Mapping[str, object]], ...]:
433 """
434 Update all deployments for the given model names to include the access group.
436 Args:
437 model_names: List of model names whose deployments should be updated
438 access_group: The access group name to add
439 prisma_client: Database client
441 Returns:
442 The (model_id, updated model_info) pair of every deployment actually written,
443 so callers can verify each one survived the post-write reload
444 """
445 deployments: Final = await _model_table(prisma_client).find_many(where={"model_name": {"in": model_names}})
446 verbose_proxy_logger.debug("Found %s deployments for model_names: %s", len(deployments), model_names)
448 found_names: Final = {deployment.model_name for deployment in deployments}
449 for model_name in model_names:
450 if model_name not in found_names:
451 raise HTTPException(
452 status_code=400,
453 detail={
454 "error": f"Can't find model '{model_name}' in Database. Access group management is only supported for database models."
455 },
456 )
458 tagged: Final = [
459 await _tag_deployment_with_access_group(
460 model_id=deployment.model_id,
461 model_info=deployment.model_info,
462 access_group=access_group,
463 prisma_client=prisma_client,
464 )
465 for deployment in deployments
466 ]
467 return tuple(pair for pair in tagged if pair is not None)
470async def update_specific_deployments_with_access_group(
471 model_ids: list[str],
472 access_group: str,
473 prisma_client: PrismaClient,
474) -> tuple[tuple[str, Mapping[str, object]], ...]:
475 """
476 Update specific deployments (by model_id) to include the access group.
478 Unlike update_deployments_with_access_group which tags ALL deployments sharing
479 a model_name, this function only tags the specific deployments identified by
480 their unique model_id. Returns the (model_id, updated model_info) pair of every
481 deployment actually written.
482 """
483 verbose_proxy_logger.debug("Updating specific deployment model_ids: %s", model_ids)
484 tagged: Final = [
485 await _tag_deployment_with_access_group(
486 model_id=model_id,
487 model_info=(await _find_deployment_or_400(model_id=model_id, prisma_client=prisma_client)),
488 access_group=access_group,
489 prisma_client=prisma_client,
490 )
491 for model_id in model_ids
492 ]
493 return tuple(pair for pair in tagged if pair is not None)
496async def _find_deployment_or_400(model_id: str, prisma_client: PrismaClient) -> object:
497 deployment: Final = await _model_table(prisma_client).find_unique(where={"model_id": model_id})
498 if deployment is None: 498 ↛ 503line 498 didn't jump to line 503 because the condition on line 498 was always true
499 raise HTTPException(
500 status_code=400,
501 detail={"error": f"Deployment with model_id '{model_id}' not found in Database."},
502 )
503 return deployment.model_info
506def remove_access_group_from_deployment(model_info: dict[str, Any], access_group: str) -> tuple[dict[str, Any], bool]:
507 """
508 Remove an access group from a deployment's model_info.
510 Args:
511 model_info: The model_info dictionary from the deployment
512 access_group: The access group name to remove
514 Returns:
515 Tuple[Dict[str, Any], bool]: (updated_model_info, was_modified)
516 """
517 access_groups: Final = model_info.get("access_groups", [])
519 # Check if access group exists
520 if access_group not in access_groups:
521 return model_info, False
523 # Remove the access group
524 access_groups.remove(access_group)
525 model_info["access_groups"] = access_groups
527 return model_info, True
530async def get_all_access_groups_from_db(
531 prisma_client: PrismaClient,
532) -> dict[str, AccessGroupInfo]:
533 """
534 Get all access groups from the database.
536 Returns:
537 Dict[str, AccessGroupInfo]: Dictionary mapping access_group name to info
538 """
539 # Get all deployments
540 deployments: Final = await ModelRepository(prisma_client).table.find_many()
542 # Build access group map
543 model_names_by_group: Final[dict[str, list[str]]] = {}
545 for deployment in deployments:
546 model_info = deployment.model_info or {}
547 access_groups = model_info.get( # pyright: ignore[reportAttributeAccessIssue] # Json reads back as a dict
548 "access_groups", []
549 )
550 model_name = deployment.model_name
552 for access_group in access_groups: 552 ↛ 553line 552 didn't jump to line 553 because the loop on line 552 never started
553 if access_group not in model_names_by_group:
554 model_names_by_group[access_group] = []
556 model_names_by_group[access_group].append(model_name)
558 # Convert to AccessGroupInfo objects
559 return {
560 access_group: AccessGroupInfo(
561 access_group=access_group,
562 model_names=sorted(frozenset(model_names)),
563 deployment_count=len(model_names),
564 )
565 for access_group, model_names in model_names_by_group.items()
566 }
569@router.post(
570 "/access_group/new",
571 dependencies=_AUTH_DEPENDENCIES,
572 response_model=NewModelGroupResponse,
573)
574async def create_model_group(
575 data: NewModelGroupRequest,
576 user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
577):
578 """
579 Create a new access group containing multiple model names.
581 An access group is a named collection of model groups that can be referenced
582 by teams/keys for simplified access control.
584 Example:
585 ```bash
586 curl -X POST 'http://localhost:4000/access_group/new' \\
587 -H 'Authorization: Bearer sk-1234' \\
588 -H 'Content-Type: application/json' \\
589 -d '{
590 "access_group": "production-models",
591 "model_names": ["gpt-4", "claude-3-opus", "gemini-pro"]
592 }'
593 ```
595 Parameters:
596 - access_group: str - The access group name (e.g., "production-models")
597 - model_names: List[str] - List of existing model groups to include
599 Returns:
600 - NewModelGroupResponse with the created access group details
602 Raises:
603 - HTTPException 400: If any model names don't exist
604 - HTTPException 500: If database operations fail
605 """
606 from litellm.proxy.proxy_server import (
607 llm_router,
608 prisma_client,
609 )
611 verbose_proxy_logger.debug("Creating access group: %s with models: %s", data.access_group, data.model_names)
613 # Validation: Check if access_group is provided
614 if not data.access_group or not data.access_group.strip():
615 raise HTTPException(
616 status_code=400,
617 detail={"error": "access_group is required and cannot be empty"},
618 )
620 # Validation: Check that at least one of model_names or model_ids is provided
621 has_model_names: Final = data.model_names and len(data.model_names) > 0
622 has_model_ids: Final = data.model_ids and len(data.model_ids) > 0
624 if not has_model_names and not has_model_ids:
625 raise HTTPException(
626 status_code=400,
627 detail={"error": "Either model_names or model_ids must be provided and non-empty"},
628 )
630 # If model_ids is provided, use it (more precise targeting)
631 use_model_ids: Final = has_model_ids
633 # Validate model_names exist in router (only if using model_names path)
634 if not use_model_ids and has_model_names:
635 assert data.model_names is not None
636 missing_models: Final = await _missing_models_after_read_through(
637 model_names=data.model_names,
638 llm_router=llm_router,
639 )
641 if missing_models: 641 ↛ 648line 641 didn't jump to line 648 because the condition on line 641 was always true
642 raise HTTPException(
643 status_code=400,
644 detail={"error": f"Model(s) not found: {', '.join(missing_models)}"},
645 )
647 # Check if database is connected
648 if prisma_client is None: 648 ↛ 649line 648 didn't jump to line 649 because the condition on line 648 was never true
649 raise HTTPException(
650 status_code=500,
651 detail={"error": "Database not connected. Cannot create access group."},
652 )
654 try:
655 # Check if access group already exists
656 existing_access_groups: Final = await get_all_access_groups_from_db(prisma_client=prisma_client)
658 if data.access_group in existing_access_groups: 658 ↛ 659line 658 didn't jump to line 659 because the condition on line 658 was never true
659 raise HTTPException(
660 status_code=409,
661 detail={
662 "error": f"Access group '{data.access_group}' already exists. Use PUT /access_group/{data.access_group}/update to modify it."
663 },
664 )
666 # Update deployments using the appropriate method
667 if use_model_ids: 667 ↛ 675line 667 didn't jump to line 675 because the condition on line 667 was always true
668 assert data.model_ids is not None
669 updated_pairs = await update_specific_deployments_with_access_group(
670 model_ids=data.model_ids,
671 access_group=data.access_group,
672 prisma_client=prisma_client,
673 )
674 else:
675 assert data.model_names is not None
676 updated_pairs = await update_deployments_with_access_group(
677 model_names=data.model_names,
678 access_group=data.access_group,
679 prisma_client=prisma_client,
680 )
681 models_updated: Final = len(updated_pairs)
683 live_before_reload: Final = live_model_ids_snapshot()
685 reload_outcome: Final = await clear_cache()
686 _raise_http_if_reload_degraded_serving(
687 before=live_before_reload,
688 written_models=updated_pairs,
689 access_group=data.access_group,
690 still_desired=reload_outcome.still_desired,
691 live_after=reload_outcome.live_after,
692 )
694 verbose_proxy_logger.info(
695 "Successfully created access group '%s' with %s models updated", data.access_group, models_updated
696 )
698 return NewModelGroupResponse(
699 access_group=data.access_group,
700 model_names=data.model_names,
701 model_ids=data.model_ids,
702 models_updated=models_updated,
703 )
705 except HTTPException:
706 raise
707 except Exception as e:
708 verbose_proxy_logger.exception("Error creating access group '%s': %s", data.access_group, e)
709 raise HTTPException(
710 status_code=500,
711 detail={"error": f"Failed to create access group: {e}"},
712 )
715@router.get(
716 "/access_group/list",
717 dependencies=_AUTH_DEPENDENCIES,
718 response_model=ListAccessGroupsResponse,
719)
720async def list_access_groups(
721 user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
722):
723 """
724 List all access groups.
726 Returns a list of all access groups with their model names, deployment counts, shared budget
727 and the spend drawn against it.
729 Example:
730 ```bash
731 curl -X GET 'http://localhost:4000/access_group/list' \\
732 -H 'Authorization: Bearer sk-1234'
733 ```
735 Returns:
736 - ListAccessGroupsResponse with all access groups
737 """
738 from litellm.proxy.proxy_server import prisma_client
740 if prisma_client is None: 740 ↛ 741line 740 didn't jump to line 741 because the condition on line 740 was never true
741 raise HTTPException(
742 status_code=500,
743 detail={"error": "Database not connected."},
744 )
746 try:
747 access_groups_map: Final = await get_all_access_groups_from_db(prisma_client=prisma_client)
748 budget_rows: Final = await _model_access_group_budget_rows(prisma_client)
750 access_groups_list: Final = sorted(
751 (_with_budget(info, budget_rows.get(info.access_group)) for info in access_groups_map.values()),
752 key=lambda group: group.access_group,
753 )
755 return ListAccessGroupsResponse(access_groups=access_groups_list)
757 except Exception as e:
758 verbose_proxy_logger.exception("Error listing access groups: %s", e)
759 raise HTTPException(
760 status_code=500,
761 detail={"error": f"Failed to list access groups: {e}"},
762 )
765@router.get(
766 "/access_group/{access_group}/info",
767 dependencies=_AUTH_DEPENDENCIES,
768 response_model=AccessGroupInfo,
769)
770async def get_access_group_info(
771 access_group: str,
772 user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
773):
774 """
775 Get information about a specific access group.
777 Example:
778 ```bash
779 curl -X GET 'http://localhost:4000/access_group/production-models/info' \\
780 -H 'Authorization: Bearer sk-1234'
781 ```
783 Parameters:
784 - access_group: str - The access group name (URL path parameter)
786 Returns:
787 - AccessGroupInfo with the access group details, its shared budget and its spend
789 Raises:
790 - HTTPException 404: If access group not found
791 """
792 from litellm.proxy.proxy_server import prisma_client
794 if prisma_client is None: 794 ↛ 795line 794 didn't jump to line 795 because the condition on line 794 was never true
795 raise HTTPException(
796 status_code=500,
797 detail={"error": "Database not connected."},
798 )
800 try:
801 access_groups_map: Final = await get_all_access_groups_from_db(prisma_client=prisma_client)
803 if access_group not in access_groups_map: 803 ↛ 809line 803 didn't jump to line 809 because the condition on line 803 was always true
804 raise HTTPException(
805 status_code=404,
806 detail={"error": f"Access group '{access_group}' not found"},
807 )
809 return _with_budget(
810 access_groups_map[access_group],
811 await _model_access_group_budget_row(access_group, prisma_client),
812 )
814 except HTTPException:
815 raise
816 except Exception as e:
817 verbose_proxy_logger.exception("Error getting access group info for '%s': %s", access_group, e)
818 raise HTTPException(
819 status_code=500,
820 detail={"error": f"Failed to get access group info: {e}"},
821 )
824@router.put(
825 "/access_group/{access_group}/update",
826 dependencies=_AUTH_DEPENDENCIES,
827 response_model=NewModelGroupResponse,
828)
829async def update_access_group(
830 access_group: str,
831 data: UpdateModelGroupRequest,
832 user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
833):
834 """
835 Update an access group's model names.
837 This will:
838 1. Remove the access group from all current deployments
839 2. Add the access group to all deployments for the new model_names list
841 Example:
842 ```bash
843 curl -X PUT 'http://localhost:4000/access_group/production-models/update' \\
844 -H 'Authorization: Bearer sk-1234' \\
845 -H 'Content-Type: application/json' \\
846 -d '{
847 "model_names": ["gpt-4", "claude-3-sonnet"]
848 }'
849 ```
851 Parameters:
852 - access_group: str - The access group name (URL path parameter)
853 - model_names: List[str] - New list of model groups to include
855 Returns:
856 - NewModelGroupResponse with the updated access group details
858 Raises:
859 - HTTPException 400: If any model names don't exist
860 - HTTPException 404: If access group not found
861 """
862 from litellm.proxy.proxy_server import llm_router, prisma_client
864 if prisma_client is None: 864 ↛ 865line 864 didn't jump to line 865 because the condition on line 864 was never true
865 raise HTTPException(
866 status_code=500,
867 detail={"error": "Database not connected."},
868 )
870 verbose_proxy_logger.debug("Updating access group: %s with models: %s", access_group, data.model_names)
872 # Validation: Check that at least one of model_names or model_ids is provided
873 has_model_names: Final = data.model_names and len(data.model_names) > 0
874 has_model_ids: Final = data.model_ids and len(data.model_ids) > 0
876 if not has_model_names and not has_model_ids:
877 raise HTTPException(
878 status_code=400,
879 detail={"error": "Either model_names or model_ids must be provided and non-empty"},
880 )
882 use_model_ids: Final = has_model_ids
884 # Validation: Check if access group exists
885 try:
886 access_groups_map: Final = await get_all_access_groups_from_db(prisma_client=prisma_client)
887 if access_group not in access_groups_map: 887 ↛ 901line 887 didn't jump to line 901 because the condition on line 887 was always true
888 raise HTTPException(
889 status_code=404,
890 detail={"error": f"Access group '{access_group}' not found"},
891 )
892 except HTTPException:
893 raise
894 except Exception as e:
895 raise HTTPException(
896 status_code=500,
897 detail={"error": f"Failed to check access group existence: {e}"},
898 )
900 # Validation: Check if all new models exist (only if using model_names path)
901 if not use_model_ids and has_model_names:
902 assert data.model_names is not None
903 missing_models: Final = await _missing_models_after_read_through(
904 model_names=data.model_names,
905 llm_router=llm_router,
906 )
908 if missing_models:
909 raise HTTPException(
910 status_code=400,
911 detail={"error": f"Model(s) not found: {', '.join(missing_models)}"},
912 )
914 try:
915 # Step 1: Remove access group from ALL DB deployments (skip config models)
916 all_deployments: Final = await _model_table(prisma_client).find_many()
918 stripped: Final = [
919 await _strip_access_group_from_deployment(
920 model_id=deployment.model_id,
921 model_info=deployment.model_info,
922 access_group=access_group,
923 prisma_client=prisma_client,
924 )
925 for deployment in all_deployments
926 ]
927 stripped_pairs: Final = tuple(pair for pair in stripped if pair is not None)
929 # Step 2: Add access group using the appropriate method
930 if use_model_ids:
931 assert data.model_ids is not None
932 updated_pairs = await update_specific_deployments_with_access_group(
933 model_ids=data.model_ids,
934 access_group=access_group,
935 prisma_client=prisma_client,
936 )
937 else:
938 assert data.model_names is not None
939 updated_pairs = await update_deployments_with_access_group(
940 model_names=data.model_names,
941 access_group=access_group,
942 prisma_client=prisma_client,
943 )
944 models_updated: Final = len(updated_pairs)
946 # Clear cache and reload models to pick up the access group changes
947 live_before_reload: Final = live_model_ids_snapshot()
948 reload_outcome: Final = await clear_cache()
949 _raise_http_if_reload_degraded_serving(
950 before=live_before_reload,
951 written_models=list({**dict(stripped_pairs), **dict(updated_pairs)}.items()),
952 access_group=access_group,
953 still_desired=reload_outcome.still_desired,
954 live_after=reload_outcome.live_after,
955 )
957 verbose_proxy_logger.info(
958 "Successfully updated access group '%s' with %s models updated", access_group, models_updated
959 )
961 return NewModelGroupResponse(
962 access_group=access_group,
963 model_names=data.model_names,
964 model_ids=data.model_ids,
965 models_updated=models_updated,
966 )
968 except HTTPException:
969 raise
970 except Exception as e:
971 verbose_proxy_logger.exception("Error updating access group '%s': %s", access_group, e)
972 raise HTTPException(
973 status_code=500,
974 detail={"error": f"Failed to update access group: {e}"},
975 )
978@router.delete(
979 "/access_group/{access_group}/delete",
980 dependencies=_AUTH_DEPENDENCIES,
981 response_model=DeleteModelGroupResponse,
982)
983async def delete_access_group(
984 access_group: str,
985 user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
986 auth_cache: Annotated[UserApiKeyCache, Depends(_auth_cache)],
987):
988 """
989 Delete an access group.
991 Removes the access group from all deployments that have it.
993 Example:
994 ```bash
995 curl -X DELETE 'http://localhost:4000/access_group/production-models/delete' \\
996 -H 'Authorization: Bearer sk-1234'
997 ```
999 Parameters:
1000 - access_group: str - The access group name (URL path parameter)
1002 Returns:
1003 - DeleteModelGroupResponse with deletion details
1005 Raises:
1006 - HTTPException 404: If access group not found
1007 """
1008 from litellm.proxy.proxy_server import prisma_client
1010 if prisma_client is None: 1010 ↛ 1011line 1010 didn't jump to line 1011 because the condition on line 1010 was never true
1011 raise HTTPException(
1012 status_code=500,
1013 detail={"error": "Database not connected."},
1014 )
1016 verbose_proxy_logger.debug("Deleting access group: %s", access_group)
1018 # Validation: Check if access group exists
1019 try:
1020 access_groups_map: Final = await get_all_access_groups_from_db(prisma_client=prisma_client)
1021 if access_group not in access_groups_map: 1021 ↛ 1034line 1021 didn't jump to line 1034 because the condition on line 1021 was always true
1022 raise HTTPException(
1023 status_code=404,
1024 detail={"error": f"Access group '{access_group}' not found"},
1025 )
1026 except HTTPException:
1027 raise
1028 except Exception as e:
1029 raise HTTPException(
1030 status_code=500,
1031 detail={"error": f"Failed to check access group existence: {e}"},
1032 )
1034 try:
1035 # Remove access group from all DB deployments (skip config models)
1036 all_deployments: Final = await _model_table(prisma_client).find_many()
1038 removed: Final = [
1039 await _strip_access_group_from_deployment(
1040 model_id=deployment.model_id,
1041 model_info=deployment.model_info,
1042 access_group=access_group,
1043 prisma_client=prisma_client,
1044 )
1045 for deployment in all_deployments
1046 ]
1047 removed_pairs: Final = tuple(pair for pair in removed if pair is not None)
1048 models_updated: Final = len(removed_pairs)
1050 # Budget last, deliberately: failing here strands a budget row for a group already on no
1051 # deployment (clutter), where the reverse order can leave a live group enforcing nothing.
1052 # The LiteLLM_BudgetTable row it linked is left alone, as /tag/delete leaves a tag's.
1053 await _delete_model_access_group_budget_row(
1054 access_group=access_group, prisma_client=prisma_client, auth_cache=auth_cache
1055 )
1057 # Clear cache and reload models to pick up the access group changes
1058 live_before_reload: Final = live_model_ids_snapshot()
1059 reload_outcome: Final = await clear_cache()
1060 _raise_http_if_reload_degraded_serving(
1061 before=live_before_reload,
1062 written_models=removed_pairs,
1063 access_group=access_group,
1064 still_desired=reload_outcome.still_desired,
1065 live_after=reload_outcome.live_after,
1066 )
1068 verbose_proxy_logger.info(
1069 "Successfully deleted access group '%s' from %s deployments", access_group, models_updated
1070 )
1072 return DeleteModelGroupResponse(
1073 access_group=access_group,
1074 models_updated=models_updated,
1075 message=f"Access group '{access_group}' deleted successfully",
1076 )
1078 except HTTPException:
1079 raise
1080 except Exception as e:
1081 verbose_proxy_logger.exception("Error deleting access group '%s': %s", access_group, e)
1082 raise HTTPException(
1083 status_code=500,
1084 detail={"error": f"Failed to delete access group: {e}"},
1085 )
1088@router.get(
1089 "/access_group/{access_group}/budget",
1090 dependencies=_AUTH_DEPENDENCIES,
1091 response_model=AccessGroupBudgetResponse,
1092)
1093async def get_access_group_budget(
1094 access_group: str,
1095) -> AccessGroupBudgetResponse:
1096 """
1097 Get the shared budget of an access group, and the spend drawn against it.
1099 Example:
1100 ```bash
1101 curl -X GET 'http://localhost:4000/access_group/production-models/budget' \\
1102 -H 'Authorization: Bearer sk-1234'
1103 ```
1105 Parameters:
1106 - access_group: str - The access group name (URL path parameter)
1108 Returns:
1109 - AccessGroupBudgetResponse; budget is null when the group has no budget set
1111 Raises:
1112 - HTTPException 404: If access group not found
1113 """
1114 prisma_client: Final = _prisma_client_or_500()
1115 await _raise_404_if_model_access_group_missing(access_group=access_group, prisma_client=prisma_client)
1117 return _budget_response(
1118 access_group=access_group,
1119 row=await _model_access_group_budget_row(access_group, prisma_client),
1120 )
1123@router.put(
1124 "/access_group/{access_group}/budget",
1125 dependencies=_AUTH_DEPENDENCIES,
1126 response_model=AccessGroupBudgetResponse,
1127)
1128async def set_access_group_budget(
1129 access_group: str,
1130 data: AccessGroupBudgetRequest,
1131 user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
1132 auth_cache: Annotated[UserApiKeyCache, Depends(_auth_cache)],
1133) -> AccessGroupBudgetResponse:
1134 """
1135 Set or replace the shared budget of an access group. Idempotent.
1137 Every key that can reach a model in the group draws from this one budget.
1139 Example:
1140 ```bash
1141 curl -X PUT 'http://localhost:4000/access_group/production-models/budget' \\
1142 -H 'Authorization: Bearer sk-1234' \\
1143 -H 'Content-Type: application/json' \\
1144 -d '{
1145 "max_budget": 100.0,
1146 "budget_duration": "30d"
1147 }'
1148 ```
1150 Parameters:
1151 - access_group: str - The access group name (URL path parameter)
1152 - max_budget: Optional[float] - Requests fail once the group's shared spend exceeds this
1153 - soft_budget: Optional[float] - Fires an alert when reached; requests still succeed
1154 - budget_duration: Optional[str] - Frequency of resetting the group's spend (e.g. '30d')
1155 - budget_id: Optional[str] - Link an existing budget instead of creating one
1157 Returns:
1158 - AccessGroupBudgetResponse with the stored budget and current spend
1160 Raises:
1161 - HTTPException 400: If no budget field is given, or budget_duration cannot be parsed
1162 - HTTPException 404: If access group not found
1163 """
1164 from litellm.proxy.proxy_server import litellm_proxy_admin_name
1166 prisma_client: Final = _prisma_client_or_500()
1167 if not data.model_dump(exclude_none=True):
1168 raise _http_error(400, "One of max_budget, soft_budget, budget_duration or budget_id is required")
1169 validate_budget_duration(data.budget_duration)
1170 await _raise_404_if_model_access_group_missing(access_group=access_group, prisma_client=prisma_client)
1172 existing_row: Final = await _model_access_group_budget_row(access_group, prisma_client)
1173 budget_id: Final = await handle_budget_for_entity(
1174 data=data,
1175 existing_budget_id=existing_row.budget_id if existing_row is not None else None,
1176 user_api_key_dict=user_api_key_dict,
1177 prisma_client=prisma_client,
1178 litellm_proxy_admin_name=litellm_proxy_admin_name,
1179 )
1180 actor: Final = user_api_key_dict.user_id or litellm_proxy_admin_name
1181 upsert_data: Final[_ModelAccessGroupBudgetUpsert] = {
1182 "create": {
1183 "access_group_name": access_group,
1184 "budget_id": budget_id,
1185 "created_by": actor,
1186 "updated_by": actor,
1187 },
1188 "update": {"budget_id": budget_id, "updated_by": actor},
1189 }
1190 where: Final[_ModelAccessGroupWhere] = {"access_group_name": access_group}
1191 include: Final[_BudgetInclude] = {"litellm_budget_table": True}
1192 row: Final = await _model_access_group_budget_table(prisma_client).upsert(
1193 where=where, data=upsert_data, include=include
1194 )
1195 await _evict_model_access_group_cache_keys(access_group, auth_cache)
1197 verbose_proxy_logger.info("Set budget %s on access group '%s'", budget_id, access_group)
1198 return _budget_response(access_group=access_group, row=row)
1201@router.delete(
1202 "/access_group/{access_group}/budget",
1203 dependencies=_AUTH_DEPENDENCIES,
1204 response_model=DeleteAccessGroupBudgetResponse,
1205)
1206async def delete_access_group_budget(
1207 access_group: str,
1208 auth_cache: Annotated[UserApiKeyCache, Depends(_auth_cache)],
1209) -> DeleteAccessGroupBudgetResponse:
1210 """
1211 Clear the shared budget of an access group, leaving the group itself in place.
1213 Example:
1214 ```bash
1215 curl -X DELETE 'http://localhost:4000/access_group/production-models/budget' \\
1216 -H 'Authorization: Bearer sk-1234'
1217 ```
1219 Parameters:
1220 - access_group: str - The access group name (URL path parameter)
1222 Returns:
1223 - DeleteAccessGroupBudgetResponse; budget_deleted is false when there was nothing to clear
1225 Raises:
1226 - HTTPException 404: If access group not found
1227 """
1228 prisma_client: Final = _prisma_client_or_500()
1229 await _raise_404_if_model_access_group_missing(access_group=access_group, prisma_client=prisma_client)
1231 budget_deleted: Final = await _delete_model_access_group_budget_row(
1232 access_group=access_group,
1233 prisma_client=prisma_client,
1234 auth_cache=auth_cache,
1235 )
1236 return DeleteAccessGroupBudgetResponse(
1237 access_group=access_group,
1238 budget_deleted=budget_deleted,
1239 message=(
1240 f"Budget for access group '{access_group}' deleted successfully"
1241 if budget_deleted
1242 else f"Access group '{access_group}' has no budget to delete"
1243 ),
1244 )