mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
feat(ptu): runtime UI flag for enable_ptu_cost_attribution + always-registered rollup
Add enable_ptu_cost_attribution to UISettings so admins can flip it from the UI's Admin Settings page without a proxy restart. Register the daily rollup job unconditionally at startup; it already no-ops on flag-off, so the runtime toggle takes effect on the next scheduled fire without any restart. Existing per-request guards in the /ptu_reservation endpoints and in run_ptu_reservation_rollup remain the source of truth. Persisted settings sync into general_settings on both PATCH and GET paths as well as on server startup via _sync_ui_settings_to_general_settings.
This commit is contained in:
parent
5808dcf079
commit
1b85426a6a
6 changed files with 183 additions and 16 deletions
|
|
@ -7930,23 +7930,24 @@ class ProxyStartupEvent:
|
|||
await cls._initialize_spend_tracking_background_jobs(scheduler=scheduler)
|
||||
|
||||
### PTU RESERVATION DAILY ROLLUP ###
|
||||
if general_settings.get("enable_ptu_cost_attribution", False):
|
||||
from litellm.proxy.spend_tracking.ptu_reservation_rollup import (
|
||||
PTU_ROLLUP_JOB_ID,
|
||||
run_ptu_reservation_rollup,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.ptu_reservation_rollup import (
|
||||
PTU_ROLLUP_JOB_ID,
|
||||
run_ptu_reservation_rollup,
|
||||
)
|
||||
|
||||
scheduler.add_job(
|
||||
run_ptu_reservation_rollup,
|
||||
"cron",
|
||||
hour=0,
|
||||
minute=15,
|
||||
args=[prisma_client],
|
||||
id=PTU_ROLLUP_JOB_ID,
|
||||
replace_existing=True,
|
||||
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
|
||||
)
|
||||
verbose_proxy_logger.info("PTU reservation rollup job scheduled at 00:15 UTC daily")
|
||||
scheduler.add_job(
|
||||
run_ptu_reservation_rollup,
|
||||
"cron",
|
||||
hour=0,
|
||||
minute=15,
|
||||
args=[prisma_client],
|
||||
id=PTU_ROLLUP_JOB_ID,
|
||||
replace_existing=True,
|
||||
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
|
||||
)
|
||||
verbose_proxy_logger.info(
|
||||
"PTU reservation rollup job scheduled at 00:15 UTC daily (no-ops until enable_ptu_cost_attribution is set)"
|
||||
)
|
||||
|
||||
### SPEND LOG CLEANUP ###
|
||||
if general_settings.get("maximum_spend_logs_retention_period") is not None:
|
||||
|
|
|
|||
|
|
@ -182,6 +182,11 @@ class UISettings(BaseModel):
|
|||
description="If true, shows the Chat page in the UI sidebar, letting users chat with an LLM and connect their own MCP server credentials via OAuth.",
|
||||
)
|
||||
|
||||
enable_ptu_cost_attribution: bool = Field(
|
||||
default=False,
|
||||
description="If true, enables admin-registered PTU reservations and daily flat-cost attribution on team daily spend. Governs the /ptu_reservation CRUD endpoints, the daily rollup job, the PTU Reservations UI page, and the Flat Cost column on the Usage page.",
|
||||
)
|
||||
|
||||
|
||||
class UISettingsResponse(SettingsResponse):
|
||||
"""Response model for UI settings"""
|
||||
|
|
@ -206,6 +211,7 @@ ALLOWED_UI_SETTINGS_FIELDS = {
|
|||
"disable_custom_api_keys",
|
||||
"disable_key_generate_for_org_admin",
|
||||
"enable_chat_ui",
|
||||
"enable_ptu_cost_attribution",
|
||||
}
|
||||
|
||||
# Flags that must be synced from the persisted UISettings into
|
||||
|
|
@ -219,6 +225,7 @@ _RUNTIME_GENERAL_SETTINGS_FLAGS = [
|
|||
"disable_vector_stores_for_internal_users",
|
||||
"allow_vector_stores_for_team_admins",
|
||||
"disable_key_generate_for_org_admin",
|
||||
"enable_ptu_cost_attribution",
|
||||
]
|
||||
|
||||
# Extension point: packages outside OSS (e.g. litellm_enterprise) can
|
||||
|
|
|
|||
|
|
@ -9718,3 +9718,41 @@ async def test_startup_survives_database_read_failure_for_coordination_redis():
|
|||
)
|
||||
|
||||
assert result is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ptu_rollup_job_registered_regardless_of_flag(monkeypatch):
|
||||
"""The PTU rollup cron must be registered at startup so a UI-time flip of enable_ptu_cost_attribution takes effect without a proxy restart. The job itself no-ops when the flag is off (asserted in test_ptu_reservation_rollup.py)."""
|
||||
monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False)
|
||||
from litellm.proxy.proxy_server import ProxyStartupEvent
|
||||
from litellm.proxy.spend_tracking.ptu_reservation_rollup import (
|
||||
PTU_ROLLUP_JOB_ID,
|
||||
)
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.db.litellm_config.find_first = AsyncMock(return_value=None)
|
||||
|
||||
mock_proxy_logging = MagicMock(spec=ProxyLogging)
|
||||
mock_proxy_logging.slack_alerting_instance = MagicMock()
|
||||
mock_proxy_config = AsyncMock()
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config),
|
||||
patch("litellm.proxy.proxy_server.store_model_in_db", True),
|
||||
patch("litellm.proxy.proxy_server.get_secret_bool", return_value=True),
|
||||
):
|
||||
# Flag OFF at startup — job must still be registered
|
||||
await ProxyStartupEvent.initialize_scheduled_background_jobs(
|
||||
general_settings={"enable_ptu_cost_attribution": False},
|
||||
prisma_client=mock_prisma_client,
|
||||
proxy_budget_rescheduler_min_time=1,
|
||||
proxy_budget_rescheduler_max_time=2,
|
||||
proxy_batch_write_at=5,
|
||||
proxy_logging_obj=mock_proxy_logging,
|
||||
)
|
||||
|
||||
import litellm.proxy.proxy_server as ps
|
||||
|
||||
assert ps.scheduler is not None
|
||||
assert ps.scheduler.get_job(PTU_ROLLUP_JOB_ID) is not None
|
||||
|
|
|
|||
|
|
@ -1336,6 +1336,61 @@ class TestProxySettingEndpoints:
|
|||
# Synced into general_settings so the enforcement helper sees it
|
||||
assert general_settings.get(flag_name) is True
|
||||
|
||||
def test_update_ui_settings_persists_and_syncs_enable_ptu_cost_attribution(
|
||||
self, mock_auth, monkeypatch
|
||||
):
|
||||
"""enable_ptu_cost_attribution must be allowlisted, persisted, and synced to general_settings so PTU endpoints and rollup pick it up without a proxy restart."""
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
mock_user_auth = UserAPIKeyAuth(
|
||||
user_id="test-user-123",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True)
|
||||
|
||||
general_settings: dict = {}
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.general_settings", general_settings
|
||||
)
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_uisettings.upsert = AsyncMock()
|
||||
mock_prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
|
||||
|
||||
flag_name = "enable_ptu_cost_attribution"
|
||||
payload = {flag_name: True}
|
||||
|
||||
try:
|
||||
response = client.patch("/update/ui_settings", json=payload)
|
||||
finally:
|
||||
app.dependency_overrides.clear()
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["settings"][flag_name] is True
|
||||
|
||||
call_args = mock_prisma.db.litellm_uisettings.upsert.call_args
|
||||
stored_settings = json.loads(call_args.kwargs["data"]["create"]["ui_settings"])
|
||||
assert stored_settings[flag_name] is True
|
||||
|
||||
assert general_settings.get(flag_name) is True
|
||||
|
||||
def test_enable_ptu_cost_attribution_registered_as_runtime_flag(self):
|
||||
"""The runtime-sync list must include the PTU flag so a UI flip lands in general_settings without a restart."""
|
||||
from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
|
||||
_RUNTIME_GENERAL_SETTINGS_FLAGS,
|
||||
ALLOWED_UI_SETTINGS_FIELDS,
|
||||
)
|
||||
|
||||
assert "enable_ptu_cost_attribution" in _RUNTIME_GENERAL_SETTINGS_FLAGS
|
||||
assert "enable_ptu_cost_attribution" in ALLOWED_UI_SETTINGS_FIELDS
|
||||
|
||||
def test_get_sso_settings_from_database(
|
||||
self, mock_proxy_config, mock_auth, monkeypatch
|
||||
):
|
||||
|
|
|
|||
|
|
@ -41,12 +41,16 @@ const buildSettingsResponse = (overrides?: Partial<Record<string, unknown>>) =>
|
|||
require_auth_for_public_ai_hub: {
|
||||
description: "Require authentication for public AI Hub",
|
||||
},
|
||||
enable_ptu_cost_attribution: {
|
||||
description: "Enable PTU cost attribution",
|
||||
},
|
||||
},
|
||||
},
|
||||
values: {
|
||||
disable_model_add_for_internal_users: false,
|
||||
disable_team_admin_delete_team_user: false,
|
||||
require_auth_for_public_ai_hub: false,
|
||||
enable_ptu_cost_attribution: false,
|
||||
},
|
||||
},
|
||||
isLoading: false,
|
||||
|
|
@ -162,4 +166,33 @@ describe("UISettings", () => {
|
|||
);
|
||||
expect(NotificationManager.success).toHaveBeenCalledWith("UI settings updated successfully");
|
||||
});
|
||||
|
||||
it("should toggle enable PTU cost attribution setting and call update with page-refresh notice", () => {
|
||||
const mutateMock = vi.fn((_settings, options) => {
|
||||
options?.onSuccess?.();
|
||||
});
|
||||
|
||||
mockUseUpdateUISettings.mockReturnValue({
|
||||
mutate: mutateMock,
|
||||
isPending: false,
|
||||
error: null,
|
||||
});
|
||||
|
||||
render(<UISettings />);
|
||||
|
||||
const toggle = screen.getByRole("switch", { name: "Enable PTU cost attribution" });
|
||||
|
||||
act(() => {
|
||||
fireEvent.click(toggle);
|
||||
});
|
||||
|
||||
expect(mutateMock).toHaveBeenCalledWith(
|
||||
{ enable_ptu_cost_attribution: true },
|
||||
expect.objectContaining({
|
||||
onSuccess: expect.any(Function),
|
||||
onError: expect.any(Function),
|
||||
}),
|
||||
);
|
||||
expect(NotificationManager.success).toHaveBeenCalledWith("UI settings updated successfully. Refreshing page...");
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ export default function UISettings() {
|
|||
const forwardLLMProviderAuthHeadersProperty = schema?.properties?.forward_llm_provider_auth_headers;
|
||||
const enableProjectsUIProperty = schema?.properties?.enable_projects_ui;
|
||||
const enableChatUIProperty = schema?.properties?.enable_chat_ui;
|
||||
const enablePtuCostAttributionProperty = schema?.properties?.enable_ptu_cost_attribution;
|
||||
const enabledPagesProperty = schema?.properties?.enabled_ui_pages_internal_users;
|
||||
const disableAgentsProperty = schema?.properties?.disable_agents_for_internal_users;
|
||||
const allowAgentsTeamAdminsProperty = schema?.properties?.allow_agents_for_team_admins;
|
||||
|
|
@ -130,6 +131,21 @@ export default function UISettings() {
|
|||
);
|
||||
};
|
||||
|
||||
const handleToggleEnablePtuCostAttribution = (checked: boolean) => {
|
||||
updateSettings(
|
||||
{ enable_ptu_cost_attribution: checked },
|
||||
{
|
||||
onSuccess: () => {
|
||||
NotificationManager.success("UI settings updated successfully. Refreshing page...");
|
||||
setTimeout(() => window.location.reload(), 1000);
|
||||
},
|
||||
onError: (error) => {
|
||||
NotificationManager.fromBackend(error);
|
||||
},
|
||||
},
|
||||
);
|
||||
};
|
||||
|
||||
const handleToggleRequireAuthForPublicAIHub = (checked: boolean) => {
|
||||
updateSettings(
|
||||
{ require_auth_for_public_ai_hub: checked },
|
||||
|
|
@ -368,6 +384,23 @@ export default function UISettings() {
|
|||
</Space>
|
||||
</Space>
|
||||
|
||||
<Space align="start" size="middle">
|
||||
<Switch
|
||||
checked={Boolean(values.enable_ptu_cost_attribution)}
|
||||
disabled={isUpdating}
|
||||
loading={isUpdating}
|
||||
onChange={handleToggleEnablePtuCostAttribution}
|
||||
aria-label={enablePtuCostAttributionProperty?.description ?? "Enable PTU cost attribution"}
|
||||
/>
|
||||
<Space direction="vertical" size={4}>
|
||||
<Typography.Text strong>Enable PTU cost attribution (page will refresh)</Typography.Text>
|
||||
<Typography.Text type="secondary">
|
||||
{enablePtuCostAttributionProperty?.description ??
|
||||
"If enabled, unlocks admin-registered PTU reservations and daily flat-cost attribution on team daily spend."}
|
||||
</Typography.Text>
|
||||
</Space>
|
||||
</Space>
|
||||
|
||||
<Divider />
|
||||
|
||||
{/* Agents access control */}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue