feat(ptu): configure provisioned-throughput flat cost on a model deployment
Some checks failed
Terraform Provider / gofmt, vet, build, test (push) Has been cancelled
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Has been cancelled

Add ptu_count, cost_per_ptu_per_hour, ptu_effective_from and ptu_effective_to to
ModelInfo so a model deployment can carry the inputs for provisioned-throughput
flat-cost attribution. ModelInfo validates per-field bounds (positive count,
non-negative rate, effective_to after effective_from); model/new and
model/{id}/update enforce the cross-field invariant (count and rate set together,
team_id required) on the effective model_info so partial updates validate the
merged result, and v1/model/info returns the fields.

LiteLLM_DailyTeamSpend gains ptu_flat_cost and ptu_source_model_id columns plus a
sentinel api_key constant; the daily rollup that writes them lands in a follow-up
PR. Adding the optional model_info fields is backward compatible; models without
them are unaffected.

ptu_effective_from is required alongside the count and rate rather than optional. Flat
cost accrues from that instant, so an absent start has to be inferred, and inferring it
let a deployment configured today be billed for days it did not exist. Both PTU validators
also run over the merged view before any write on the update path, beside the premium check the create path
already runs there: the team ACL update below autocommits, so a validator raising further
down left the team mutated and the deployment row never written.

The update path validates the model_info a patch would store rather than the patch
alone. An invariant holds over the deployment as it will exist, not over whichever
subset of fields a caller sent, and validating the patch rejected raising the rate on
an already configured model because that patch carries no start of its own.
This commit is contained in:
Yucheng Wang 2026-07-31 01:03:18 -07:00 • committed by Yucheng Zhu
parent f6df762b25
commit d2be1ce6d0
9 changed files with 550 additions and 4 deletions

View file

@ -0,0 +1,2 @@
-- AlterTable
ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN IF NOT EXISTS "ptu_flat_cost" DOUBLE PRECISION NOT NULL DEFAULT 0.0;

View file

@ -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

View file

@ -1718,3 +1718,9 @@ BROWSER_SECURITY_HEADERS: Final[frozenset[str]] = frozenset(
)
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"

View file

@ -15,6 +15,7 @@ import datetime
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 @@ from litellm.types.router import (
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

View file

@ -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

View file

@ -5,7 +5,7 @@ litellm.Router Types - includes RouterConfig, UpdateRouterConfig, ModelInfo etc
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 @@ class ModelInfo(MirroredPricingParams):
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:

View file

@ -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

View file

@ -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

View file

@ -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 */