mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
feat(ptu): gate PTU flat-cost attribution behind an opt-in env var (#36138)
LITELLM_ENABLE_PTU_COST_ATTRIBUTION, read through get_secret_bool and defaulting to
false, makes the whole PTU flat-cost feature inert unless an operator opts in. The
daily rollup cron is not registered at all, so no sentinel row is ever written;
/model/new and /model/{id}/update reject a request that carries any PTU model_info
field with a 400 naming the fields and the env var rather than dropping them; the
daily activity read path reports zero flat cost; and the model add and edit forms
hide the four PTU inputs.
The read gate lives where flat cost enters SpendMetrics rather than in the aggregated
SQL select. /team/daily/activity, the endpoint the Usage page reads, is served by the
paginated find_many path and never runs that query, so forcing the select to a
constant zero would have left the reporting surface that matters still showing flat
cost.
Sentinel row filtering is deliberately not gated. An operator can enable the flag,
accrue rows under the __ptu_flat_cost__ api_key, then disable it, and those rows stay
in LiteLLM_DailyTeamSpend; gating the filter too would surface the sentinel as a bogus
api_key and mint a provider bucket for its empty provider. Response fields keep their
shape and report 0.0, so typed clients are unaffected, and the migration and the
ModelInfo field declarations are untouched.
The write gate reads the incoming request rather than the merged deployment, so a
model configured during an earlier opt-in stays editable, and the edit form drops the
PTU keys from the payload instead of sending nulls that would clear stored config.
The dashboard reads the flag from a read-only enable_ptu_cost_attribution key on
/get/ui_settings, computed from the environment on every read. It is deliberately not
an allowlisted persisted setting, and PATCH /update/ui_settings rejects it with a 400,
so an admin cannot flip an env-gated feature from the UI.
Two review findings on the gate itself. The PTU clear loop now runs only when the
feature is enabled: the write gate rejects a value but lets an explicit null through,
and a client round-tripping a model_info blob sends the PTU keys as nulls, so a
disabled proxy would have quietly erased a billing configuration set up during an
earlier opt-in. Disabling pauses PTU rather than discarding its setup. And the
dashboard flag is re-read every thirty seconds instead of the hour the other UI settings
use, since those are persisted records while this one tracks the proxy process; a
restart that flips the variable would otherwise leave the model form offering inputs
the backend now rejects. The flag is polled rather than only marked stale, since a form
that stays mounted and focused never refetches on its own.
The read gate checks the row before the flag. It runs once per metric accumulation and a
record fans out across roughly a dozen breakdowns, while the flag reads through the secret
manager uncached, so consulting it for every accumulation put thousands of lookups on a
shared endpoint that made none before. Only a row actually carrying flat cost reaches it.
This commit is contained in:
parent
9de3315dad
commit
e014b341c8
21 changed files with 1513 additions and 184 deletions
|
|
@ -10,6 +10,7 @@ from typing_extensions import TypedDict
|
|||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import PTU_SENTINEL_API_KEY
|
||||
from litellm.proxy._types import CommonProxyErrors
|
||||
from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
from litellm.repositories.table_repositories import DeletedVerificationTokenRepository
|
||||
from litellm.repositories.verification_token_repository import (
|
||||
|
|
@ -141,6 +142,28 @@ class _GroupingSetsRow(SimpleNamespace):
|
|||
failed_requests: int | None
|
||||
|
||||
|
||||
def _reported_flat_cost(record: DailySpendRecord | _GroupingSetsRow) -> float:
|
||||
"""Flat cost a daily row reports, which is zero unless PTU cost attribution is enabled.
|
||||
|
||||
Both read paths funnel through here: the paginated path reads the ``ptu_flat_cost``
|
||||
column straight off the row, and the aggregated path reads the SUM() alias. Rows an
|
||||
operator accrued during an earlier opt-in stay in the table, so the gate lives on the
|
||||
read rather than on the query that produced the rows.
|
||||
|
||||
The row is checked before the flag because this runs once per metric accumulation, and
|
||||
a record fans out across roughly a dozen breakdowns. The flag reads through the secret
|
||||
manager, uncached, so consulting it for every accumulation put thousands of lookups on
|
||||
a shared endpoint that made none before. Only a row actually carrying flat cost, which
|
||||
is a sentinel row, reaches it now.
|
||||
"""
|
||||
raw: Final = getattr(record, "ptu_flat_cost", None) or 0.0
|
||||
if not raw:
|
||||
return 0.0
|
||||
if not is_ptu_cost_attribution_enabled():
|
||||
return 0.0
|
||||
return raw
|
||||
|
||||
|
||||
def update_metrics(existing_metrics: SpendMetrics, record: DailySpendRecord) -> SpendMetrics:
|
||||
"""Update metrics with new record data.
|
||||
|
||||
|
|
@ -151,7 +174,7 @@ def update_metrics(existing_metrics: SpendMetrics, record: DailySpendRecord) ->
|
|||
prompt_tokens: Final = record.prompt_tokens or 0
|
||||
completion_tokens: Final = record.completion_tokens or 0
|
||||
existing_metrics.spend += record.spend or 0.0
|
||||
existing_metrics.flat_cost += getattr(record, "ptu_flat_cost", None) or 0.0
|
||||
existing_metrics.flat_cost += _reported_flat_cost(record)
|
||||
existing_metrics.prompt_tokens += prompt_tokens
|
||||
existing_metrics.completion_tokens += completion_tokens
|
||||
existing_metrics.total_tokens += prompt_tokens + completion_tokens
|
||||
|
|
@ -784,7 +807,7 @@ def _record_to_spend_metrics(record: _GroupingSetsRow) -> SpendMetrics:
|
|||
completion_tokens: Final = record.completion_tokens or 0
|
||||
return SpendMetrics(
|
||||
spend=record.spend or 0.0,
|
||||
flat_cost=getattr(record, "ptu_flat_cost", None) or 0.0,
|
||||
flat_cost=_reported_flat_cost(record),
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
total_tokens=prompt_tokens + completion_tokens,
|
||||
|
|
|
|||
|
|
@ -56,6 +56,10 @@ from litellm.proxy.management_endpoints.team_endpoints import (
|
|||
update_team as _legacy_update_team,
|
||||
)
|
||||
from litellm.proxy.management_helpers.audit_logs import create_object_audit_log
|
||||
from litellm.proxy.spend_tracking.ptu_feature_flag import (
|
||||
PTU_COST_ATTRIBUTION_ENV_VAR,
|
||||
is_ptu_cost_attribution_enabled,
|
||||
)
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from litellm.repositories.model_repository import ModelRepository
|
||||
from litellm.repositories.table_repositories import ModelTableRepository
|
||||
|
|
@ -239,8 +243,12 @@ _PTU_MODEL_INFO_FIELDS: Final = ("ptu_count", "cost_per_ptu_per_hour", "ptu_effe
|
|||
|
||||
|
||||
def _explicitly_cleared_ptu_fields(model_info: ModelInfo | None) -> frozenset[str]:
|
||||
"""The PTU fields a patch sends as an explicit null, which update_db_model drops."""
|
||||
if model_info is None:
|
||||
"""The PTU fields a patch sends as an explicit null, which update_db_model drops.
|
||||
|
||||
Empty while the feature is off, so disabling pauses PTU rather than letting a client
|
||||
that round-trips a model_info blob erase a configuration set up during an earlier opt-in.
|
||||
"""
|
||||
if model_info is None or not is_ptu_cost_attribution_enabled():
|
||||
return frozenset()
|
||||
return frozenset(
|
||||
field
|
||||
|
|
@ -262,6 +270,32 @@ def _merged_ptu_model_info(*, db_model: Deployment, patch_data: updateDeployment
|
|||
return MappingProxyType({k: v for k, v in {**stored, **incoming}.items() if k not in cleared})
|
||||
|
||||
|
||||
def _raise_if_ptu_cost_attribution_disabled(incoming_model_info: Mapping[str, object]) -> None:
|
||||
"""Reject PTU model_info fields unless the operator opted into PTU cost attribution.
|
||||
|
||||
Takes the incoming request's model_info rather than the merged deployment, so an
|
||||
unrelated patch of a model that still stores PTU config from an earlier opt-in is
|
||||
left alone. The fields are rejected rather than dropped so a caller never believes
|
||||
a flat cost was configured while the rollup that would price it is not running.
|
||||
|
||||
Only a value is rejected. An explicit null reaches the clear loop, which is gated on
|
||||
the same flag, so a disabled proxy neither writes PTU config nor erases what an
|
||||
earlier opt-in stored. Disabling pauses the feature rather than discarding its setup.
|
||||
"""
|
||||
if is_ptu_cost_attribution_enabled():
|
||||
return
|
||||
supplied: Final = tuple(field for field in _PTU_MODEL_INFO_FIELDS if incoming_model_info.get(field) is not None)
|
||||
if not supplied:
|
||||
return
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=(
|
||||
f"PTU cost attribution is disabled, so {', '.join(supplied)} cannot be set. "
|
||||
f"Set {PTU_COST_ATTRIBUTION_ENV_VAR}=true to enable it."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _validate_ptu_model_info(model_info: Mapping[str, object]) -> None:
|
||||
"""Enforce the PTU cross-field invariant on the effective model_info.
|
||||
|
||||
|
|
@ -326,6 +360,8 @@ def _coerce_ptu_datetime(value: object) -> datetime.datetime | None:
|
|||
|
||||
|
||||
def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> PrismaCompatibleUpdateDBModel:
|
||||
if updated_patch.model_info is not None:
|
||||
_raise_if_ptu_cost_attribution_disabled(updated_patch.model_info.model_dump(exclude_none=True))
|
||||
merged_model_name: Final = updated_patch.model_name or db_model.model_name
|
||||
merged_litellm_params: Final = db_model.litellm_params.model_dump(exclude_none=True)
|
||||
merged_model_info: Final = db_model.model_info.model_dump(exclude_none=True)
|
||||
|
|
@ -821,6 +857,7 @@ async def _update_team_model_in_db(
|
|||
# raising the rate on a configured model carries no ptu_effective_from, which the
|
||||
# stored row supplies.
|
||||
if patch_data.model_info is not None:
|
||||
_raise_if_ptu_cost_attribution_disabled(patch_data.model_info.model_dump(exclude_none=True))
|
||||
_validate_ptu_model_info(_merged_ptu_model_info(db_model=db_model, patch_data=patch_data))
|
||||
|
||||
patch_team_id: Final = patch_data.model_info.team_id if patch_data.model_info else None
|
||||
|
|
@ -1531,7 +1568,9 @@ async def add_new_model(
|
|||
|
||||
model_response: LiteLLM_ProxyModelTable | None = None
|
||||
# update DB
|
||||
_validate_ptu_model_info(model_params.model_info.model_dump(exclude_none=True))
|
||||
incoming_model_info: Final = model_params.model_info.model_dump(exclude_none=True)
|
||||
_raise_if_ptu_cost_attribution_disabled(incoming_model_info)
|
||||
_validate_ptu_model_info(incoming_model_info)
|
||||
|
||||
if store_model_in_db is True:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -8471,40 +8471,45 @@ class ProxyStartupEvent:
|
|||
await cls._initialize_spend_tracking_background_jobs(scheduler=scheduler)
|
||||
|
||||
### PTU DAILY ROLLUP ###
|
||||
from litellm.proxy.spend_tracking.ptu_flat_cost_rollup import (
|
||||
PTU_ROLLUP_JOB_ID,
|
||||
run_scheduled_ptu_rollup,
|
||||
from litellm.proxy.spend_tracking.ptu_feature_flag import (
|
||||
is_ptu_cost_attribution_enabled,
|
||||
)
|
||||
|
||||
async def _alert_ptu_rollup_failure(message: str) -> None:
|
||||
await proxy_logging_obj.alerting_handler(
|
||||
message=message,
|
||||
level="High",
|
||||
alert_type=AlertType.failed_tracking_spend,
|
||||
if is_ptu_cost_attribution_enabled():
|
||||
from litellm.proxy.spend_tracking.ptu_flat_cost_rollup import (
|
||||
PTU_ROLLUP_JOB_ID,
|
||||
run_scheduled_ptu_rollup,
|
||||
)
|
||||
|
||||
async def _scheduled_ptu_rollup() -> None:
|
||||
# Reuse the PodLockManager from db_spend_update_writer so only one pod
|
||||
# reconciles a day; a multi-pod race could prune another pod's fresh rows
|
||||
await run_scheduled_ptu_rollup(
|
||||
prisma_client,
|
||||
pod_lock_manager=proxy_logging_obj.db_spend_update_writer.pod_lock_manager,
|
||||
alert=_alert_ptu_rollup_failure,
|
||||
)
|
||||
async def _alert_ptu_rollup_failure(message: str) -> None:
|
||||
await proxy_logging_obj.alerting_handler(
|
||||
message=message,
|
||||
level="High",
|
||||
alert_type=AlertType.failed_tracking_spend,
|
||||
)
|
||||
|
||||
scheduler.add_job(
|
||||
_scheduled_ptu_rollup,
|
||||
"cron",
|
||||
hour=0,
|
||||
minute=15,
|
||||
timezone="UTC",
|
||||
id=PTU_ROLLUP_JOB_ID,
|
||||
replace_existing=True,
|
||||
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
|
||||
)
|
||||
verbose_proxy_logger.info(
|
||||
"PTU rollup job scheduled at 00:15 UTC daily (only models with PTU config accrue flat cost)"
|
||||
)
|
||||
async def _scheduled_ptu_rollup() -> None:
|
||||
# Reuse the PodLockManager from db_spend_update_writer so only one pod
|
||||
# reconciles a day; a multi-pod race could prune another pod's fresh rows
|
||||
await run_scheduled_ptu_rollup(
|
||||
prisma_client,
|
||||
pod_lock_manager=proxy_logging_obj.db_spend_update_writer.pod_lock_manager,
|
||||
alert=_alert_ptu_rollup_failure,
|
||||
)
|
||||
|
||||
scheduler.add_job(
|
||||
_scheduled_ptu_rollup,
|
||||
"cron",
|
||||
hour=0,
|
||||
minute=15,
|
||||
timezone="UTC",
|
||||
id=PTU_ROLLUP_JOB_ID,
|
||||
replace_existing=True,
|
||||
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
|
||||
)
|
||||
verbose_proxy_logger.info(
|
||||
"PTU rollup job scheduled at 00:15 UTC daily (only models with PTU config accrue flat cost)"
|
||||
)
|
||||
|
||||
### SPEND LOG CLEANUP ###
|
||||
if (
|
||||
|
|
|
|||
18
litellm/proxy/spend_tracking/ptu_feature_flag.py
Normal file
18
litellm/proxy/spend_tracking/ptu_feature_flag.py
Normal file
|
|
@ -0,0 +1,18 @@
|
|||
"""Opt-in flag for PTU (provisioned throughput unit) flat-cost attribution.
|
||||
|
||||
The whole feature is inert unless an operator sets
|
||||
``LITELLM_ENABLE_PTU_COST_ATTRIBUTION``: the daily rollup is not scheduled, the
|
||||
model endpoints reject PTU config, the daily activity read path reports zero flat
|
||||
cost, and the model form hides the PTU inputs.
|
||||
"""
|
||||
|
||||
from typing import Final
|
||||
|
||||
from litellm.secret_managers.main import get_secret_bool
|
||||
|
||||
PTU_COST_ATTRIBUTION_ENV_VAR: Final = "LITELLM_ENABLE_PTU_COST_ATTRIBUTION"
|
||||
|
||||
|
||||
def is_ptu_cost_attribution_enabled() -> bool:
|
||||
"""Report whether this deployment opted into PTU flat-cost attribution."""
|
||||
return get_secret_bool(PTU_COST_ATTRIBUTION_ENV_VAR, False) is True
|
||||
|
|
@ -27,6 +27,7 @@ from litellm.constants import (
|
|||
PTU_ROLLUP_MAX_BACKFILL_DAYS,
|
||||
PTU_SENTINEL_API_KEY,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled
|
||||
from litellm.types.router import ModelInfo
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -512,7 +513,15 @@ async def run_scheduled_ptu_rollup(
|
|||
duplicate work rather than correctness: the upserts are idempotent on the sentinel
|
||||
key and the prune reads only the row's own timestamp, so a second pod arriving
|
||||
mid-run cannot corrupt the day.
|
||||
|
||||
Returns None without touching the database when PTU cost attribution is off. Proxy
|
||||
startup already skips scheduling the cron, so this guards the function itself rather
|
||||
than its one caller, and a deployment that never opted in accrues nothing whatever
|
||||
reaches it.
|
||||
"""
|
||||
if not is_ptu_cost_attribution_enabled():
|
||||
return None
|
||||
|
||||
if pod_lock_manager is None or pod_lock_manager.redis_cache is None:
|
||||
return await _run_and_alert(prisma_client, target_date=target_date, alert=alert, may_prune=False)
|
||||
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ from litellm.proxy.config_resolvers.sso import (
|
|||
SSO_SECRET_FIELDS,
|
||||
resolve_sso_config,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled
|
||||
from litellm.proxy.utils import invalidate_config_param
|
||||
from litellm.repositories.config_repository import ConfigRepository
|
||||
from litellm.repositories.organization_repository import OrganizationRepository
|
||||
|
|
@ -307,6 +308,27 @@ ALLOWED_UI_SETTINGS_FIELDS: Final = {
|
|||
"enable_chat_ui",
|
||||
}
|
||||
|
||||
ENABLE_PTU_COST_ATTRIBUTION_UI_SETTING: Final = "enable_ptu_cost_attribution"
|
||||
|
||||
# UI settings derived from the deployment environment. Deliberately kept out of
|
||||
# ALLOWED_UI_SETTINGS_FIELDS: they are read-only, never persisted, and PATCH
|
||||
# rejects them so an admin cannot flip an env-gated feature at runtime.
|
||||
_DERIVED_UI_SETTINGS_FIELDS: Final[frozenset[str]] = frozenset({ENABLE_PTU_COST_ATTRIBUTION_UI_SETTING})
|
||||
|
||||
|
||||
def _derived_ui_setting_value(key: str) -> object:
|
||||
"""The environment-derived value GET reports for ``key``.
|
||||
|
||||
PATCH compares against this rather than rejecting the key outright, so the body GET
|
||||
hands back is still a valid PATCH body. Rejecting on presence broke read-modify-write:
|
||||
a client that edited one setting and sent the rest back unchanged got a 400 and lost
|
||||
the edit it actually wanted.
|
||||
"""
|
||||
if key == ENABLE_PTU_COST_ATTRIBUTION_UI_SETTING:
|
||||
return is_ptu_cost_attribution_enabled()
|
||||
return None
|
||||
|
||||
|
||||
# Flags that must be synced from the persisted UISettings into
|
||||
# general_settings at runtime (on both read and write).
|
||||
_RUNTIME_GENERAL_SETTINGS_FLAGS: Final = [
|
||||
|
|
@ -1345,21 +1367,15 @@ async def get_ui_settings():
|
|||
detail={"error": "Database not connected. Please connect a database."},
|
||||
)
|
||||
|
||||
ui_settings: Mapping[str, JsonValue] = {}
|
||||
|
||||
db_record: Final = await _ui_settings_db(UISettingsRepository(prisma_client)).find_unique(
|
||||
where={"id": "ui_settings"}
|
||||
)
|
||||
|
||||
if db_record and db_record.ui_settings:
|
||||
ui_settings_json: Final = db_record.ui_settings
|
||||
if isinstance(ui_settings_json, str):
|
||||
ui_settings = json.loads(ui_settings_json)
|
||||
else:
|
||||
ui_settings = dict(ui_settings_json)
|
||||
stored: Final = (db_record.ui_settings if db_record else None) or "{}"
|
||||
parsed: Final = json.loads(stored) if isinstance(stored, str) else stored
|
||||
|
||||
# Sanitize any unexpected keys from persisted config before returning
|
||||
ui_settings = {k: v for k, v in ui_settings.items() if k in ALLOWED_UI_SETTINGS_FIELDS}
|
||||
ui_settings: Final = {k: v for k, v in parsed.items() if k in ALLOWED_UI_SETTINGS_FIELDS}
|
||||
|
||||
# Sync runtime flags into general_settings so the proxy picks them up
|
||||
# at runtime (covers server restart scenarios).
|
||||
|
|
@ -1377,11 +1393,18 @@ async def get_ui_settings():
|
|||
# Build config-like object for schema helper
|
||||
config: Final[dict[str, object]] = {"litellm_settings": {"ui_settings": ui_settings}}
|
||||
|
||||
return await _get_settings_with_schema(
|
||||
settings: Final = await _get_settings_with_schema(
|
||||
settings_key="ui_settings",
|
||||
settings_class=_get_effective_ui_settings_class(),
|
||||
config=config,
|
||||
)
|
||||
return UISettingsResponse(
|
||||
values={
|
||||
**settings["values"],
|
||||
ENABLE_PTU_COST_ATTRIBUTION_UI_SETTING: is_ptu_cost_attribution_enabled(),
|
||||
},
|
||||
field_schema=settings["field_schema"],
|
||||
)
|
||||
|
||||
|
||||
@router.patch(
|
||||
|
|
@ -1418,6 +1441,20 @@ async def update_ui_settings(
|
|||
detail={"error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature."},
|
||||
)
|
||||
|
||||
conflicting_keys: Final = sorted(
|
||||
key
|
||||
for key, value in settings_body.items()
|
||||
if key in _DERIVED_UI_SETTINGS_FIELDS and value != _derived_ui_setting_value(key)
|
||||
)
|
||||
if conflicting_keys:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=(
|
||||
f"Setting(s) {conflicting_keys} are derived from the deployment environment "
|
||||
"and cannot be changed from the UI."
|
||||
),
|
||||
)
|
||||
|
||||
# Validate against the same effective class GET advertises, so
|
||||
# enterprise-registered fields are typed consistently on both sides.
|
||||
effective_cls: Final = _get_effective_ui_settings_class()
|
||||
|
|
|
|||
|
|
@ -7,6 +7,8 @@ from unittest.mock import AsyncMock, MagicMock
|
|||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../../..")) # Adds the parent directory to the system path
|
||||
|
||||
from litellm.proxy.management_endpoints.common_daily_activity import (
|
||||
|
|
@ -1142,6 +1144,11 @@ class TestEverySavingsDriverSurvivesTheReadPath:
|
|||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ptu_cost_attribution_enabled(monkeypatch):
|
||||
monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true")
|
||||
|
||||
|
||||
def _spend_record(api_key, *, model="gpt-4o-mini-ptu", spend=0.0, ptu_flat_cost=0.0):
|
||||
return SimpleNamespace(
|
||||
api_key=api_key,
|
||||
|
|
@ -1167,13 +1174,13 @@ def _spend_record(api_key, *, model="gpt-4o-mini-ptu", spend=0.0, ptu_flat_cost=
|
|||
)
|
||||
|
||||
|
||||
def test_update_metrics_accumulates_ptu_flat_cost():
|
||||
def test_update_metrics_accumulates_ptu_flat_cost(ptu_cost_attribution_enabled):
|
||||
metrics = update_metrics(SpendMetrics(), _spend_record("real-key", spend=1.0, ptu_flat_cost=240.0))
|
||||
assert metrics.flat_cost == 240.0
|
||||
assert metrics.spend == 1.0
|
||||
|
||||
|
||||
def test_ptu_sentinel_excluded_from_key_breakdown_but_flat_cost_aggregates():
|
||||
def test_ptu_sentinel_excluded_from_key_breakdown_but_flat_cost_aggregates(ptu_cost_attribution_enabled):
|
||||
from litellm.constants import PTU_SENTINEL_API_KEY
|
||||
from litellm.proxy.management_endpoints.common_daily_activity import update_breakdown_metrics
|
||||
from litellm.types.proxy.management_endpoints.common_daily_activity import BreakdownMetrics
|
||||
|
|
@ -1230,7 +1237,7 @@ def _grouping_row(
|
|||
)
|
||||
|
||||
|
||||
def test_grouping_sets_dispatcher_excludes_ptu_sentinel_from_key_breakdowns():
|
||||
def test_grouping_sets_dispatcher_excludes_ptu_sentinel_from_key_breakdowns(ptu_cost_attribution_enabled):
|
||||
"""The GROUPING SETS path must mirror the per-row path: the flat-cost sentinel
|
||||
aggregates into the date/model/total metrics but never surfaces as an api_key."""
|
||||
from litellm.constants import PTU_SENTINEL_API_KEY
|
||||
|
|
@ -1267,7 +1274,7 @@ def test_grouping_sets_dispatcher_excludes_ptu_sentinel_from_key_breakdowns():
|
|||
assert "real-key" in model_bucket.api_key_breakdown
|
||||
|
||||
|
||||
def test_grouping_sets_dispatcher_populates_every_breakdown_level():
|
||||
def test_grouping_sets_dispatcher_populates_every_breakdown_level(ptu_cost_attribution_enabled):
|
||||
"""Every GROUPING SETS level lands in its bucket, and the flat-cost sentinel
|
||||
is kept out of the model_group and provider api_key sub-breakdowns too."""
|
||||
from litellm.constants import PTU_SENTINEL_API_KEY
|
||||
|
|
@ -1359,7 +1366,7 @@ def test_grouping_sets_dispatcher_keeps_a_real_provider_row_that_shares_the_sent
|
|||
assert unknown.metrics.flat_cost == 0.0
|
||||
|
||||
|
||||
def test_update_breakdown_metrics_covers_mcp_endpoint_and_entity():
|
||||
def test_update_breakdown_metrics_covers_mcp_endpoint_and_entity(ptu_cost_attribution_enabled):
|
||||
"""A full request record fans out into the mcp, endpoint, provider and entity
|
||||
breakdowns, while the flat-cost sentinel stays out of the entity api_key sub-map."""
|
||||
from litellm.constants import PTU_SENTINEL_API_KEY
|
||||
|
|
@ -1432,6 +1439,10 @@ class TestSentinelRowsDisplayTheirModelName:
|
|||
"""A sentinel row keys on the deployment id so a rename cannot move it. The usage views
|
||||
render the breakdown key directly as a label, so the read path has to show the name."""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _enabled(self, ptu_cost_attribution_enabled):
|
||||
"""Flat cost is gated off by default, and these assert on the amounts."""
|
||||
|
||||
@staticmethod
|
||||
def _breakdown(records):
|
||||
from litellm.proxy.management_endpoints.common_daily_activity import update_breakdown_metrics
|
||||
|
|
@ -1485,3 +1496,226 @@ class TestSentinelRowsDisplayTheirModelName:
|
|||
models = self._breakdown([self._sentinel(model_id="dep-1", model_group=None)]).models
|
||||
|
||||
assert models["dep-1"].metrics.flat_cost == pytest.approx(480.0)
|
||||
|
||||
|
||||
def _daily_team_row(api_key, *, spend=0.0, ptu_flat_cost=0.0):
|
||||
"""A LiteLLM_DailyTeamSpend row as the paginated read path receives it from find_many."""
|
||||
base: Final = _spend_record(api_key, spend=spend, ptu_flat_cost=ptu_flat_cost)
|
||||
return SimpleNamespace(**{**base.__dict__, "date": "2026-07-01", "team_id": "team-1"})
|
||||
|
||||
|
||||
class TestPtuCostAttributionDisabled:
|
||||
"""With LITELLM_ENABLE_PTU_COST_ATTRIBUTION unset, both read paths report zero flat
|
||||
cost, while the sentinel filtering that keeps ``__ptu_flat_cost__`` out of the
|
||||
breakdowns keeps running.
|
||||
|
||||
Filtering is deliberately not gated: an operator can enable the flag, accrue
|
||||
sentinel rows, then disable it, and those rows stay in LiteLLM_DailyTeamSpend
|
||||
forever. Gating the filter too would surface the sentinel as a bogus api_key and
|
||||
mint a provider bucket for its empty provider.
|
||||
"""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _flag_off(self, monkeypatch):
|
||||
monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False)
|
||||
|
||||
def test_paginated_path_reports_zero_flat_cost(self):
|
||||
metrics = update_metrics(SpendMetrics(), _spend_record("real-key", spend=1.0, ptu_flat_cost=240.0))
|
||||
|
||||
assert metrics.flat_cost == 0.0
|
||||
assert metrics.spend == 1.0
|
||||
|
||||
def test_aggregated_path_reports_zero_flat_cost(self):
|
||||
from litellm.proxy.management_endpoints.common_daily_activity import _GROUP_GRAND_TOTAL
|
||||
|
||||
metrics = _record_to_spend_metrics(_grouping_row(_GROUP_GRAND_TOTAL, spend=5.0, ptu_flat_cost=240.0))
|
||||
|
||||
assert metrics.flat_cost == 0.0
|
||||
assert metrics.spend == 5.0
|
||||
|
||||
def test_aggregated_totals_and_buckets_report_zero_flat_cost(self):
|
||||
from litellm.constants import PTU_SENTINEL_API_KEY
|
||||
from litellm.proxy.management_endpoints.common_daily_activity import (
|
||||
_GROUP_DATE_API_KEY,
|
||||
_GROUP_DATE_MODEL,
|
||||
_GROUP_GRAND_TOTAL,
|
||||
_aggregate_grouping_sets_records_sync,
|
||||
)
|
||||
|
||||
records = [
|
||||
_grouping_row(_GROUP_DATE_API_KEY, api_key=PTU_SENTINEL_API_KEY, ptu_flat_cost=240.0),
|
||||
_grouping_row(_GROUP_DATE_MODEL, model="gpt-4o-mini-ptu", spend=5.0, ptu_flat_cost=240.0),
|
||||
_grouping_row(_GROUP_GRAND_TOTAL, spend=5.0, ptu_flat_cost=240.0),
|
||||
]
|
||||
|
||||
aggregated = _aggregate_grouping_sets_records_sync(records=records, api_key_metadata={})
|
||||
|
||||
assert aggregated["totals"].flat_cost == 0.0
|
||||
assert aggregated["totals"].spend == 5.0
|
||||
assert aggregated["results"][0].breakdown.models["gpt-4o-mini-ptu"].metrics.flat_cost == 0.0
|
||||
|
||||
def test_sentinel_still_excluded_from_the_api_key_breakdown(self):
|
||||
from litellm.constants import PTU_SENTINEL_API_KEY
|
||||
from litellm.proxy.management_endpoints.common_daily_activity import update_breakdown_metrics
|
||||
from litellm.types.proxy.management_endpoints.common_daily_activity import BreakdownMetrics
|
||||
|
||||
breakdown = BreakdownMetrics()
|
||||
update_breakdown_metrics(breakdown, _spend_record("real-key", spend=5.0), {}, {}, {})
|
||||
update_breakdown_metrics(
|
||||
breakdown, _spend_record(PTU_SENTINEL_API_KEY, ptu_flat_cost=240.0), {}, {}, {}, entity_id_field="team_id"
|
||||
)
|
||||
|
||||
assert PTU_SENTINEL_API_KEY not in breakdown.api_keys
|
||||
assert PTU_SENTINEL_API_KEY not in breakdown.models["gpt-4o-mini-ptu"].api_key_breakdown
|
||||
assert "real-key" in breakdown.models["gpt-4o-mini-ptu"].api_key_breakdown
|
||||
|
||||
def test_sentinel_still_excluded_from_the_provider_breakdown(self):
|
||||
from litellm.constants import PTU_SENTINEL_API_KEY
|
||||
from litellm.proxy.management_endpoints.common_daily_activity import update_breakdown_metrics
|
||||
from litellm.types.proxy.management_endpoints.common_daily_activity import BreakdownMetrics
|
||||
|
||||
breakdown = BreakdownMetrics()
|
||||
update_breakdown_metrics(breakdown, _spend_record(PTU_SENTINEL_API_KEY, ptu_flat_cost=240.0), {}, {}, {})
|
||||
|
||||
assert breakdown.providers == {}
|
||||
|
||||
def test_grouping_sets_sentinel_still_excluded_from_breakdowns(self):
|
||||
from litellm.constants import PTU_SENTINEL_API_KEY
|
||||
from litellm.proxy.management_endpoints.common_daily_activity import (
|
||||
_GROUP_DATE_API_KEY,
|
||||
_GROUP_DATE_MODEL,
|
||||
_GROUP_DATE_MODEL_API_KEY,
|
||||
_GROUP_DATE_PROVIDER,
|
||||
_aggregate_grouping_sets_records_sync,
|
||||
)
|
||||
|
||||
records = [
|
||||
_grouping_row(_GROUP_DATE_API_KEY, api_key=PTU_SENTINEL_API_KEY, ptu_flat_cost=240.0),
|
||||
_grouping_row(_GROUP_DATE_MODEL, model="gpt-4o-mini-ptu", spend=5.0, ptu_flat_cost=240.0),
|
||||
_grouping_row(
|
||||
_GROUP_DATE_MODEL_API_KEY, model="gpt-4o-mini-ptu", api_key=PTU_SENTINEL_API_KEY, ptu_flat_cost=240.0
|
||||
),
|
||||
_grouping_row(_GROUP_DATE_PROVIDER, custom_llm_provider="", ptu_flat_cost=240.0),
|
||||
]
|
||||
|
||||
day = _aggregate_grouping_sets_records_sync(records=records, api_key_metadata={})["results"][0]
|
||||
|
||||
assert PTU_SENTINEL_API_KEY not in day.breakdown.api_keys
|
||||
assert PTU_SENTINEL_API_KEY not in day.breakdown.models["gpt-4o-mini-ptu"].api_key_breakdown
|
||||
assert sum(bucket.metrics.flat_cost for bucket in day.breakdown.providers.values()) == 0.0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_daily_activity_endpoint_reports_zero_flat_cost(self):
|
||||
"""/team/daily/activity reads rows with find_many rather than the aggregated SQL, so
|
||||
forcing the SQL select to a constant zero would leave this path reporting flat cost."""
|
||||
from litellm.constants import PTU_SENTINEL_API_KEY
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db = MagicMock()
|
||||
mock_table = MagicMock()
|
||||
mock_table.count = AsyncMock(return_value=2)
|
||||
mock_table.find_many = AsyncMock(
|
||||
return_value=[
|
||||
_daily_team_row("real-key", spend=5.0),
|
||||
_daily_team_row(PTU_SENTINEL_API_KEY, ptu_flat_cost=240.0),
|
||||
]
|
||||
)
|
||||
mock_prisma.db.litellm_verificationtoken = MagicMock()
|
||||
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
|
||||
mock_prisma.db.litellm_dailyteamspend = mock_table
|
||||
|
||||
result = await get_daily_activity(
|
||||
prisma_client=mock_prisma,
|
||||
table_name="litellm_dailyteamspend",
|
||||
entity_id_field="team_id",
|
||||
entity_id="team-1",
|
||||
entity_metadata_field=None,
|
||||
start_date="2026-07-01",
|
||||
end_date="2026-07-01",
|
||||
model=None,
|
||||
api_key=None,
|
||||
page=1,
|
||||
page_size=50,
|
||||
)
|
||||
|
||||
assert result.metadata.total_flat_cost == 0.0
|
||||
assert result.metadata.total_spend == 5.0
|
||||
assert PTU_SENTINEL_API_KEY not in result.results[0].breakdown.api_keys
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_daily_activity_endpoint_reports_flat_cost_once_enabled(self, monkeypatch):
|
||||
from litellm.constants import PTU_SENTINEL_API_KEY
|
||||
|
||||
monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true")
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db = MagicMock()
|
||||
mock_table = MagicMock()
|
||||
mock_table.count = AsyncMock(return_value=2)
|
||||
mock_table.find_many = AsyncMock(
|
||||
return_value=[
|
||||
_daily_team_row("real-key", spend=5.0),
|
||||
_daily_team_row(PTU_SENTINEL_API_KEY, ptu_flat_cost=240.0),
|
||||
]
|
||||
)
|
||||
mock_prisma.db.litellm_verificationtoken = MagicMock()
|
||||
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
|
||||
mock_prisma.db.litellm_dailyteamspend = mock_table
|
||||
|
||||
result = await get_daily_activity(
|
||||
prisma_client=mock_prisma,
|
||||
table_name="litellm_dailyteamspend",
|
||||
entity_id_field="team_id",
|
||||
entity_id="team-1",
|
||||
entity_metadata_field=None,
|
||||
start_date="2026-07-01",
|
||||
end_date="2026-07-01",
|
||||
model=None,
|
||||
api_key=None,
|
||||
page=1,
|
||||
page_size=50,
|
||||
)
|
||||
|
||||
assert result.metadata.total_flat_cost == 240.0
|
||||
assert PTU_SENTINEL_API_KEY not in result.results[0].breakdown.api_keys
|
||||
|
||||
|
||||
class TestFlagIsNotReadOnTheHotPath:
|
||||
"""update_metrics runs once per accumulation and a record fans out across roughly a
|
||||
dozen breakdowns, so a flag that reads through the secret manager must not be consulted
|
||||
for rows that carry no flat cost at all."""
|
||||
|
||||
@staticmethod
|
||||
def _count_flag_reads(records):
|
||||
import litellm.proxy.management_endpoints.common_daily_activity as cda
|
||||
from litellm.types.proxy.management_endpoints.common_daily_activity import BreakdownMetrics
|
||||
|
||||
reads = []
|
||||
real = cda.is_ptu_cost_attribution_enabled
|
||||
|
||||
def counted():
|
||||
reads.append(1)
|
||||
return real()
|
||||
|
||||
cda.is_ptu_cost_attribution_enabled = counted
|
||||
try:
|
||||
breakdown = BreakdownMetrics()
|
||||
for record in records:
|
||||
cda.update_breakdown_metrics(breakdown, record, {}, {}, {})
|
||||
finally:
|
||||
cda.is_ptu_cost_attribution_enabled = real
|
||||
return len(reads)
|
||||
|
||||
def test_a_request_row_never_reads_the_flag(self):
|
||||
reads = self._count_flag_reads([_spend_record("real-key", spend=5.0, ptu_flat_cost=0.0)])
|
||||
assert reads == 0, f"{reads} secret-manager lookups for a row with no flat cost"
|
||||
|
||||
def test_a_page_of_request_rows_never_reads_the_flag(self):
|
||||
rows = [_spend_record(f"key-{i}", spend=1.0, ptu_flat_cost=0.0) for i in range(50)]
|
||||
assert self._count_flag_reads(rows) == 0
|
||||
|
||||
def test_a_sentinel_row_still_consults_the_flag(self):
|
||||
from litellm.constants import PTU_SENTINEL_API_KEY
|
||||
|
||||
reads = self._count_flag_reads([_spend_record(PTU_SENTINEL_API_KEY, spend=0.0, ptu_flat_cost=240.0)])
|
||||
assert reads > 0
|
||||
|
|
|
|||
|
|
@ -1,34 +1,57 @@
|
|||
import datetime
|
||||
import json
|
||||
|
||||
"""Tests for PTU config on the model deployment (v1 model-settings design)."""
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
import datetime
|
||||
import json
|
||||
from contextlib import ExitStack
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy._types import LiteLLM_ProxyModelTable, LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
_merged_ptu_model_info,
|
||||
_raise_if_ptu_cost_attribution_disabled,
|
||||
_validate_ptu_model_info,
|
||||
add_new_model,
|
||||
update_db_model,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR
|
||||
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo, updateDeployment
|
||||
|
||||
|
||||
def test_model_info_accepts_valid_ptu_fields():
|
||||
info = ModelInfo(id="x", team_id="t", ptu_count=5, cost_per_ptu_per_hour=2.0)
|
||||
info = ModelInfo(
|
||||
id="x",
|
||||
team_id="t",
|
||||
ptu_count=5,
|
||||
cost_per_ptu_per_hour=2.0,
|
||||
ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc),
|
||||
)
|
||||
assert info.ptu_count == 5
|
||||
assert info.cost_per_ptu_per_hour == 2.0
|
||||
|
||||
|
||||
def test_model_info_rejects_non_positive_count():
|
||||
with pytest.raises(ValueError):
|
||||
ModelInfo(id="x", team_id="t", ptu_count=0, cost_per_ptu_per_hour=2.0)
|
||||
ModelInfo(
|
||||
id="x",
|
||||
team_id="t",
|
||||
ptu_count=0,
|
||||
cost_per_ptu_per_hour=2.0,
|
||||
ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc),
|
||||
)
|
||||
|
||||
|
||||
def test_model_info_rejects_negative_rate():
|
||||
with pytest.raises(ValueError):
|
||||
ModelInfo(id="x", team_id="t", ptu_count=5, cost_per_ptu_per_hour=-1.0)
|
||||
ModelInfo(
|
||||
id="x",
|
||||
team_id="t",
|
||||
ptu_count=5,
|
||||
cost_per_ptu_per_hour=-1.0,
|
||||
ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc),
|
||||
)
|
||||
|
||||
|
||||
def test_model_info_rejects_a_count_beyond_the_cap():
|
||||
|
|
@ -223,6 +246,11 @@ class TestPartialPtuEditsUseTheMergedView:
|
|||
"""A PTU invariant holds over the deployment as it will exist, not over whichever
|
||||
subset of fields a caller sent. Validating the patch alone rejected an ordinary edit."""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _enabled(self, monkeypatch):
|
||||
"""PTU writes are gated off by default; these are about the validator, not the gate."""
|
||||
monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true")
|
||||
|
||||
@staticmethod
|
||||
def _configured():
|
||||
return Deployment(
|
||||
|
|
@ -310,6 +338,11 @@ class TestTeamModelUpdateValidatesBeforeWriting:
|
|||
"""Drives the endpoint path itself, not the helpers. The validator sits above the team
|
||||
ACL write, which autocommits, so what it validates has to be right at that call site."""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _enabled(self, monkeypatch):
|
||||
"""PTU writes are gated off by default; these are about the validator, not the gate."""
|
||||
monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true")
|
||||
|
||||
@staticmethod
|
||||
async def _run(db_model, patch_data, monkeypatch, touched=None):
|
||||
import litellm.proxy.management_endpoints.model_management_endpoints as mme
|
||||
|
|
@ -352,6 +385,33 @@ class TestTeamModelUpdateValidatesBeforeWriting:
|
|||
|
||||
assert "ptu_effective_from is required" in exc.value.detail
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_gate_refuses_before_the_team_write(self, monkeypatch):
|
||||
"""The gate lived inside update_db_model, which runs after the team ACL write, so a
|
||||
rejected edit still moved the model between teams."""
|
||||
monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False)
|
||||
db_model = Deployment(
|
||||
model_name="gpt-4o",
|
||||
litellm_params=LiteLLM_Params(model="openai/gpt-4o"),
|
||||
model_info=ModelInfo(id="dep-0", team_id="team-A"),
|
||||
)
|
||||
patch = updateDeployment(
|
||||
model_info=ModelInfo(
|
||||
id="dep-0",
|
||||
team_id="team-B",
|
||||
ptu_count=15,
|
||||
cost_per_ptu_per_hour=2.0,
|
||||
ptu_effective_from=datetime.datetime(2026, 8, 1, tzinfo=datetime.timezone.utc),
|
||||
)
|
||||
)
|
||||
touched = []
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await self._run(db_model, patch, monkeypatch, touched)
|
||||
|
||||
assert PTU_COST_ATTRIBUTION_ENV_VAR in exc.value.detail
|
||||
assert touched == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_clearing_half_the_pair_is_refused_before_the_team_write(self, monkeypatch):
|
||||
"""The write drops the nulled field, so validating against the stored one let a
|
||||
|
|
@ -380,3 +440,270 @@ class TestTeamModelUpdateValidatesBeforeWriting:
|
|||
stored = json.loads(result["model_info"])
|
||||
assert "ptu_count" not in stored
|
||||
assert "cost_per_ptu_per_hour" not in stored
|
||||
|
||||
|
||||
class TestPtuCostAttributionGate:
|
||||
"""PTU config is only writable once an operator sets LITELLM_ENABLE_PTU_COST_ATTRIBUTION.
|
||||
|
||||
The fields are rejected rather than dropped: a silent accept-and-drop would let a
|
||||
caller believe a flat cost was configured while the rollup that prices it is not
|
||||
even scheduled.
|
||||
"""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _flag_off(self, monkeypatch):
|
||||
monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False)
|
||||
|
||||
@pytest.fixture
|
||||
def flag_on(self, monkeypatch):
|
||||
monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true")
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_info",
|
||||
[
|
||||
{"team_id": "t", "ptu_count": 5, "cost_per_ptu_per_hour": 2.0},
|
||||
{"ptu_count": 5},
|
||||
{"cost_per_ptu_per_hour": 2.0},
|
||||
{"ptu_effective_from": "2026-08-01T00:00:00Z"},
|
||||
{"ptu_effective_to": "2026-08-02T00:00:00Z"},
|
||||
],
|
||||
)
|
||||
def test_rejects_any_ptu_field_while_disabled(self, model_info):
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
_raise_if_ptu_cost_attribution_disabled(model_info)
|
||||
assert exc.value.status_code == 400
|
||||
assert PTU_COST_ATTRIBUTION_ENV_VAR in exc.value.detail
|
||||
|
||||
def test_names_every_offending_field(self):
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
_raise_if_ptu_cost_attribution_disabled({"ptu_count": 5, "cost_per_ptu_per_hour": 2.0})
|
||||
assert "ptu_count" in exc.value.detail
|
||||
assert "cost_per_ptu_per_hour" in exc.value.detail
|
||||
|
||||
def test_allows_a_request_without_ptu_fields_while_disabled(self):
|
||||
_raise_if_ptu_cost_attribution_disabled({"team_id": "t", "access_groups": ["a"]})
|
||||
|
||||
def test_allows_every_ptu_field_once_enabled(self, flag_on):
|
||||
_raise_if_ptu_cost_attribution_disabled(
|
||||
{
|
||||
"team_id": "t",
|
||||
"ptu_count": 5,
|
||||
"cost_per_ptu_per_hour": 2.0,
|
||||
"ptu_effective_from": "2026-08-01T00:00:00Z",
|
||||
"ptu_effective_to": "2026-08-02T00:00:00Z",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _deployment_without_ptu() -> Deployment:
|
||||
return Deployment(
|
||||
model_name="gpt-4o",
|
||||
litellm_params=LiteLLM_Params(model="openai/gpt-4o"),
|
||||
model_info=ModelInfo(id="dep-0", team_id="t"),
|
||||
)
|
||||
|
||||
|
||||
def _deployment_with_stored_ptu() -> Deployment:
|
||||
return Deployment(
|
||||
model_name="gpt-4o",
|
||||
litellm_params=LiteLLM_Params(model="openai/gpt-4o"),
|
||||
model_info=ModelInfo(
|
||||
id="dep-0",
|
||||
team_id="t",
|
||||
ptu_count=15,
|
||||
cost_per_ptu_per_hour=2.0,
|
||||
ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class TestUpdateDbModelPtuGate:
|
||||
@pytest.fixture(autouse=True)
|
||||
def _flag_off(self, monkeypatch):
|
||||
monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False)
|
||||
|
||||
def test_patch_carrying_ptu_config_is_rejected(self):
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
update_db_model(
|
||||
db_model=_deployment_without_ptu(),
|
||||
updated_patch=updateDeployment(model_info=ModelInfo(id="dep-0", team_id="t", ptu_count=15)),
|
||||
)
|
||||
assert exc.value.status_code == 400
|
||||
|
||||
def test_patch_that_touches_nothing_ptu_still_succeeds(self):
|
||||
result = update_db_model(
|
||||
db_model=_deployment_without_ptu(),
|
||||
updated_patch=updateDeployment(model_info=ModelInfo(id="dep-0", access_groups=["a"])),
|
||||
)
|
||||
assert json.loads(result["model_info"])["access_groups"] == ["a"]
|
||||
|
||||
def test_unrelated_patch_of_a_model_that_stores_ptu_config_is_not_blocked(self):
|
||||
"""A deployment configured during an earlier opt-in stays editable: the gate reads the
|
||||
incoming patch, not the merged deployment, so the stored config is left in place."""
|
||||
result = update_db_model(
|
||||
db_model=_deployment_with_stored_ptu(),
|
||||
updated_patch=updateDeployment(model_name="gpt-4o-renamed"),
|
||||
)
|
||||
assert result["model_name"] == "gpt-4o-renamed"
|
||||
|
||||
def test_explicit_nulls_do_not_erase_stored_ptu_config_while_disabled(self):
|
||||
"""A client round-tripping a model_info blob sends the PTU keys as nulls. While the
|
||||
feature is disabled those nulls must not reach the clear loop: disabling pauses PTU,
|
||||
it does not silently discard a billing configuration the operator set up earlier."""
|
||||
result = update_db_model(
|
||||
db_model=_deployment_with_stored_ptu(),
|
||||
updated_patch=updateDeployment(
|
||||
model_info=ModelInfo(id="dep-0", ptu_count=None, cost_per_ptu_per_hour=None)
|
||||
),
|
||||
)
|
||||
stored = json.loads(result["model_info"])
|
||||
assert stored["ptu_count"] == 15
|
||||
assert stored["cost_per_ptu_per_hour"] == 2.0
|
||||
|
||||
def test_the_merged_view_agrees_with_the_write_while_disabled(self):
|
||||
"""The validator sees what the write will store. If the merged view honoured a null the
|
||||
clear loop ignores, a round-tripped blob would 400 on a half-set pair that never forms."""
|
||||
merged = _merged_ptu_model_info(
|
||||
db_model=_deployment_with_stored_ptu(),
|
||||
patch_data=updateDeployment(model_info=ModelInfo(id="dep-0", ptu_count=None)),
|
||||
)
|
||||
assert merged["ptu_count"] == 15
|
||||
_validate_ptu_model_info(merged)
|
||||
|
||||
def test_explicit_nulls_still_clear_once_enabled(self, monkeypatch):
|
||||
"""Clearing remains available to an operator who opted in, which is how PTU config is
|
||||
removed from a deployment."""
|
||||
monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true")
|
||||
result = update_db_model(
|
||||
db_model=_deployment_with_stored_ptu(),
|
||||
updated_patch=updateDeployment(
|
||||
model_info=ModelInfo(id="dep-0", ptu_count=None, cost_per_ptu_per_hour=None)
|
||||
),
|
||||
)
|
||||
stored = json.loads(result["model_info"])
|
||||
assert "ptu_count" not in stored
|
||||
assert "cost_per_ptu_per_hour" not in stored
|
||||
|
||||
def test_patch_carrying_ptu_config_is_accepted_once_enabled(self, monkeypatch):
|
||||
monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true")
|
||||
result = update_db_model(
|
||||
db_model=_deployment_without_ptu(),
|
||||
updated_patch=updateDeployment(
|
||||
model_info=ModelInfo(
|
||||
id="dep-0",
|
||||
team_id="t",
|
||||
ptu_count=15,
|
||||
cost_per_ptu_per_hour=2.0,
|
||||
ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc),
|
||||
)
|
||||
),
|
||||
)
|
||||
stored = json.loads(result["model_info"])
|
||||
assert stored["ptu_count"] == 15
|
||||
assert stored["cost_per_ptu_per_hour"] == 2.0
|
||||
|
||||
|
||||
class TestAddNewModelPtuGate:
|
||||
@pytest.fixture(autouse=True)
|
||||
def _flag_off(self, monkeypatch):
|
||||
monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False)
|
||||
|
||||
@staticmethod
|
||||
def _patched_proxy(model_id: str):
|
||||
"""Patch everything /model/new touches except the PTU gate, and hand back the DB writers."""
|
||||
db_row = LiteLLM_ProxyModelTable(
|
||||
model_id=model_id,
|
||||
model_name="ptu-model",
|
||||
litellm_params={"model": "openai/gpt-4.1-nano"},
|
||||
model_info={"id": model_id},
|
||||
created_by="test-admin",
|
||||
updated_by="test-admin",
|
||||
)
|
||||
add_model_to_db = AsyncMock(return_value=db_row)
|
||||
add_team_model_to_db = AsyncMock(return_value=db_row)
|
||||
|
||||
mock_proxy_config = MagicMock()
|
||||
mock_proxy_config.add_deployment = AsyncMock(return_value=None)
|
||||
|
||||
mock_router = MagicMock()
|
||||
mock_router.get_model_ids.return_value = [model_id]
|
||||
|
||||
proxy_server = "litellm.proxy.proxy_server"
|
||||
endpoints = "litellm.proxy.management_endpoints.model_management_endpoints"
|
||||
return (add_model_to_db, add_team_model_to_db), [
|
||||
patch(f"{proxy_server}.prisma_client", MagicMock()),
|
||||
patch(f"{proxy_server}.store_model_in_db", True),
|
||||
patch(f"{proxy_server}.proxy_config", mock_proxy_config),
|
||||
patch(f"{proxy_server}.proxy_logging_obj", MagicMock()),
|
||||
patch(f"{proxy_server}.general_settings", {}),
|
||||
patch(f"{proxy_server}.premium_user", True),
|
||||
patch(f"{proxy_server}.llm_router", mock_router),
|
||||
patch(
|
||||
f"{endpoints}.ModelManagementAuthChecks.can_user_make_model_call",
|
||||
AsyncMock(return_value=True),
|
||||
),
|
||||
patch(f"{endpoints}._add_model_to_db", add_model_to_db),
|
||||
patch(f"{endpoints}._add_team_model_to_db", add_team_model_to_db),
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def _ptu_deployment(model_id: str) -> Deployment:
|
||||
return Deployment(
|
||||
model_name="ptu-model",
|
||||
litellm_params=LiteLLM_Params(model="openai/gpt-4.1-nano", api_key="fake-key"),
|
||||
model_info=ModelInfo(
|
||||
id=model_id,
|
||||
team_id="team-1",
|
||||
ptu_count=15,
|
||||
cost_per_ptu_per_hour=2.0,
|
||||
ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc),
|
||||
),
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_new_rejects_ptu_config_while_disabled(self):
|
||||
(add_model_to_db, add_team_model_to_db), patches = self._patched_proxy("ptu-gate-model")
|
||||
admin = UserAPIKeyAuth(user_id="test-admin", user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
|
||||
with ExitStack() as stack:
|
||||
for active_patch in patches:
|
||||
stack.enter_context(active_patch)
|
||||
with pytest.raises(Exception) as exc:
|
||||
await add_new_model(model_params=self._ptu_deployment("ptu-gate-model"), user_api_key_dict=admin)
|
||||
|
||||
assert PTU_COST_ATTRIBUTION_ENV_VAR in str(exc.value)
|
||||
add_model_to_db.assert_not_called()
|
||||
add_team_model_to_db.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_new_accepts_a_deployment_without_ptu_config_while_disabled(self):
|
||||
_, patches = self._patched_proxy("plain-model")
|
||||
admin = UserAPIKeyAuth(user_id="test-admin", user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
|
||||
with ExitStack() as stack:
|
||||
for active_patch in patches:
|
||||
stack.enter_context(active_patch)
|
||||
result = await add_new_model(
|
||||
model_params=Deployment(
|
||||
model_name="ptu-model",
|
||||
litellm_params=LiteLLM_Params(model="openai/gpt-4.1-nano", api_key="fake-key"),
|
||||
model_info=ModelInfo(id="plain-model"),
|
||||
),
|
||||
user_api_key_dict=admin,
|
||||
)
|
||||
|
||||
assert result.model_id == "plain-model"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_new_accepts_ptu_config_once_enabled(self, monkeypatch):
|
||||
monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true")
|
||||
(_, add_team_model_to_db), patches = self._patched_proxy("ptu-gate-model")
|
||||
admin = UserAPIKeyAuth(user_id="test-admin", user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
|
||||
with ExitStack() as stack:
|
||||
for active_patch in patches:
|
||||
stack.enter_context(active_patch)
|
||||
result = await add_new_model(model_params=self._ptu_deployment("ptu-gate-model"), user_api_key_dict=admin)
|
||||
|
||||
assert result.model_id == "ptu-gate-model"
|
||||
add_team_model_to_db.assert_called_once()
|
||||
|
|
|
|||
|
|
@ -0,0 +1,33 @@
|
|||
"""Tests for the opt-in flag that gates PTU flat-cost attribution."""
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.spend_tracking.ptu_feature_flag import (
|
||||
PTU_COST_ATTRIBUTION_ENV_VAR,
|
||||
is_ptu_cost_attribution_enabled,
|
||||
)
|
||||
|
||||
|
||||
def test_disabled_when_env_var_is_unset(monkeypatch):
|
||||
monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False)
|
||||
assert is_ptu_cost_attribution_enabled() is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", ["true", "True", "TRUE", " true "])
|
||||
def test_enabled_for_the_values_the_house_helper_recognises(monkeypatch, value):
|
||||
monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, value)
|
||||
assert is_ptu_cost_attribution_enabled() is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", ["false", "False", "0", "1", "", "yes", "off", "maybe"])
|
||||
def test_disabled_for_everything_else(monkeypatch, value):
|
||||
monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, value)
|
||||
assert is_ptu_cost_attribution_enabled() is False
|
||||
|
||||
|
||||
def test_reads_the_env_var_on_every_call(monkeypatch):
|
||||
monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False)
|
||||
assert is_ptu_cost_attribution_enabled() is False
|
||||
|
||||
monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true")
|
||||
assert is_ptu_cost_attribution_enabled() is True
|
||||
|
|
@ -8,6 +8,7 @@ import pytest
|
|||
|
||||
import litellm.proxy.spend_tracking.ptu_flat_cost_rollup as ptu_rollup
|
||||
from litellm.constants import PTU_ROLLUP_MAX_BACKFILL_DAYS, PTU_SENTINEL_API_KEY
|
||||
from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR
|
||||
from litellm.types.router import ModelInfo
|
||||
from litellm.proxy.spend_tracking.ptu_flat_cost_rollup import (
|
||||
PTUModel,
|
||||
|
|
@ -29,6 +30,13 @@ TODAY = date(2026, 7, 31)
|
|||
_DEFAULT_PTU_START = "2020-01-01T00:00:00Z"
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _ptu_enabled(monkeypatch):
|
||||
"""PTU is gated off by default. These cover the rollup's mechanics, not the gate, so
|
||||
they run with it on; the gate itself is covered by its own test below."""
|
||||
monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true")
|
||||
|
||||
|
||||
_VALID_PTU = {"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "team_id": "t"}
|
||||
|
||||
|
||||
|
|
@ -1479,3 +1487,18 @@ async def test_the_prune_cutoff_allows_for_clock_skew_between_hosts():
|
|||
"a charge written 30s ago by a lagging pod was swept"
|
||||
)
|
||||
assert ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-stale") not in table.rows
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scheduled_rollup_writes_nothing_when_ptu_attribution_is_disabled(monkeypatch):
|
||||
"""Startup already skips scheduling the cron, so this guards the function itself: a
|
||||
deployment that never opted in accrues nothing whatever route reaches the rollup."""
|
||||
monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False)
|
||||
table = _FakeSentinelTable()
|
||||
prisma = _prisma_for([_model_row(model_info=_VALID_PTU)], table)
|
||||
|
||||
result = await run_scheduled_ptu_rollup(prisma, pod_lock_manager=None, alert=None)
|
||||
|
||||
assert result is None
|
||||
assert table.rows == {}
|
||||
assert table.upsert_keys == []
|
||||
|
|
|
|||
|
|
@ -11280,14 +11280,8 @@ async def test_setup_prisma_client_returns_none_when_connect_itself_fails(monkey
|
|||
assert mock_client.health_check.await_count == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ptu_rollup_job_registered_at_startup(monkeypatch):
|
||||
"""The PTU rollup cron is registered at startup; only models with PTU config accrue flat cost (asserted in test_ptu_flat_cost_rollup.py)."""
|
||||
monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False)
|
||||
async def _run_scheduled_background_jobs():
|
||||
from litellm.proxy.proxy_server import ProxyStartupEvent
|
||||
from litellm.proxy.spend_tracking.ptu_flat_cost_rollup import (
|
||||
PTU_ROLLUP_JOB_ID,
|
||||
)
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
||||
mock_prisma_client = MagicMock()
|
||||
|
|
@ -11311,7 +11305,41 @@ async def test_ptu_rollup_job_registered_at_startup(monkeypatch):
|
|||
proxy_logging_obj=mock_proxy_logging,
|
||||
)
|
||||
|
||||
import litellm.proxy.proxy_server as ps
|
||||
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
|
||||
assert ps.scheduler is not None
|
||||
return ps.scheduler
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ptu_rollup_job_registered_at_startup(monkeypatch):
|
||||
"""The PTU rollup cron is registered once an operator opts in; only models with PTU config accrue flat cost (asserted in test_ptu_flat_cost_rollup.py)."""
|
||||
monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False)
|
||||
from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR
|
||||
from litellm.proxy.spend_tracking.ptu_flat_cost_rollup import (
|
||||
PTU_ROLLUP_JOB_ID,
|
||||
)
|
||||
|
||||
monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true")
|
||||
|
||||
scheduler = await _run_scheduled_background_jobs()
|
||||
|
||||
assert scheduler.get_job(PTU_ROLLUP_JOB_ID) is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ptu_rollup_job_not_registered_without_opt_in(monkeypatch):
|
||||
"""Without LITELLM_ENABLE_PTU_COST_ATTRIBUTION the rollup never runs, so no sentinel row
|
||||
is ever written. This is the gate that keeps the whole feature inert by default."""
|
||||
monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False)
|
||||
from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR
|
||||
from litellm.proxy.spend_tracking.ptu_flat_cost_rollup import (
|
||||
PTU_ROLLUP_JOB_ID,
|
||||
)
|
||||
|
||||
monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False)
|
||||
|
||||
scheduler = await _run_scheduled_background_jobs()
|
||||
|
||||
assert scheduler.get_job(PTU_ROLLUP_JOB_ID) is None
|
||||
assert len(scheduler.get_jobs()) > 0
|
||||
|
|
|
|||
|
|
@ -2928,3 +2928,143 @@ def test_update_mcp_semantic_filter_settings_requires_proxy_admin(monkeypatch):
|
|||
assert "proxy admin" in resp.json()["detail"].lower()
|
||||
finally:
|
||||
app.dependency_overrides.pop(user_api_key_auth, None)
|
||||
|
||||
|
||||
class TestPtuCostAttributionUISetting:
|
||||
"""``enable_ptu_cost_attribution`` is derived from the environment on every GET.
|
||||
|
||||
It is deliberately not an allowlisted, persisted setting: the point of gating PTU
|
||||
flat cost on an env var is that an admin cannot flip it at runtime from the UI.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def _mock_prisma(monkeypatch, stored=None):
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_record = None
|
||||
if stored is not None:
|
||||
mock_record = MagicMock()
|
||||
mock_record.ui_settings = stored
|
||||
mock_prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=mock_record)
|
||||
mock_prisma.db.litellm_uisettings.upsert = AsyncMock()
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
|
||||
return mock_prisma
|
||||
|
||||
def test_reported_false_when_the_env_var_is_unset(self, mock_auth, monkeypatch):
|
||||
from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR
|
||||
|
||||
monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False)
|
||||
self._mock_prisma(monkeypatch)
|
||||
|
||||
response = client.get("/get/ui_settings")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["values"]["enable_ptu_cost_attribution"] is False
|
||||
|
||||
def test_reported_true_once_the_env_var_is_set(self, mock_auth, monkeypatch):
|
||||
from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR
|
||||
|
||||
monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true")
|
||||
self._mock_prisma(monkeypatch)
|
||||
|
||||
response = client.get("/get/ui_settings")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["values"]["enable_ptu_cost_attribution"] is True
|
||||
|
||||
def test_a_persisted_true_cannot_forge_the_derived_value(self, mock_auth, monkeypatch):
|
||||
"""A row written before the allowlist existed must not be able to turn the feature on."""
|
||||
from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR
|
||||
|
||||
monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False)
|
||||
self._mock_prisma(monkeypatch, stored={"enable_ptu_cost_attribution": True})
|
||||
|
||||
response = client.get("/get/ui_settings")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["values"]["enable_ptu_cost_attribution"] is False
|
||||
|
||||
def test_is_not_an_allowlisted_persisted_setting(self):
|
||||
from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
|
||||
ALLOWED_UI_SETTINGS_FIELDS,
|
||||
)
|
||||
|
||||
assert "enable_ptu_cost_attribution" not in ALLOWED_UI_SETTINGS_FIELDS
|
||||
|
||||
def test_the_body_get_returns_is_a_valid_patch_body(self, mock_auth, monkeypatch):
|
||||
"""Read-modify-write is how a client edits one setting. GET injects the derived key,
|
||||
so rejecting it on presence made GET's own output an invalid PATCH body: the caller
|
||||
got a 400 and silently lost the edit it actually wanted."""
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR
|
||||
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_id="test-user-123",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True)
|
||||
monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False)
|
||||
mock_prisma = self._mock_prisma(monkeypatch)
|
||||
|
||||
try:
|
||||
round_tripped = client.get("/get/ui_settings").json()["values"]
|
||||
assert "enable_ptu_cost_attribution" in round_tripped
|
||||
response = client.patch("/update/ui_settings", json=round_tripped)
|
||||
finally:
|
||||
app.dependency_overrides.clear()
|
||||
|
||||
assert response.status_code == 200
|
||||
assert mock_prisma.db.litellm_uisettings.upsert.called
|
||||
|
||||
def test_a_co_submitted_setting_still_applies_alongside_the_derived_key(self, mock_auth, monkeypatch):
|
||||
"""The derived key riding along must not discard the caller's real edit."""
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR
|
||||
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_id="test-user-123",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True)
|
||||
monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False)
|
||||
mock_prisma = self._mock_prisma(monkeypatch)
|
||||
|
||||
try:
|
||||
response = client.patch(
|
||||
"/update/ui_settings",
|
||||
json={"enable_ptu_cost_attribution": False, "enable_chat_ui": True},
|
||||
)
|
||||
finally:
|
||||
app.dependency_overrides.clear()
|
||||
|
||||
assert response.status_code == 200
|
||||
upsert_data = mock_prisma.db.litellm_uisettings.upsert.call_args.kwargs["data"]
|
||||
persisted = json.loads(upsert_data["create"]["ui_settings"])
|
||||
assert persisted["enable_chat_ui"] is True
|
||||
assert "enable_ptu_cost_attribution" not in persisted
|
||||
|
||||
def test_patch_rejects_the_derived_setting(self, mock_auth, monkeypatch):
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
user_id="test-user-123",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True)
|
||||
mock_prisma = self._mock_prisma(monkeypatch)
|
||||
|
||||
try:
|
||||
response = client.patch(
|
||||
"/update/ui_settings",
|
||||
json={"enable_ptu_cost_attribution": True},
|
||||
)
|
||||
finally:
|
||||
app.dependency_overrides.clear()
|
||||
|
||||
assert response.status_code == 400
|
||||
assert "enable_ptu_cost_attribution" in str(response.json()["detail"])
|
||||
assert not mock_prisma.db.litellm_uisettings.upsert.called
|
||||
|
|
|
|||
|
|
@ -0,0 +1,135 @@
|
|||
import { getUiSettings } from "@/components/networking";
|
||||
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
|
||||
import { renderHook, waitFor } from "@testing-library/react";
|
||||
import React, { ReactNode } from "react";
|
||||
import { beforeEach, describe, expect, it, vi } from "vitest";
|
||||
import { PTU_FLAG_REFRESH_MS, usePtuCostAttributionEnabled } from "./usePtuCostAttributionEnabled";
|
||||
import { useUISettings } from "./useUISettings";
|
||||
|
||||
vi.mock("@/components/networking", () => ({
|
||||
getUiSettings: vi.fn(),
|
||||
}));
|
||||
|
||||
describe("usePtuCostAttributionEnabled", () => {
|
||||
let queryClient: QueryClient;
|
||||
|
||||
beforeEach(() => {
|
||||
queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } });
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
const wrapper = ({ children }: { children: ReactNode }) =>
|
||||
React.createElement(QueryClientProvider, { client: queryClient }, children);
|
||||
|
||||
/** Read the flag alongside the query it derives from, so assertions wait for a settled fetch. */
|
||||
const renderSettledFlag = async (settings: unknown) => {
|
||||
(getUiSettings as any).mockResolvedValue(settings);
|
||||
const { result } = renderHook(() => ({ enabled: usePtuCostAttributionEnabled(), query: useUISettings() }), {
|
||||
wrapper,
|
||||
});
|
||||
await waitFor(() => {
|
||||
expect(result.current.query.isSuccess).toBe(true);
|
||||
});
|
||||
return result;
|
||||
};
|
||||
|
||||
it("is true only when the proxy reports the flag as enabled", async () => {
|
||||
const result = await renderSettledFlag({ values: { enable_ptu_cost_attribution: true } });
|
||||
expect(result.current.enabled).toBe(true);
|
||||
});
|
||||
|
||||
it("is false when the proxy reports the flag as disabled", async () => {
|
||||
const result = await renderSettledFlag({ values: { enable_ptu_cost_attribution: false } });
|
||||
expect(result.current.enabled).toBe(false);
|
||||
});
|
||||
|
||||
it("is false when the proxy omits the flag entirely", async () => {
|
||||
const result = await renderSettledFlag({ values: { enable_chat_ui: true } });
|
||||
expect(result.current.enabled).toBe(false);
|
||||
});
|
||||
|
||||
it("is false when the proxy returns no values at all", async () => {
|
||||
const result = await renderSettledFlag({});
|
||||
expect(result.current.enabled).toBe(false);
|
||||
});
|
||||
|
||||
it("does not treat a truthy non-boolean as enabled", async () => {
|
||||
const result = await renderSettledFlag({ values: { enable_ptu_cost_attribution: "false" } });
|
||||
expect(result.current.enabled).toBe(false);
|
||||
});
|
||||
|
||||
it("does not treat the string 'true' as enabled, since the proxy sends a real boolean", async () => {
|
||||
const result = await renderSettledFlag({ values: { enable_ptu_cost_attribution: "true" } });
|
||||
expect(result.current.enabled).toBe(false);
|
||||
});
|
||||
|
||||
it("is false before the settings request resolves", () => {
|
||||
(getUiSettings as any).mockReturnValue(new Promise(() => {}));
|
||||
const { result } = renderHook(() => usePtuCostAttributionEnabled(), { wrapper });
|
||||
expect(result.current).toBe(false);
|
||||
});
|
||||
|
||||
it("is false when the settings request fails", async () => {
|
||||
(getUiSettings as any).mockRejectedValue(new Error("boom"));
|
||||
const { result } = renderHook(() => ({ enabled: usePtuCostAttributionEnabled(), query: useUISettings() }), {
|
||||
wrapper,
|
||||
});
|
||||
await waitFor(() => {
|
||||
expect(result.current.query.isError).toBe(true);
|
||||
});
|
||||
expect(result.current.enabled).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe("staleness", () => {
|
||||
let queryClient: QueryClient;
|
||||
|
||||
beforeEach(() => {
|
||||
queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } });
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
const wrapper = ({ children }: { children: ReactNode }) =>
|
||||
React.createElement(QueryClientProvider, { client: queryClient }, children);
|
||||
|
||||
it("polls the flag once it is on, so an already-open dashboard notices it going off", async () => {
|
||||
(getUiSettings as any).mockResolvedValue({ values: { enable_ptu_cost_attribution: true } });
|
||||
const { result } = renderHook(() => usePtuCostAttributionEnabled(), { wrapper });
|
||||
await waitFor(() => {
|
||||
expect(result.current).toBe(true);
|
||||
});
|
||||
|
||||
const observers = queryClient.getQueryCache().getAll()[0].observers;
|
||||
const polling = observers.filter((o: any) => o.options.refetchInterval === PTU_FLAG_REFRESH_MS);
|
||||
expect(polling.length).toBeGreaterThan(0);
|
||||
expect(polling[0].options.staleTime).toBe(PTU_FLAG_REFRESH_MS);
|
||||
expect(PTU_FLAG_REFRESH_MS).toBeLessThan(60 * 60 * 1000);
|
||||
});
|
||||
|
||||
it("does not poll while the flag is off, which is every deployment that never opted in", async () => {
|
||||
// The hook cannot gate on the flag before reading it, so it starts on the shared
|
||||
// one-hour cache and only escalates once it has seen the feature enabled. Polling
|
||||
// unconditionally made a disabled deployment re-fetch settings 120x more often.
|
||||
(getUiSettings as any).mockResolvedValue({ values: { enable_ptu_cost_attribution: false } });
|
||||
const { result } = renderHook(() => usePtuCostAttributionEnabled(), { wrapper });
|
||||
await waitFor(() => {
|
||||
expect(result.current).toBe(false);
|
||||
});
|
||||
|
||||
const observers = queryClient.getQueryCache().getAll()[0].observers;
|
||||
expect(observers.every((o: any) => o.options.refetchInterval === undefined)).toBe(true);
|
||||
expect(observers.every((o: any) => o.options.staleTime === 60 * 60 * 1000)).toBe(true);
|
||||
});
|
||||
|
||||
it("leaves the default alone for every other settings consumer", async () => {
|
||||
(getUiSettings as any).mockResolvedValue({ values: {} });
|
||||
const { result } = renderHook(() => useUISettings(), { wrapper });
|
||||
await waitFor(() => {
|
||||
expect(result.current.isSuccess).toBe(true);
|
||||
});
|
||||
|
||||
const observers = queryClient.getQueryCache().getAll()[0].observers;
|
||||
expect(observers[0].options.staleTime).toBe(60 * 60 * 1000);
|
||||
expect(observers[0].options.refetchInterval).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,26 @@
|
|||
import { useUISettings } from "./useUISettings";
|
||||
|
||||
export const PTU_COST_ATTRIBUTION_SETTING_KEY = "enable_ptu_cost_attribution";
|
||||
|
||||
/**
|
||||
* Whether the proxy opted into PTU flat-cost attribution.
|
||||
*
|
||||
* Derived on the proxy from LITELLM_ENABLE_PTU_COST_ATTRIBUTION and returned read-only on
|
||||
* /get/ui_settings, so it is not editable from the UI. Anything other than an explicit
|
||||
* true (including a settings fetch that has not resolved) counts as off.
|
||||
*
|
||||
* Polled only once the flag has been seen on. This tracks the proxy process rather than a
|
||||
* persisted setting, so an already-open dashboard has to notice a restart that turns the
|
||||
* feature off, and a form that stays mounted and focused never refetches on staleTime
|
||||
* alone. A deployment that never opts in is the common case and gets the shared one-hour
|
||||
* cache, so the poll costs nothing where the feature is unused; the trade is that turning
|
||||
* it on reaches an open dashboard on the next natural refetch rather than within 30s.
|
||||
*/
|
||||
export const PTU_FLAG_REFRESH_MS = 30 * 1000;
|
||||
|
||||
export const usePtuCostAttributionEnabled = (): boolean => {
|
||||
const { data } = useUISettings();
|
||||
const enabled = data?.values?.[PTU_COST_ATTRIBUTION_SETTING_KEY] === true;
|
||||
useUISettings(enabled ? { staleTime: PTU_FLAG_REFRESH_MS, refetchInterval: PTU_FLAG_REFRESH_MS } : undefined);
|
||||
return enabled;
|
||||
};
|
||||
|
|
@ -4,11 +4,21 @@ import { createQueryKeys } from "../common/queryKeysFactory";
|
|||
|
||||
const uiSettingsKeys = createQueryKeys("uiSettings");
|
||||
|
||||
export const useUISettings = () => {
|
||||
/**
|
||||
* UI settings, cached for an hour by default because they rarely change.
|
||||
*
|
||||
* Both options are per observer in react-query, so a caller reading a value that tracks
|
||||
* proxy process state, rather than a persisted setting, can refresh it on its own cadence
|
||||
* without changing how long every other caller caches. `staleTime` alone only marks the
|
||||
* cached copy stale; a screen that stays mounted and focused never refetches on its own,
|
||||
* so a caller that needs to notice a change also has to poll.
|
||||
*/
|
||||
export const useUISettings = (options?: { staleTime?: number; refetchInterval?: number }) => {
|
||||
return useQuery<Record<string, any>>({
|
||||
queryKey: uiSettingsKeys.list({}),
|
||||
queryFn: async () => await getUiSettings(),
|
||||
staleTime: 60 * 60 * 1000, // 1 hour - data rarely changes
|
||||
staleTime: options?.staleTime ?? 60 * 60 * 1000, // 1 hour - data rarely changes
|
||||
gcTime: 60 * 60 * 1000, // 1 hour - keep in cache for 1 hour
|
||||
refetchInterval: options?.refetchInterval,
|
||||
});
|
||||
};
|
||||
|
|
|
|||
|
|
@ -2,32 +2,37 @@ import { act, fireEvent, render, waitFor } from "@testing-library/react";
|
|||
import { beforeEach, describe, expect, it, vi } from "vitest";
|
||||
import AdvancedSettings from "./advanced_settings";
|
||||
|
||||
const mockUsePtuCostAttributionEnabled = vi.fn();
|
||||
|
||||
vi.mock("@/app/(dashboard)/hooks/uiSettings/usePtuCostAttributionEnabled", () => ({
|
||||
usePtuCostAttributionEnabled: () => mockUsePtuCostAttributionEnabled(),
|
||||
}));
|
||||
|
||||
const PTU_LABELS = ["PTU Count", "Calculated Cost per PTU / Hour (USD)", "PTU Effective From (UTC)"];
|
||||
|
||||
const renderAdvancedSettings = () =>
|
||||
render(
|
||||
<AdvancedSettings
|
||||
showAdvancedSettings={true}
|
||||
setShowAdvancedSettings={() => {}}
|
||||
guardrailsList={[]}
|
||||
tagsList={{}}
|
||||
accessToken="test-token"
|
||||
/>,
|
||||
);
|
||||
|
||||
describe("AdvancedSettings", () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
mockUsePtuCostAttributionEnabled.mockReturnValue(false);
|
||||
});
|
||||
|
||||
it("should render", () => {
|
||||
render(
|
||||
<AdvancedSettings
|
||||
showAdvancedSettings={true}
|
||||
setShowAdvancedSettings={() => {}}
|
||||
guardrailsList={[]}
|
||||
tagsList={{}}
|
||||
accessToken="test-token"
|
||||
/>,
|
||||
);
|
||||
renderAdvancedSettings();
|
||||
});
|
||||
|
||||
it("should render tags list", async () => {
|
||||
const { getByText } = render(
|
||||
<AdvancedSettings
|
||||
showAdvancedSettings={true}
|
||||
setShowAdvancedSettings={() => {}}
|
||||
guardrailsList={[]}
|
||||
tagsList={{}}
|
||||
accessToken="test-token"
|
||||
/>,
|
||||
);
|
||||
const { getByText } = renderAdvancedSettings();
|
||||
fireEvent.click(getByText("Advanced Settings"));
|
||||
await waitFor(() => {
|
||||
expect(getByText("Tags")).toBeInTheDocument();
|
||||
|
|
@ -35,15 +40,7 @@ describe("AdvancedSettings", () => {
|
|||
});
|
||||
|
||||
it("should render the litellm params", async () => {
|
||||
const { getByText } = render(
|
||||
<AdvancedSettings
|
||||
showAdvancedSettings={true}
|
||||
setShowAdvancedSettings={() => {}}
|
||||
guardrailsList={[]}
|
||||
tagsList={{}}
|
||||
accessToken="test-token"
|
||||
/>,
|
||||
);
|
||||
const { getByText } = renderAdvancedSettings();
|
||||
act(() => {
|
||||
fireEvent.click(getByText("Advanced Settings"));
|
||||
});
|
||||
|
|
@ -51,4 +48,35 @@ describe("AdvancedSettings", () => {
|
|||
expect(getByText("LiteLLM Params")).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
||||
it("hides every PTU field when PTU cost attribution is disabled", async () => {
|
||||
const { getByText, queryByText } = renderAdvancedSettings();
|
||||
act(() => {
|
||||
fireEvent.click(getByText("Advanced Settings"));
|
||||
});
|
||||
await waitFor(() => {
|
||||
expect(getByText("Tags")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
for (const label of PTU_LABELS) {
|
||||
expect(queryByText(label)).not.toBeInTheDocument();
|
||||
}
|
||||
expect(queryByText("PTU Effective To (UTC)")).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("shows every PTU field when PTU cost attribution is enabled", async () => {
|
||||
mockUsePtuCostAttributionEnabled.mockReturnValue(true);
|
||||
const { getByText } = renderAdvancedSettings();
|
||||
act(() => {
|
||||
fireEvent.click(getByText("Advanced Settings"));
|
||||
});
|
||||
|
||||
await waitFor(() => {
|
||||
expect(getByText("PTU Count")).toBeInTheDocument();
|
||||
});
|
||||
for (const label of PTU_LABELS) {
|
||||
expect(getByText(label)).toBeInTheDocument();
|
||||
}
|
||||
expect(getByText("PTU Effective To (UTC)")).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ import {
|
|||
ptuWindowOrderRule,
|
||||
PTU_END_FIELD,
|
||||
} from "../../utils/ptuValidation";
|
||||
import { usePtuCostAttributionEnabled } from "@/app/(dashboard)/hooks/uiSettings/usePtuCostAttributionEnabled";
|
||||
const { Link } = Typography;
|
||||
|
||||
interface AdvancedSettingsProps {
|
||||
|
|
@ -43,6 +44,7 @@ const AdvancedSettings: React.FC<AdvancedSettingsProps> = ({
|
|||
const [customPricing, setCustomPricing] = React.useState(false);
|
||||
const [pricingModel, setPricingModel] = React.useState<"per_token" | "per_second">("per_token");
|
||||
const [showCacheControl, setShowCacheControl] = React.useState(false);
|
||||
const ptuCostAttributionEnabled = usePtuCostAttributionEnabled();
|
||||
|
||||
// Add validation function for numbers
|
||||
const validateNumber = (_: any, value: string) => {
|
||||
|
|
@ -193,49 +195,53 @@ const AdvancedSettings: React.FC<AdvancedSettingsProps> = ({
|
|||
/>
|
||||
</Form.Item>
|
||||
|
||||
<Form.Item
|
||||
label="PTU Count"
|
||||
name={PTU_COUNT_FIELD}
|
||||
dependencies={[PTU_RATE_FIELD]}
|
||||
rules={[{ validator: validateNumber }, ...ptuCountRules, ptuPairRule(PTU_RATE_FIELD)]}
|
||||
tooltip="Provisioned throughput units for this deployment. Set together with Cost per PTU / Hour and a Team to attribute a flat daily cost."
|
||||
className="mb-4"
|
||||
>
|
||||
<TextInput placeholder="e.g. 15" />
|
||||
</Form.Item>
|
||||
{ptuCostAttributionEnabled && (
|
||||
<>
|
||||
<Form.Item
|
||||
label="PTU Count"
|
||||
name={PTU_COUNT_FIELD}
|
||||
dependencies={[PTU_RATE_FIELD]}
|
||||
rules={[{ validator: validateNumber }, ...ptuCountRules, ptuPairRule(PTU_RATE_FIELD)]}
|
||||
tooltip="Provisioned throughput units for this deployment. Set together with Cost per PTU / Hour and a Team to attribute a flat daily cost."
|
||||
className="mb-4"
|
||||
>
|
||||
<TextInput placeholder="e.g. 15" />
|
||||
</Form.Item>
|
||||
|
||||
<Form.Item
|
||||
label="Calculated Cost per PTU / Hour (USD)"
|
||||
name={PTU_RATE_FIELD}
|
||||
dependencies={[PTU_COUNT_FIELD]}
|
||||
rules={[{ validator: validateNumber }, ...ptuRateRules, ptuPairRule(PTU_COUNT_FIELD)]}
|
||||
tooltip="Flat cost = PTU count * this rate * active hours, attributed to the deployment's team."
|
||||
className="mb-4"
|
||||
>
|
||||
<TextInput placeholder="e.g. 2.00" />
|
||||
</Form.Item>
|
||||
<Form.Item
|
||||
label="Calculated Cost per PTU / Hour (USD)"
|
||||
name={PTU_RATE_FIELD}
|
||||
dependencies={[PTU_COUNT_FIELD]}
|
||||
rules={[{ validator: validateNumber }, ...ptuRateRules, ptuPairRule(PTU_COUNT_FIELD)]}
|
||||
tooltip="Flat cost = PTU count * this rate * active hours, attributed to the deployment's team."
|
||||
className="mb-4"
|
||||
>
|
||||
<TextInput placeholder="e.g. 2.00" />
|
||||
</Form.Item>
|
||||
|
||||
<Form.Item
|
||||
label="PTU Effective From (UTC)"
|
||||
name={PTU_START_FIELD}
|
||||
dependencies={[PTU_COUNT_FIELD, PTU_END_FIELD]}
|
||||
rules={[ptuStartRequiredRule(PTU_COUNT_FIELD), ptuWindowOrderRule(PTU_END_FIELD, "start")]}
|
||||
tooltip="Start of the PTU window, required when PTU Count is set. Flat cost accrues by the hour within the window; a window opening at 23:00 charges one hour that day."
|
||||
className="mb-4"
|
||||
>
|
||||
<DatePicker showTime style={{ width: "100%" }} />
|
||||
</Form.Item>
|
||||
<Form.Item
|
||||
label="PTU Effective From (UTC)"
|
||||
name={PTU_START_FIELD}
|
||||
dependencies={[PTU_COUNT_FIELD, PTU_END_FIELD]}
|
||||
rules={[ptuStartRequiredRule(PTU_COUNT_FIELD), ptuWindowOrderRule(PTU_END_FIELD, "start")]}
|
||||
tooltip="Start of the PTU window, required when PTU Count is set. Flat cost accrues by the hour within the window; a window opening at 23:00 charges one hour that day."
|
||||
className="mb-4"
|
||||
>
|
||||
<DatePicker showTime style={{ width: "100%" }} />
|
||||
</Form.Item>
|
||||
|
||||
<Form.Item
|
||||
label="PTU Effective To (UTC)"
|
||||
name={PTU_END_FIELD}
|
||||
dependencies={[PTU_START_FIELD]}
|
||||
rules={[ptuWindowOrderRule(PTU_START_FIELD, "end")]}
|
||||
tooltip="Optional end of the PTU window (exclusive). Leave blank for open-ended."
|
||||
className="mb-4"
|
||||
>
|
||||
<DatePicker showTime style={{ width: "100%" }} />
|
||||
</Form.Item>
|
||||
<Form.Item
|
||||
label="PTU Effective To (UTC)"
|
||||
name={PTU_END_FIELD}
|
||||
dependencies={[PTU_START_FIELD]}
|
||||
rules={[ptuWindowOrderRule(PTU_START_FIELD, "end")]}
|
||||
tooltip="Optional end of the PTU window (exclusive). Leave blank for open-ended."
|
||||
className="mb-4"
|
||||
>
|
||||
<DatePicker showTime style={{ width: "100%" }} />
|
||||
</Form.Item>
|
||||
</>
|
||||
)}
|
||||
|
||||
{customPricing && (
|
||||
<div className="ml-6 pl-4 border-l-2 border-gray-200">
|
||||
|
|
|
|||
|
|
@ -47,6 +47,11 @@ vi.mock("@/app/(dashboard)/hooks/models/useModelCostMap", () => ({
|
|||
useModelCostMap: (...args: any[]) => mockUseModelCostMap(...args),
|
||||
}));
|
||||
|
||||
const mockUsePtuCostAttributionEnabled = vi.fn();
|
||||
vi.mock("@/app/(dashboard)/hooks/uiSettings/usePtuCostAttributionEnabled", () => ({
|
||||
usePtuCostAttributionEnabled: () => mockUsePtuCostAttributionEnabled(),
|
||||
}));
|
||||
|
||||
const mockNotificationsManager = vi.mocked(NotificationsManager);
|
||||
const mockModelInfoV1Call = vi.mocked(networking.modelInfoV1Call);
|
||||
const mockCredentialGetCall = vi.mocked(networking.credentialGetCall);
|
||||
|
|
@ -99,6 +104,7 @@ describe("ModelInfoView", () => {
|
|||
},
|
||||
});
|
||||
vi.clearAllMocks();
|
||||
mockUsePtuCostAttributionEnabled.mockReturnValue(false);
|
||||
|
||||
mockUseModelsInfo.mockReturnValue({
|
||||
data: {
|
||||
|
|
@ -608,6 +614,100 @@ describe("ModelInfoView", () => {
|
|||
expect(updatePayload.litellm_params).not.toHaveProperty("vector_store_ids");
|
||||
});
|
||||
|
||||
describe("PTU cost attribution gate", () => {
|
||||
const ptuModelData = {
|
||||
...defaultModelData,
|
||||
model_info: {
|
||||
...defaultModelData.model_info,
|
||||
team_id: "team-1",
|
||||
ptu_count: 15,
|
||||
cost_per_ptu_per_hour: 2,
|
||||
ptu_effective_from: "2026-07-01T00:00:00+00:00",
|
||||
ptu_effective_to: "2026-08-01T00:00:00+00:00",
|
||||
},
|
||||
};
|
||||
|
||||
const renderWithPtuModel = () => {
|
||||
mockUseModelsInfo.mockReturnValue({ data: { data: [ptuModelData] }, isLoading: false, error: null });
|
||||
mockModelInfoV1Call.mockResolvedValue({ data: [ptuModelData] });
|
||||
return render(<ModelInfoView {...DEFAULT_ADMIN_PROPS} />, { wrapper });
|
||||
};
|
||||
|
||||
it("hides the PTU fields when disabled, even for a model that already stores PTU config", async () => {
|
||||
renderWithPtuModel();
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByText("Model Settings")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
expect(screen.queryByText("PTU Count")).not.toBeInTheDocument();
|
||||
expect(screen.queryByText("Cost per PTU / Hour (USD)")).not.toBeInTheDocument();
|
||||
expect(screen.queryByText("PTU Effective From (UTC)")).not.toBeInTheDocument();
|
||||
expect(screen.queryByText("PTU Effective To (UTC)")).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("shows the PTU fields when enabled", async () => {
|
||||
mockUsePtuCostAttributionEnabled.mockReturnValue(true);
|
||||
renderWithPtuModel();
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByText("PTU Count")).toBeInTheDocument();
|
||||
});
|
||||
expect(screen.getByText("Cost per PTU / Hour (USD)")).toBeInTheDocument();
|
||||
expect(screen.getByText("PTU Effective From (UTC)")).toBeInTheDocument();
|
||||
expect(screen.getByText("PTU Effective To (UTC)")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("omits PTU fields from the save payload when disabled, so an unrelated edit cannot clear stored config", async () => {
|
||||
const user = userEvent.setup();
|
||||
renderWithPtuModel();
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByRole("button", { name: /edit settings/i })).toBeInTheDocument();
|
||||
});
|
||||
await user.click(screen.getByRole("button", { name: /edit settings/i }));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByRole("button", { name: /save changes/i })).toBeInTheDocument();
|
||||
});
|
||||
await user.click(screen.getByRole("button", { name: /save changes/i }));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(mockModelPatchUpdateCall).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
const modelInfo = mockModelPatchUpdateCall.mock.calls[0][1].model_info;
|
||||
expect(modelInfo).not.toHaveProperty("ptu_count");
|
||||
expect(modelInfo).not.toHaveProperty("cost_per_ptu_per_hour");
|
||||
expect(modelInfo).not.toHaveProperty("ptu_effective_from");
|
||||
expect(modelInfo).not.toHaveProperty("ptu_effective_to");
|
||||
});
|
||||
|
||||
it("sends the PTU fields on save when enabled", async () => {
|
||||
mockUsePtuCostAttributionEnabled.mockReturnValue(true);
|
||||
const user = userEvent.setup();
|
||||
renderWithPtuModel();
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByRole("button", { name: /edit settings/i })).toBeInTheDocument();
|
||||
});
|
||||
await user.click(screen.getByRole("button", { name: /edit settings/i }));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(screen.getByRole("button", { name: /save changes/i })).toBeInTheDocument();
|
||||
});
|
||||
await user.click(screen.getByRole("button", { name: /save changes/i }));
|
||||
|
||||
await waitFor(() => {
|
||||
expect(mockModelPatchUpdateCall).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
const modelInfo = mockModelPatchUpdateCall.mock.calls[0][1].model_info;
|
||||
expect(modelInfo.ptu_count).toBe(15);
|
||||
expect(modelInfo.cost_per_ptu_per_hour).toBe(2);
|
||||
});
|
||||
});
|
||||
|
||||
it("should not include input_cost_per_token or output_cost_per_token in update payload when user does not touch cost fields", async () => {
|
||||
// Regression: editing a model without touching cost fields used to inject
|
||||
// input_cost_per_token: 0 and output_cost_per_token: 0 into litellm_params,
|
||||
|
|
|
|||
|
|
@ -18,7 +18,9 @@ import {
|
|||
Button as TremorButton,
|
||||
} from "@tremor/react";
|
||||
import { Button, DatePicker, Form, Input, Modal, Select, Tooltip } from "antd";
|
||||
import { formatPtuUtcDisplay, ptuPickerToUtcIso, utcIsoToPickerValue } from "../utils/ptuDatetime";
|
||||
import { formatPtuUtcDisplay, utcIsoToPickerValue } from "../utils/ptuDatetime";
|
||||
import { applyPtuModelInfo } from "../utils/ptuModelInfo";
|
||||
import { usePtuCostAttributionEnabled } from "@/app/(dashboard)/hooks/uiSettings/usePtuCostAttributionEnabled";
|
||||
import {
|
||||
PTU_COUNT_FIELD,
|
||||
PTU_RATE_FIELD,
|
||||
|
|
@ -224,6 +226,7 @@ export default function ModelInfoView({
|
|||
const { data: modelCostMapData } = useModelCostMap();
|
||||
const { data: modelHubData } = useModelHub();
|
||||
const { data: teams } = useTeams();
|
||||
const ptuCostAttributionEnabled = usePtuCostAttributionEnabled();
|
||||
|
||||
// Transform the model data
|
||||
const getProviderFromModel = (model: string) => {
|
||||
|
|
@ -495,15 +498,7 @@ export default function ModelInfoView({
|
|||
health_check_model: values.health_check_model,
|
||||
};
|
||||
}
|
||||
const ptuNumber = (val: string | number | null | undefined): number | null =>
|
||||
val !== undefined && val !== null && val !== "" ? Number(val) : null;
|
||||
updatedModelInfo = {
|
||||
...updatedModelInfo,
|
||||
ptu_count: ptuNumber(values.ptu_count),
|
||||
cost_per_ptu_per_hour: ptuNumber(values.cost_per_ptu_per_hour),
|
||||
ptu_effective_from: ptuPickerToUtcIso(values.ptu_effective_from),
|
||||
ptu_effective_to: ptuPickerToUtcIso(values.ptu_effective_to),
|
||||
};
|
||||
updatedModelInfo = applyPtuModelInfo(updatedModelInfo, values, ptuCostAttributionEnabled);
|
||||
} catch (e) {
|
||||
NotificationsManager.fromBackend("Invalid JSON in Model Info");
|
||||
return;
|
||||
|
|
@ -953,45 +948,46 @@ export default function ModelInfoView({
|
|||
)}
|
||||
</div>
|
||||
|
||||
{PTU_EDIT_FIELDS.map((ptuField) => {
|
||||
const { name, label, input, placeholder, isCount, isRate, isStart, pairedWith } = ptuField;
|
||||
const { windowPeer, bound } = ptuField;
|
||||
return (
|
||||
<div key={name}>
|
||||
<Text className="font-medium">{label}</Text>
|
||||
{isEditing ? (
|
||||
<Form.Item
|
||||
name={name}
|
||||
className="mb-0"
|
||||
dependencies={ptuFieldDependencies(ptuField)}
|
||||
rules={[
|
||||
...(isCount ? ptuCountRules : []),
|
||||
...(isRate ? ptuRateRules : []),
|
||||
...(isStart ? [ptuStartRequiredRule(PTU_COUNT_FIELD)] : []),
|
||||
...(pairedWith ? [ptuPairRule(pairedWith)] : []),
|
||||
...(windowPeer && bound ? [ptuWindowOrderRule(windowPeer, bound)] : []),
|
||||
]}
|
||||
>
|
||||
{input === "number" ? (
|
||||
<NumericalInput
|
||||
placeholder={placeholder}
|
||||
step={isCount ? 1 : undefined}
|
||||
min={isCount ? 1 : 0}
|
||||
/>
|
||||
) : (
|
||||
<DatePicker showTime style={{ width: "100%" }} />
|
||||
)}
|
||||
</Form.Item>
|
||||
) : (
|
||||
<div className="mt-1 p-2 bg-gray-50 rounded-sm">
|
||||
{(input === "datetime"
|
||||
? formatPtuUtcDisplay(localModelData?.model_info?.[name])
|
||||
: localModelData?.model_info?.[name]) ?? "Not Set"}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
})}
|
||||
{ptuCostAttributionEnabled &&
|
||||
PTU_EDIT_FIELDS.map((ptuField) => {
|
||||
const { name, label, input, placeholder, isCount, isRate, isStart, pairedWith } = ptuField;
|
||||
const { windowPeer, bound } = ptuField;
|
||||
return (
|
||||
<div key={name}>
|
||||
<Text className="font-medium">{label}</Text>
|
||||
{isEditing ? (
|
||||
<Form.Item
|
||||
name={name}
|
||||
className="mb-0"
|
||||
dependencies={ptuFieldDependencies(ptuField)}
|
||||
rules={[
|
||||
...(isCount ? ptuCountRules : []),
|
||||
...(isRate ? ptuRateRules : []),
|
||||
...(isStart ? [ptuStartRequiredRule(PTU_COUNT_FIELD)] : []),
|
||||
...(pairedWith ? [ptuPairRule(pairedWith)] : []),
|
||||
...(windowPeer && bound ? [ptuWindowOrderRule(windowPeer, bound)] : []),
|
||||
]}
|
||||
>
|
||||
{input === "number" ? (
|
||||
<NumericalInput
|
||||
placeholder={placeholder}
|
||||
step={isCount ? 1 : undefined}
|
||||
min={isCount ? 1 : 0}
|
||||
/>
|
||||
) : (
|
||||
<DatePicker showTime style={{ width: "100%" }} />
|
||||
)}
|
||||
</Form.Item>
|
||||
) : (
|
||||
<div className="mt-1 p-2 bg-gray-50 rounded-sm">
|
||||
{(input === "datetime"
|
||||
? formatPtuUtcDisplay(localModelData?.model_info?.[name])
|
||||
: localModelData?.model_info?.[name]) ?? "Not Set"}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
})}
|
||||
|
||||
<div>
|
||||
<Text className="font-medium">Cache Read Cost (per 1M tokens)</Text>
|
||||
|
|
|
|||
68
ui/litellm-dashboard/src/utils/ptuModelInfo.test.ts
Normal file
68
ui/litellm-dashboard/src/utils/ptuModelInfo.test.ts
Normal file
|
|
@ -0,0 +1,68 @@
|
|||
import dayjs from "dayjs";
|
||||
import utc from "dayjs/plugin/utc";
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { applyPtuModelInfo, PTU_MODEL_INFO_FIELDS } from "./ptuModelInfo";
|
||||
|
||||
dayjs.extend(utc);
|
||||
|
||||
const storedModelInfo = () => ({
|
||||
id: "model-1",
|
||||
team_id: "team-1",
|
||||
ptu_count: 15,
|
||||
cost_per_ptu_per_hour: 2,
|
||||
ptu_effective_from: "2026-07-01T00:00:00.000Z",
|
||||
ptu_effective_to: "2026-08-01T00:00:00.000Z",
|
||||
});
|
||||
|
||||
describe("applyPtuModelInfo", () => {
|
||||
it("folds the form values into model_info when PTU cost attribution is enabled", () => {
|
||||
const result = applyPtuModelInfo(
|
||||
{ id: "model-1", team_id: "team-1" },
|
||||
{
|
||||
ptu_count: "20",
|
||||
cost_per_ptu_per_hour: "3.5",
|
||||
ptu_effective_from: dayjs.utc("2026-09-01T00:00:00.000Z"),
|
||||
ptu_effective_to: null,
|
||||
},
|
||||
true,
|
||||
);
|
||||
|
||||
expect(result).toEqual({
|
||||
id: "model-1",
|
||||
team_id: "team-1",
|
||||
ptu_count: 20,
|
||||
cost_per_ptu_per_hour: 3.5,
|
||||
ptu_effective_from: "2026-09-01T00:00:00.000Z",
|
||||
ptu_effective_to: null,
|
||||
});
|
||||
});
|
||||
|
||||
it("sends an explicit null for a field the operator cleared while enabled", () => {
|
||||
const result = applyPtuModelInfo(storedModelInfo(), { ptu_count: "", cost_per_ptu_per_hour: "" }, true);
|
||||
|
||||
expect(result.ptu_count).toBeNull();
|
||||
expect(result.cost_per_ptu_per_hour).toBeNull();
|
||||
});
|
||||
|
||||
it("strips every PTU field from the payload when PTU cost attribution is disabled", () => {
|
||||
const result = applyPtuModelInfo(storedModelInfo(), { ptu_count: "20", cost_per_ptu_per_hour: "3.5" }, false);
|
||||
|
||||
for (const field of PTU_MODEL_INFO_FIELDS) {
|
||||
expect(Object.keys(result)).not.toContain(field);
|
||||
}
|
||||
expect(result).toEqual({ id: "model-1", team_id: "team-1" });
|
||||
});
|
||||
|
||||
it("never sends a null PTU field when disabled, so an unrelated save cannot clear stored config", () => {
|
||||
const result = applyPtuModelInfo(storedModelInfo(), {}, false);
|
||||
|
||||
expect(Object.values(result)).not.toContain(null);
|
||||
expect("ptu_count" in result).toBe(false);
|
||||
});
|
||||
|
||||
it("leaves non-PTU model_info untouched when disabled", () => {
|
||||
const result = applyPtuModelInfo({ id: "model-1", access_groups: ["a"], health_check_model: "gpt-5.2" }, {}, false);
|
||||
|
||||
expect(result).toEqual({ id: "model-1", access_groups: ["a"], health_check_model: "gpt-5.2" });
|
||||
});
|
||||
});
|
||||
44
ui/litellm-dashboard/src/utils/ptuModelInfo.ts
Normal file
44
ui/litellm-dashboard/src/utils/ptuModelInfo.ts
Normal file
|
|
@ -0,0 +1,44 @@
|
|||
import { Dayjs } from "dayjs";
|
||||
import { ptuPickerToUtcIso } from "./ptuDatetime";
|
||||
import { PTU_COUNT_FIELD, PTU_RATE_FIELD } from "./ptuValidation";
|
||||
|
||||
export const PTU_MODEL_INFO_FIELDS: readonly string[] = [
|
||||
PTU_COUNT_FIELD,
|
||||
PTU_RATE_FIELD,
|
||||
"ptu_effective_from",
|
||||
"ptu_effective_to",
|
||||
];
|
||||
|
||||
export interface PtuFormValues {
|
||||
ptu_count?: string | number | null;
|
||||
cost_per_ptu_per_hour?: string | number | null;
|
||||
ptu_effective_from?: Dayjs | null;
|
||||
ptu_effective_to?: Dayjs | null;
|
||||
}
|
||||
|
||||
const ptuNumber = (value: string | number | null | undefined): number | null =>
|
||||
value !== undefined && value !== null && value !== "" ? Number(value) : null;
|
||||
|
||||
/**
|
||||
* Fold the PTU form values into the model_info an edit is about to save.
|
||||
*
|
||||
* When PTU cost attribution is off the four fields are stripped rather than sent as null:
|
||||
* the form does not render them, so a null would be an explicit clear of config the operator
|
||||
* never saw, and any PTU field present in the payload is rejected by the proxy.
|
||||
*/
|
||||
export const applyPtuModelInfo = (
|
||||
modelInfo: Record<string, unknown>,
|
||||
values: PtuFormValues,
|
||||
enabled: boolean,
|
||||
): Record<string, unknown> => {
|
||||
if (!enabled) {
|
||||
return Object.fromEntries(Object.entries(modelInfo).filter(([key]) => !PTU_MODEL_INFO_FIELDS.includes(key)));
|
||||
}
|
||||
return {
|
||||
...modelInfo,
|
||||
ptu_count: ptuNumber(values.ptu_count),
|
||||
cost_per_ptu_per_hour: ptuNumber(values.cost_per_ptu_per_hour),
|
||||
ptu_effective_from: ptuPickerToUtcIso(values.ptu_effective_from),
|
||||
ptu_effective_to: ptuPickerToUtcIso(values.ptu_effective_to),
|
||||
};
|
||||
};
|
||||
Loading…
Add table
Reference in a new issue