Litellm budgets

backend/.env.example

@@ -30,5 +30,5 @@ SUPABASE_SERVICE_ROLE_KEY=your_supabase_service_role_key
 # LiteLLM Proxy Configuration
 LITELLM_URL=http://litellm:4000
 LITELLM_MASTER_KEY=sk-litellm-change-me
-DEFAULT_USER_BUDGET=10.00
+DEFAULT_USER_BUDGET=2.00
 DEFAULT_BUDGET_DURATION=30d

backend/app/controllers/litellm_controller.py

@@ -0,0 +1,38 @@
+import httpx
+from fastapi import HTTPException
+
+from app.services.litellm_service import list_managed_users, update_all_user_budgets
+
+
+async def list_managed_users_controller() -> dict:
+    """
+    List the LiteLLM users a bulk budget update would apply to.
+    """
+
+    try:
+        users = await list_managed_users()
+
+    except httpx.HTTPError as error:
+        raise HTTPException(
+            status_code=502,
+            detail=f"Failed to fetch users from LiteLLM: {error}",
+        )
+
+    return {"users": users, "total": len(users)}
+
+
+async def update_all_budgets_controller(
+    max_budget: float, budget_duration: str
+) -> dict:
+    """
+    Apply the same max budget and reset frequency to every managed LiteLLM user.
+    """
+
+    try:
+        return await update_all_user_budgets(max_budget, budget_duration)
+
+    except httpx.HTTPError as error:
+        raise HTTPException(
+            status_code=502,
+            detail=f"Failed to update budgets in LiteLLM: {error}",
+        )

backend/app/core/config.py

@@ -43,7 +43,7 @@ class Settings(BaseSettings):
     # LiteLLM Proxy Config
     LITELLM_URL: str = "http://litellm:4000"
     LITELLM_MASTER_KEY: Optional[str] = None
-    DEFAULT_USER_BUDGET: Decimal = Decimal("10.00")
+    DEFAULT_USER_BUDGET: Decimal = Decimal("2.00")
     DEFAULT_BUDGET_DURATION: str = "30d"
 
 

backend/app/main.py

@@ -16,6 +16,7 @@
 from app.routes.knowledge_base_router import router as knowledge_base_router
 from app.routes.users_router import router as users_router
 from app.routes.vcell_identity_router import router as vcell_identity_router
+from app.routes.litellm_router import router as litellm_router
 
 ascii_art = """
 ╔════════════════════════════════════════════════════════════════════════════════════╗
@@ -60,6 +61,7 @@ async def startup_event():
 app.include_router(qdrant_router, tags=["Qdrant Vector DB"], prefix="/qdrant")
 app.include_router(users_router, tags=["Users"])
 app.include_router(vcell_identity_router, tags=["VCell Account Linking"])
+app.include_router(litellm_router, tags=["LiteLLM Admin"], prefix="/litellm")
 
 if __name__ == "__main__":
     uvicorn.run("app.main:app", host="0.0.0.0", port=8000, reload=True)

backend/app/routes/litellm_router.py

@@ -0,0 +1,39 @@
+from fastapi import APIRouter, Depends
+
+from app.controllers.litellm_controller import (
+    list_managed_users_controller,
+    update_all_budgets_controller,
+)
+from app.core.auth import require_admin
+from app.schemas.litellm_schema import (
+    ManagedUsersResponse,
+    UpdateAllBudgetsRequest,
+    UpdateAllBudgetsResponse,
+)
+
+# All LiteLLM management endpoints are admin-only.
+router = APIRouter(dependencies=[Depends(require_admin)])
+
+
+@router.get("/users", response_model=ManagedUsersResponse)
+async def list_litellm_users():
+    """
+    Endpoint to list the LiteLLM users a bulk budget update applies to.
+
+    Excludes the LiteLLM proxy admin / default user, whose budget stays
+    unlimited.
+    """
+
+    return await list_managed_users_controller()
+
+
+@router.post("/budgets", response_model=UpdateAllBudgetsResponse)
+async def update_all_litellm_budgets(request: UpdateAllBudgetsRequest):
+    """
+    Endpoint to set the same max budget and reset frequency on every LiteLLM
+    user except the proxy admin / default user.
+    """
+
+    return await update_all_budgets_controller(
+        request.max_budget, request.budget_duration
+    )

backend/app/schemas/litellm_schema.py

@@ -0,0 +1,33 @@
+from typing import Literal, Optional
+
+from pydantic import BaseModel, Field
+
+# Mirrors LiteLLM's budget reset options (hourly / daily / weekly / monthly).
+BudgetDuration = Literal["1h", "24h", "7d", "30d"]
+
+
+class UpdateAllBudgetsRequest(BaseModel):
+    max_budget: float = Field(ge=0, description="Max budget in USD per user")
+    budget_duration: BudgetDuration = Field(
+        description="How often each user's budget resets"
+    )
+
+
+class ManagedUser(BaseModel):
+    user_id: str
+    user_email: Optional[str] = None
+    spend: float
+    max_budget: Optional[float] = None
+    budget_duration: Optional[str] = None
+
+
+class ManagedUsersResponse(BaseModel):
+    users: list[ManagedUser]
+    total: int
+
+
+class UpdateAllBudgetsResponse(BaseModel):
+    total_users: int
+    successful_updates: int
+    failed_updates: int
+    failed_user_ids: list[str] = []

backend/app/services/litellm_service.py

@@ -98,3 +98,129 @@ async def get_user_budget_info(auth0_sub: str) -> dict:
         "max_budget": max_budget,
         "remaining_budget": remaining_budget,
     }
+
+
+# LiteLLM's built-in proxy admin. Its budget must stay unlimited, so it is
+# always excluded from bulk budget updates.
+LITELLM_EXCLUDED_USER_IDS = {"default_user_id"}
+LITELLM_EXCLUDED_ROLES = {"proxy_admin", "proxy_admin_viewer"}
+
+# /user/list caps page_size at 100.
+_USER_LIST_PAGE_SIZE = 100
+
+
+def _is_excluded_from_bulk_update(user: dict) -> bool:
+    """
+    Return True for LiteLLM users that must never be touched by a bulk budget
+    update (the proxy admin / default user).
+    """
+    return (
+        user.get("user_id") in LITELLM_EXCLUDED_USER_IDS
+        or user.get("user_role") in LITELLM_EXCLUDED_ROLES
+    )
+
+
+async def list_managed_users() -> list[dict]:
+    """
+    List every LiteLLM user eligible for bulk budget updates, i.e. all users
+    except the proxy admin / default user.
+
+    Returns:
+        list[dict]: user_id, user_email, spend, max_budget and budget_duration
+            for each managed user.
+    """
+    users: list[dict] = []
+    page = 1
+
+    async with httpx.AsyncClient(timeout=30.0) as client:
+        while True:
+            response = await client.get(
+                f"{settings.LITELLM_URL}/user/list",
+                headers={"Authorization": f"Bearer {settings.LITELLM_MASTER_KEY}"},
+                params={"page": page, "page_size": _USER_LIST_PAGE_SIZE},
+            )
+            response.raise_for_status()
+            data = response.json()
+
+            users.extend(data.get("users") or [])
+
+            if page >= (data.get("total_pages") or 1):
+                break
+            page += 1
+
+    return [
+        {
+            "user_id": user.get("user_id"),
+            "user_email": user.get("user_email"),
+            "spend": user.get("spend") or 0.0,
+            "max_budget": user.get("max_budget"),
+            "budget_duration": user.get("budget_duration"),
+        }
+        for user in users
+        if not _is_excluded_from_bulk_update(user)
+    ]
+
+
+async def update_all_user_budgets(max_budget: float, budget_duration: str) -> dict:
+    """
+    Set the same max budget and reset frequency on every managed LiteLLM user.
+
+    LiteLLM's own "update all users" option would also overwrite the proxy
+    admin's unlimited budget, so the user list is fetched and filtered first and
+    the updates are sent as an explicit per-user batch.
+
+    Args:
+        max_budget (float): Max budget in USD to apply to each managed user.
+        budget_duration (str): LiteLLM duration string the budget resets on
+            (e.g. "1h", "24h", "7d", "30d").
+
+    Returns:
+        dict: total_users, successful_updates, failed_updates and the user_ids
+            of any users that could not be updated.
+    """
+    managed_users = await list_managed_users()
+
+    if not managed_users:
+        return {
+            "total_users": 0,
+            "successful_updates": 0,
+            "failed_updates": 0,
+            "failed_user_ids": [],
+        }
+
+    async with httpx.AsyncClient(timeout=60.0) as client:
+        response = await client.post(
+            f"{settings.LITELLM_URL}/user/bulk_update",
+            headers={"Authorization": f"Bearer {settings.LITELLM_MASTER_KEY}"},
+            json={
+                "users": [
+                    {
+                        "user_id": user["user_id"],
+                        "max_budget": max_budget,
+                        "budget_duration": budget_duration,
+                    }
+                    for user in managed_users
+                ]
+            },
+        )
+        response.raise_for_status()
+        data = response.json()
+
+    failed_user_ids = [
+        result.get("user_id")
+        for result in (data.get("results") or [])
+        if not result.get("success")
+    ]
+
+    logger.info(
+        f"Bulk budget update: {data.get('successful_updates')} succeeded, "
+        f"{data.get('failed_updates')} failed "
+        f"(max_budget={max_budget}, budget_duration={budget_duration})"
+    )
+
+    return {
+        "total_users": data.get("total_requested") or len(managed_users),
+        "successful_updates": data.get("successful_updates") or 0,
+        "failed_updates": data.get("failed_updates") or 0,
+        "failed_user_ids": failed_user_ids,
+    }