diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260713010000_add_ptu_columns_to_daily_team_spend/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260713010000_add_ptu_columns_to_daily_team_spend/migration.sql new file mode 100644 index 000000000000..89a0494431b5 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260713010000_add_ptu_columns_to_daily_team_spend/migration.sql @@ -0,0 +1,2 @@ +-- AlterTable +ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN IF NOT EXISTS "ptu_flat_cost" DOUBLE PRECISION NOT NULL DEFAULT 0.0; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 9c871b65f401..0a8afd4b4475 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -30,7 +30,7 @@ model LiteLLM_BudgetTable { end_users LiteLLM_EndUserTable[] // multiple end-users can have the same budget tags LiteLLM_TagTable[] // multiple tags can have the same budget team_membership LiteLLM_TeamMembership[] // budgets of Users within a Team - organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization + organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization } // Models on proxy @@ -893,6 +893,7 @@ model LiteLLM_DailyTeamSpend { api_requests BigInt @default(0) successful_requests BigInt @default(0) failed_requests BigInt @default(0) + ptu_flat_cost Float @default(0.0) created_at DateTime @default(now()) updated_at DateTime @updatedAt diff --git a/litellm/constants.py b/litellm/constants.py index 30d3bb1f26e7..8da7091933be 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1718,3 +1718,9 @@ ) UNSAFE_PROXY_RESPONSE_HEADERS: Final[frozenset[str]] = HTTP_FRAMING_HEADERS | BROWSER_SECURITY_HEADERS + +# PTU reservation rollup writes rows to LiteLLM_DailyTeamSpend with this +# sentinel api_key so PTU flat cost stays distinguishable from real per-request +# spend under the table's composite unique constraint. +PTU_SENTINEL_API_KEY: Final[str] = "__ptu_flat_cost__" +PTU_ROLLUP_JOB_ID: Final[str] = "ptu_flat_cost_rollup_job" diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index c20870058632..cf6af66d5a9e 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -15,6 +15,7 @@ import json from collections.abc import Awaitable, Mapping, Sequence from json import JSONDecodeError +from types import MappingProxyType from typing import Final, Literal, Protocol, cast from fastapi import APIRouter, Depends, Header, HTTPException, Request, status @@ -79,6 +80,7 @@ SPECIAL_MODEL_INFO_PARAMS, Deployment, GenericLiteLLMParams, + ModelInfo, updateDeployment, ) from litellm.utils import get_utc_datetime @@ -233,6 +235,96 @@ def _raise_on_strategy_router_write_violation( ) +_PTU_MODEL_INFO_FIELDS: Final = ("ptu_count", "cost_per_ptu_per_hour", "ptu_effective_from", "ptu_effective_to") + + +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: + return frozenset() + return frozenset( + field + for field in _PTU_MODEL_INFO_FIELDS + if field in model_info.model_fields_set and getattr(model_info, field) is None + ) + + +def _merged_ptu_model_info(*, db_model: Deployment, patch_data: updateDeployment) -> Mapping[str, object]: + """The model_info a patch would store, which is the stored blob updated by the patch. + + A PTU invariant holds over the deployment as it will exist, not over whichever subset + of fields a caller happened to send. + """ + empty: Final[Mapping[str, object]] = MappingProxyType({}) + stored: Final = db_model.model_info.model_dump(exclude_none=True) if db_model.model_info else empty + incoming: Final = patch_data.model_info.model_dump(exclude_none=True) if patch_data.model_info else empty + cleared: Final = _explicitly_cleared_ptu_fields(patch_data.model_info) + return MappingProxyType({k: v for k, v in {**stored, **incoming}.items() if k not in cleared}) + + +def _validate_ptu_model_info(model_info: Mapping[str, object]) -> None: + """Enforce the PTU cross-field invariant on the effective model_info. + + ptu_count and cost_per_ptu_per_hour must be set together, and a team_id and a + ptu_effective_from are required when they are. The start is mandatory rather than + defaulted because flat cost accrues from it: inferring one would let a deployment + configured today be billed for days it did not exist. Per-field bounds (positive + count, non-negative rate) are enforced by ModelInfo itself. + + Window ordering is checked before the count/rate gate. A patch that touches only one + end of the window carries no count or rate, and ModelInfo sees one field at a time, so + leaving it to either would let an inverted window reach the row; the next load then + fails to parse it and drops the deployment out of the router, where no further patch + can repair it because each one re-parses the stored value first. + """ + effective_from: Final = _coerce_ptu_datetime(model_info.get("ptu_effective_from")) + effective_to: Final = _coerce_ptu_datetime(model_info.get("ptu_effective_to")) + if effective_from is not None and effective_to is not None and effective_to <= effective_from: + raise HTTPException(status_code=400, detail="ptu_effective_to must be after ptu_effective_from") + + has_count: Final = model_info.get("ptu_count") is not None + has_rate: Final = model_info.get("cost_per_ptu_per_hour") is not None + if not has_count and not has_rate: + return + if has_count != has_rate: + raise HTTPException(status_code=400, detail="ptu_count and cost_per_ptu_per_hour must be set together") + if effective_from is None: + raise HTTPException( + status_code=400, + detail=( + "ptu_effective_from is required when PTU fields are set. Flat cost accrues from that " + "instant, so without it the start would have to be inferred and a deployment configured " + "today could be billed for days it did not exist" + ), + ) + if not model_info.get("team_id"): + raise HTTPException( + status_code=400, detail="team_id is required when PTU fields are set (one model maps to one team)" + ) + + +def _parse_ptu_datetime(value: object) -> datetime.datetime | None: + """``value`` as a datetime, parsing an ISO string, else None.""" + if isinstance(value, datetime.datetime): + return value + if not isinstance(value, str): + return None + try: + return datetime.datetime.fromisoformat(value.replace("Z", "+00:00")) + except ValueError: + return None + + +def _coerce_ptu_datetime(value: object) -> datetime.datetime | None: + """Coerce a model_info effective-window value (datetime or ISO string) to UTC, else None.""" + parsed: Final = _parse_ptu_datetime(value) + if parsed is None: + return None + if parsed.tzinfo is None: + return parsed.replace(tzinfo=datetime.timezone.utc) + return parsed.astimezone(datetime.timezone.utc) + + def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> PrismaCompatibleUpdateDBModel: 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) @@ -270,6 +362,10 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr if field in SPECIAL_MODEL_INFO_PARAMS and getattr(updated_patch.model_info, field) is None: merged_model_info.pop(field, None) merged_litellm_params.pop(field, None) + for field in _explicitly_cleared_ptu_fields(updated_patch.model_info): + merged_model_info.pop(field, None) + + _validate_ptu_model_info(merged_model_info) # convert to prisma compatible format @@ -716,6 +812,17 @@ async def _update_team_model_in_db( premium_user=premium_user, ) + # Validated before any write, beside the premium check the create path already runs + # here. The team ACL is updated below and autocommits, so a validator that raises + # further down would leave the team mutated and the deployment row never written. + # + # The merged view is what gets stored, so that is what has to satisfy the invariants. + # Validating the patch alone rejected a partial edit of an already valid deployment: + # 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: + _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 # No team_id in patch, proceed with standard update @@ -1424,6 +1531,8 @@ 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)) + if store_model_in_db is True: """ - store model_list in db diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 9c871b65f401..0a8afd4b4475 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -30,7 +30,7 @@ model LiteLLM_BudgetTable { end_users LiteLLM_EndUserTable[] // multiple end-users can have the same budget tags LiteLLM_TagTable[] // multiple tags can have the same budget team_membership LiteLLM_TeamMembership[] // budgets of Users within a Team - organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization + organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization } // Models on proxy @@ -893,6 +893,7 @@ model LiteLLM_DailyTeamSpend { api_requests BigInt @default(0) successful_requests BigInt @default(0) failed_requests BigInt @default(0) + ptu_flat_cost Float @default(0.0) created_at DateTime @default(now()) updated_at DateTime @updatedAt diff --git a/litellm/types/router.py b/litellm/types/router.py index 4280da08cbbf..8c9080225965 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -5,7 +5,7 @@ import datetime import enum from dataclasses import dataclass -from typing import Any, Final, Generic, Literal, TypeVar, get_type_hints +from typing import Any, ClassVar, Final, Generic, Literal, TypeVar, get_type_hints import httpx from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator @@ -127,6 +127,14 @@ class UpdateRouterConfig(BaseModel): model_config = ConfigDict(protected_namespaces=()) +def _as_utc(value: datetime.datetime | None) -> datetime.datetime | None: + if value is None: + return None + if value.tzinfo is None: + return value.replace(tzinfo=datetime.timezone.utc) + return value.astimezone(datetime.timezone.utc) + + class ModelInfo(MirroredPricingParams): id: str | None # Allow id to be optional on input, but it will always be present as a str in the model instance db_model: bool = False # used for proxy - to separate models which are stored in the db vs. config. @@ -151,6 +159,17 @@ class ModelInfo(MirroredPricingParams): # admin-toggled pause flag; mirrors LiteLLM_ProxyModelTable.blocked blocked: bool | None = None + # Bounds live on the model rather than litellm.constants: names there reach + # litellm/__init__ through several modules' star re-exports, and a Final rebound that + # way trips the basedpyright gate. + MAX_PTU_COUNT: ClassVar[int] = 1_000_000 + MAX_COST_PER_PTU_PER_HOUR: ClassVar[float] = 1_000_000.0 + + ptu_count: int | None = None + cost_per_ptu_per_hour: float | None = None + ptu_effective_from: datetime.datetime | None = None + ptu_effective_to: datetime.datetime | None = None + def __init__(self, id: str | int | None = None, **params) -> None: if id is None: id = str(uuid.uuid4()) # Generate a UUID if id is None or not provided @@ -158,6 +177,23 @@ def __init__(self, id: str | int | None = None, **params) -> None: id = str(id) super().__init__(id=id, **params) + @model_validator(mode="after") + def _validate_ptu_bounds(self) -> "ModelInfo": + if self.ptu_count is not None and not 0 < self.ptu_count <= self.MAX_PTU_COUNT: + raise ValueError(f"ptu_count must be a positive integer no greater than {self.MAX_PTU_COUNT}") + if ( + self.cost_per_ptu_per_hour is not None + and not 0 <= self.cost_per_ptu_per_hour <= self.MAX_COST_PER_PTU_PER_HOUR + ): + raise ValueError( + f"cost_per_ptu_per_hour must be a finite number between 0 and {self.MAX_COST_PER_PTU_PER_HOUR}" + ) + start: Final = _as_utc(self.ptu_effective_from) + end: Final = _as_utc(self.ptu_effective_to) + if start is not None and end is not None and end <= start: + raise ValueError("ptu_effective_to must be after ptu_effective_from") + return self + model_config = ConfigDict(extra="allow") def __contains__(self, key) -> bool: diff --git a/schema.prisma b/schema.prisma index 9c871b65f401..0a8afd4b4475 100644 --- a/schema.prisma +++ b/schema.prisma @@ -30,7 +30,7 @@ model LiteLLM_BudgetTable { end_users LiteLLM_EndUserTable[] // multiple end-users can have the same budget tags LiteLLM_TagTable[] // multiple tags can have the same budget team_membership LiteLLM_TeamMembership[] // budgets of Users within a Team - organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization + organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization } // Models on proxy @@ -893,6 +893,7 @@ model LiteLLM_DailyTeamSpend { api_requests BigInt @default(0) successful_requests BigInt @default(0) failed_requests BigInt @default(0) + ptu_flat_cost Float @default(0.0) created_at DateTime @default(now()) updated_at DateTime @updatedAt diff --git a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py new file mode 100644 index 000000000000..e8131854acfe --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py @@ -0,0 +1,382 @@ +import datetime +import json + +"""Tests for PTU config on the model deployment (v1 model-settings design).""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi import HTTPException + +from litellm.proxy.management_endpoints.model_management_endpoints import ( + _merged_ptu_model_info, + _validate_ptu_model_info, +) +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) + 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) + + +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) + + +def test_model_info_rejects_a_count_beyond_the_cap(): + """flat cost multiplies the count by a float, and an unbounded int overflows that + conversion, which aborted the rollup for every team rather than skipping one model.""" + with pytest.raises(ValueError): + ModelInfo(id="x", team_id="t", ptu_count=10**400, cost_per_ptu_per_hour=2.0) + + +def test_model_info_accepts_a_count_at_the_cap(): + info = ModelInfo(id="x", team_id="t", ptu_count=ModelInfo.MAX_PTU_COUNT, cost_per_ptu_per_hour=2.0) + assert info.ptu_count == ModelInfo.MAX_PTU_COUNT + + +@pytest.mark.parametrize("rate", [float("nan"), float("inf"), float("-inf")]) +def test_model_info_rejects_a_non_finite_rate(rate): + """NaN compares False against every bound, so a bare `< 0` check let it through and the + deployment then accrued a flat cost of nan.""" + with pytest.raises(ValueError): + ModelInfo(id="x", team_id="t", ptu_count=5, cost_per_ptu_per_hour=rate) + + +def test_model_info_rejects_a_rate_beyond_the_cap(): + with pytest.raises(ValueError): + ModelInfo(id="x", team_id="t", ptu_count=5, cost_per_ptu_per_hour=ModelInfo.MAX_COST_PER_PTU_PER_HOUR * 2) + + +def test_model_info_allows_partial_delta_for_patch(): + # A PATCH delta may carry only one field; bounds-only validation must not reject it. + info = ModelInfo(id="x", ptu_count=5) + assert info.ptu_count == 5 + assert info.cost_per_ptu_per_hour is None + + +def test_validate_helper_no_ptu_is_noop(): + _validate_ptu_model_info({"team_id": "t"}) + + +def test_validate_helper_requires_both_fields(): + with pytest.raises(HTTPException) as exc: + _validate_ptu_model_info({"team_id": "t", "ptu_count": 5}) + assert exc.value.status_code == 400 + assert "set together" in exc.value.detail + + +def test_validate_helper_requires_team_id(): + with pytest.raises(HTTPException) as exc: + _validate_ptu_model_info( + {"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "ptu_effective_from": "2026-08-01T00:00:00Z"} + ) + assert exc.value.status_code == 400 + assert "team_id" in exc.value.detail + + +def test_validate_helper_requires_an_effective_start(): + """Flat cost accrues from the start, so it cannot be inferred.""" + with pytest.raises(HTTPException) as exc: + _validate_ptu_model_info({"team_id": "t", "ptu_count": 5, "cost_per_ptu_per_hour": 2.0}) + assert exc.value.status_code == 400 + assert "ptu_effective_from is required" in exc.value.detail + + +def test_validate_helper_passes_full_config(): + _validate_ptu_model_info( + {"team_id": "t", "ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "ptu_effective_from": "2026-08-01T00:00:00Z"} + ) + + +def test_model_info_rejects_effective_to_before_from(): + import datetime + + with pytest.raises(ValueError): + ModelInfo( + id="x", + team_id="t", + ptu_count=5, + cost_per_ptu_per_hour=2.0, + ptu_effective_from=datetime.datetime(2026, 7, 30, tzinfo=datetime.timezone.utc), + ptu_effective_to=datetime.datetime(2026, 7, 29, tzinfo=datetime.timezone.utc), + ) + + +def test_model_info_accepts_valid_effective_window(): + import datetime + + info = ModelInfo( + id="x", + team_id="t", + ptu_count=5, + cost_per_ptu_per_hour=2.0, + ptu_effective_from=datetime.datetime(2026, 7, 30, tzinfo=datetime.timezone.utc), + ptu_effective_to=datetime.datetime(2026, 8, 30, tzinfo=datetime.timezone.utc), + ) + assert info.ptu_effective_from is not None + + +def test_model_info_compares_mixed_naive_and_aware_timestamps(): + import datetime + + info = ModelInfo( + id="x", + team_id="t", + ptu_count=5, + cost_per_ptu_per_hour=2.0, + ptu_effective_from=datetime.datetime(2026, 7, 30, 23, 0), + ptu_effective_to=datetime.datetime(2026, 7, 31, 0, 0, tzinfo=datetime.timezone.utc), + ) + assert info.ptu_effective_to is not None + + with pytest.raises(ValueError): + ModelInfo( + id="x", + team_id="t", + ptu_count=5, + cost_per_ptu_per_hour=2.0, + ptu_effective_from=datetime.datetime(2026, 7, 31, 2, 0), + ptu_effective_to=datetime.datetime(2026, 7, 31, 0, 0, tzinfo=datetime.timezone.utc), + ) + + +def test_validate_helper_rejects_effective_to_before_from(): + with pytest.raises(HTTPException) as exc: + _validate_ptu_model_info( + { + "team_id": "t", + "ptu_count": 5, + "cost_per_ptu_per_hour": 2.0, + "ptu_effective_from": "2026-07-30T00:00:00Z", + "ptu_effective_to": "2026-07-29T00:00:00Z", + } + ) + assert exc.value.status_code == 400 + assert "ptu_effective_to" in exc.value.detail + + +def test_validate_helper_accepts_valid_window_on_merged_info(): + _validate_ptu_model_info( + { + "team_id": "t", + "ptu_count": 5, + "cost_per_ptu_per_hour": 2.0, + "ptu_effective_from": "2026-07-30T00:00:00Z", + "ptu_effective_to": "2026-08-30T00:00:00Z", + } + ) + + +def test_validate_helper_rejects_inverted_window_without_count_or_rate(): + """A patch that touches only one end of the window merges to a model_info with no count + or rate. Returning early on that shape let an inverted window reach the row, and the next + load then failed to parse it and dropped the deployment out of the router.""" + with pytest.raises(HTTPException) as exc: + _validate_ptu_model_info( + { + "team_id": "t", + "ptu_effective_from": "2026-08-02T00:00:00Z", + "ptu_effective_to": "2026-08-01T00:00:00Z", + } + ) + assert exc.value.status_code == 400 + assert "ptu_effective_to" in exc.value.detail + + +def test_validate_helper_rejects_equal_window_bounds_without_count_or_rate(): + with pytest.raises(HTTPException) as exc: + _validate_ptu_model_info( + { + "ptu_effective_from": "2026-08-01T00:00:00Z", + "ptu_effective_to": "2026-08-01T00:00:00Z", + } + ) + assert exc.value.status_code == 400 + + +def test_validate_helper_accepts_ordered_window_without_count_or_rate(): + """Window-only edits stay legal; only the ordering is enforced, and no team_id is + demanded while the deployment carries no priced PTU config.""" + _validate_ptu_model_info( + { + "ptu_effective_from": "2026-08-01T00:00:00Z", + "ptu_effective_to": "2026-08-02T00:00:00Z", + } + ) + + +def test_validate_helper_accepts_a_single_open_ended_bound(): + _validate_ptu_model_info({"ptu_effective_from": "2026-08-01T00:00:00Z"}) + _validate_ptu_model_info({"ptu_effective_to": "2026-08-02T00:00:00Z"}) + + +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.""" + + @staticmethod + def _configured(): + 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=10, + cost_per_ptu_per_hour=2.0, + ptu_effective_from=datetime.datetime(2026, 7, 1, tzinfo=datetime.timezone.utc), + ), + ) + + def test_raising_the_rate_on_a_configured_model_is_allowed(self): + """The patch carries no start; the stored row supplies it.""" + merged = _merged_ptu_model_info( + db_model=self._configured(), + patch_data=updateDeployment(model_info=ModelInfo(id="dep-0", ptu_count=10, cost_per_ptu_per_hour=3.0)), + ) + _validate_ptu_model_info(merged) + assert merged["cost_per_ptu_per_hour"] == 3.0 + assert merged["ptu_effective_from"] is not None + + def test_a_genuinely_startless_configuration_is_still_rejected(self): + """Merging must not become a way to smuggle PTU config in without a start.""" + bare = Deployment(model_name="gpt-4o", litellm_params=LiteLLM_Params(model="openai/gpt-4o")) + merged = _merged_ptu_model_info( + db_model=bare, + patch_data=updateDeployment( + model_info=ModelInfo(id="dep-0", team_id="t", ptu_count=10, cost_per_ptu_per_hour=2.0) + ), + ) + with pytest.raises(HTTPException) as exc: + _validate_ptu_model_info(merged) + assert "ptu_effective_from is required" in exc.value.detail + + def test_the_patch_still_wins_over_the_stored_value(self): + merged = _merged_ptu_model_info( + db_model=self._configured(), + patch_data=updateDeployment(model_info=ModelInfo(id="dep-0", ptu_count=25)), + ) + assert merged["ptu_count"] == 25 + + def test_an_explicit_null_clears_the_stored_field(self): + """update_db_model drops a PTU field a patch sends as null, so the merged view has to + drop it too. Carrying the stored value forward validated a deployment that never + existed.""" + merged = _merged_ptu_model_info( + db_model=self._configured(), + patch_data=updateDeployment(model_info=ModelInfo(id="dep-0", ptu_count=None)), + ) + assert "ptu_count" not in merged + + def test_clearing_one_half_of_the_pair_is_rejected(self): + """The write leaves a rate with no count. Merging on the stored count hid that.""" + merged = _merged_ptu_model_info( + db_model=self._configured(), + patch_data=updateDeployment(model_info=ModelInfo(id="dep-0", ptu_count=None)), + ) + with pytest.raises(HTTPException) as exc: + _validate_ptu_model_info(merged) + assert "must be set together" in exc.value.detail + + def test_clearing_the_whole_pair_is_allowed(self): + """Turning PTU off on a deployment is a legitimate edit.""" + merged = _merged_ptu_model_info( + db_model=self._configured(), + patch_data=updateDeployment(model_info=ModelInfo(id="dep-0", ptu_count=None, cost_per_ptu_per_hour=None)), + ) + _validate_ptu_model_info(merged) + assert "ptu_count" not in merged + assert "cost_per_ptu_per_hour" not in merged + + def test_an_omitted_field_is_not_a_clear(self): + """A partial edit that never mentions the count keeps it. Only an explicit null clears.""" + merged = _merged_ptu_model_info( + db_model=self._configured(), + patch_data=updateDeployment(model_info=ModelInfo(id="dep-0", cost_per_ptu_per_hour=3.0)), + ) + assert merged["ptu_count"] == 10 + + +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.""" + + @staticmethod + async def _run(db_model, patch_data, monkeypatch, touched=None): + import litellm.proxy.management_endpoints.model_management_endpoints as mme + + touched = [] if touched is None else touched + + async def _never(*args, **kwargs): + touched.append("team_write") + + monkeypatch.setattr(mme, "_setup_new_team_model_assignment", _never) + monkeypatch.setattr(mme, "_update_existing_team_model_assignment", _never) + monkeypatch.setattr(mme.ModelManagementAuthChecks, "allow_team_model_action", AsyncMock(return_value=True)) + result = await mme._update_team_model_in_db( + db_model=db_model, + patch_data=patch_data, + user_api_key_dict=MagicMock(), + prisma_client=MagicMock(), + ) + return result, touched + + @pytest.mark.asyncio + async def test_raising_the_rate_on_a_configured_model_reaches_the_write(self, monkeypatch): + """The patch carries no start. Validating it alone rejected this ordinary edit.""" + db_model = TestPartialPtuEditsUseTheMergedView._configured() + patch = updateDeployment(model_info=ModelInfo(id="dep-0", team_id="t", ptu_count=10, cost_per_ptu_per_hour=3.0)) + + result, touched = await self._run(db_model, patch, monkeypatch) + + assert touched == ["team_write"] + assert json.loads(result["model_info"])["cost_per_ptu_per_hour"] == 3.0 + + @pytest.mark.asyncio + async def test_a_startless_configuration_is_refused_before_the_team_write(self, monkeypatch): + """And the refusal still lands before anything is committed.""" + bare = Deployment(model_name="gpt-4o", litellm_params=LiteLLM_Params(model="openai/gpt-4o")) + patch = updateDeployment(model_info=ModelInfo(id="dep-0", team_id="t", ptu_count=10, cost_per_ptu_per_hour=2.0)) + + with pytest.raises(HTTPException) as exc: + await self._run(bare, patch, monkeypatch) + + assert "ptu_effective_from is required" in exc.value.detail + + @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 + deployment with a rate and no count commit.""" + db_model = TestPartialPtuEditsUseTheMergedView._configured() + patch = updateDeployment(model_info=ModelInfo(id="dep-0", team_id="t", ptu_count=None)) + touched = [] + + with pytest.raises(HTTPException) as exc: + await self._run(db_model, patch, monkeypatch, touched) + + assert "must be set together" in exc.value.detail + assert touched == [] + + @pytest.mark.asyncio + async def test_clearing_the_whole_pair_reaches_the_write_and_stores_neither_field(self, monkeypatch): + """What the validator approved is what the write persists.""" + db_model = TestPartialPtuEditsUseTheMergedView._configured() + patch = updateDeployment( + model_info=ModelInfo(id="dep-0", team_id="t", ptu_count=None, cost_per_ptu_per_hour=None) + ) + + result, touched = await self._run(db_model, patch, monkeypatch) + + assert touched == ["team_write"] + stored = json.loads(result["model_info"]) + assert "ptu_count" not in stored + assert "cost_per_ptu_per_hour" not in stored diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index a75c23da1cf8..e6c3627a230a 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -35298,6 +35298,8 @@ export interface components { cache_creation_input_token_cost?: number | null; /** Cache Read Input Token Cost */ cache_read_input_token_cost?: number | null; + /** Cost Per Ptu Per Hour */ + cost_per_ptu_per_hour?: number | null; /** Created At */ created_at?: string | null; /** Created By */ @@ -35317,6 +35319,12 @@ export interface components { output_cost_per_character?: number | null; /** Output Cost Per Token */ output_cost_per_token?: number | null; + /** Ptu Count */ + ptu_count?: number | null; + /** Ptu Effective From */ + ptu_effective_from?: string | null; + /** Ptu Effective To */ + ptu_effective_to?: string | null; /** Team Id */ team_id?: string | null; /** Team Public Model Name */