feat(proxy): split Azure PTU deployments across teams with per-team PTU ceilings and PTU-hour reporting (#43043)

* feat(proxy): split Azure PTU deployments across teams by share

A PTU deployment's model_info takes ptu_shares (team_id -> whole PTUs adding up to
ptu_count) in place of a single team_id. With LITELLM_ENABLE_PTU_COST_ATTRIBUTION on,
routing hides the deployment from teams holding no share, the v3 limiter enforces each
share as a per-team TPM ceiling from Azure's per-model TPM-per-PTU sizing, the daily
flat-cost rollup splits the hourly cost by share, and /team/daily/activity plus the Team
Usage page report PTU-hours next to tokens.

* fix(proxy): weight cached input in PTU ceilings and size team-scoped deployments by public name

The PTU flag is read from the process environment on every check instead of
through the secret-manager path, so a proxy with a hosted secret manager does
not pay a round trip per request. A team's ceiling now carries the model's
cached-input ratio and settlement charges cache reads at that ratio, matching
how Azure sizes a PTU. A deployment registered for one team is found by its
team_public_model_name for PTU-hours and the ceiling, since its model_name is
rewritten to a unique internal name. ptu_shares are StrictInt on ModelInfo and
parsed_ptu_shares refuses booleans, strings, and non-string team ids, a
fractional ptu_count reserves nothing, and models_processed counts distinct
deployments rather than holdings.

* fix(proxy): reserve PTU ceilings in normalized tokens, keep share routing flag-independent, and warn only shared or Azure reservations

* fix(proxy): size the uncapped output floor against the PTU share in raw output tokens and escape the exported team id

* fix(ptu): refuse ptu_shares naming a team that does not exist

* refactor(ptu): keep Azure detection under llms/azure and check share teams with one query

* fix(proxy): drop the reverted TeamIdSearchFilter import after merging main

Main reverted the team key search change that added TeamIdSearchFilter, so the PTU share team check types its where clause locally

* fix(proxy): chunk the ptu_shares team lookup under the IN-list bound

* fix(proxy): count a request naming a shared PTU deployment directly against its group's ceiling

* fix(proxy): resolve every name a shared PTU deployment is served under to its group's ceiling

* fix(ptu): refuse ptu_shares declared without ptu_count and cost_per_ptu_per_hour

* fix(ui): seed total_ptu_hours in the empty daily activity metadata

* fix(ptu): size usage rows by the name the request used

* fix(ptu): size a team's usage rows by the deployment it is served from

* fix(ptu): never size a team's usage row by another team's reservation

* fix(ptu): reserve an uncapped request against the cap the proxy writes, not its prompt length

* fix(ptu): read a share map only through the terms registration requires

A ptu_shares map stored without ptu_count, cost_per_ptu_per_hour and
ptu_effective_from, or one that does not add up to ptu_count, is refused at
registration but was enforced once loaded: every other team got the 400 and
the share holder a ceiling while no flat cost accrued. _deployment_shares now
reads the shares through ptu_terms, so such a row reserves nothing, the same
as before the map was read.

* fix(ptu): keep a declared split as the access rule and check a config split at boot with the flag off

6fb24d1d read ptu_shares through the pricing terms so a termless map reserved
nothing, but with the flag off the router skips ptu_config_error for config.yaml
entries, so such an entry loaded and served every team, against the rule that a
declared split restricts access whether or not cost attribution is on.
_deployment_shares reads the declared map again, and the router now raises
ptu_config_error for a config-sourced entry declaring ptu_shares regardless of
the flag, so a malformed split is refused at boot instead of being honoured or
ignored silently.

* fix(ptu): charge wildcard-served names, pass-through usage, and team-less settlements against the team's PTU ceiling

A name served by a wildcard route now resolves to the shared deployment behind it, so the
limiter's ceiling and the team daily activity sizing both find it. Pass-through responses
settle the PTU counter from combined_usage_object in Azure-normalized tokens, and a
reservation taken at admission is settled even when the success metadata lost the team.

* fix(ptu): resolve a wildcard-served name to the one pattern the router picks, and a deployment id before any wildcard

A team's wildcard reservation behind an open proxy-wide * route was charged and sized for traffic the * route served, and naming a shared deployment by id next to a catch-all wildcard escaped its per-team ceiling. Router.wildcard_route_deployments now owns the most-specific-proxy-wide-else-team precedence, and both the limiter and the consumption page resolve through it after listed rows and deployment ids.

* test(router): pin the wildcard route precedence that PTU sizing and the router share

* fix(proxy): accept whole-number PTU shares written as floats, as ptu_count already does

A share of 2.0 now reads as 2 on POST /model/new and in config.yaml, where it was
refused with a 422 and dropped at load. Fractional, infinite, boolean and string
shares stay refused. Adds unit tests for the malformed-rate, provider-over-prefix,
zero-token Responses usage, fractional-count and sizing-warning paths.

* fix(proxy): settle team PTU reservations in their own minute and keep billed PTU days when shares move

* fix(proxy): write each PTU deployment day's team charges in one transaction

The backfill counts a deployment day as priced once any team's row exists for it,
so a rollup that wrote team A's share and lost team C's left team C unbilled for
good. Each deployment day's team rows now go through one Prisma batch with the
retry around the whole batch, so a day is fully written or not written at all and
the next backfill prices it for every team.

* fix(proxy): read the PTU flag from a secret manager, serve share holders from their share, and refuse unfittable calls and share-holding team deletes

* test(router): list the cached PTU ceiling helpers in the router coverage check

* fix(proxy): read the PTU flag only for share holders, guard org and SCIM team deletes, keep share holders on their shared deployment through retries and the scheduler, and date PTU-hours by the terms in force

* fix(proxy): bill a rerun PTU day to the current share holders alone

* fix(proxy): replace a PTU day's rows by day and deployment, not a negated team list

* fix(proxy): read the PTU flag from the secret manager first, like get_secret

* fix(router): scope retry health lookups to the caller only for PTU share holders

* fix(router): move the share-holder retry check into ptu_shares

Adds ptu_hours and total_ptu_hours to the daily activity integration goldens and fixtures, since every route now reports them

* test(integration): cover PTU shares across routing, delete guards, cost attribution and config regex timing

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-10-15 18:44:30 +00:00 • committed by GitHub
parent 3151c47960
commit e344dbd8db
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
50 changed files with 6186 additions and 256 deletions

View file

@ -6,23 +6,41 @@ together because they have to agree: a deployment the rollup declines to charge
router prices at zero serves its traffic for free.
"""
import os
from collections.abc import Mapping
from dataclasses import dataclass
from datetime import date, datetime, time, timezone
from datetime import date, datetime, time, timedelta, timezone
from functools import lru_cache
from time import monotonic
from types import MappingProxyType
from typing import Final
from litellm.secret_managers.main import get_secret_bool
from typing_extensions import TypeIs # noqa: TID251 # TypeIs reaches typing only on 3.13
from litellm.secret_managers.main import get_secret_bool, secret_manager_would_be_consulted, str_to_bool
from litellm.types.router import ModelInfo
from litellm.types.utils import AzureSpillover, CustomPricingLiteLLMParams, MirroredPricingParams
PTU_COST_ATTRIBUTION_ENV_VAR: Final = "LITELLM_ENABLE_PTU_COST_ATTRIBUTION"
SECRET_MANAGER_FLAG_REFRESH_SECONDS: Final = 60
AZURE_SPILLOVER_HEADER: Final = "x-ms-is-spilled-over"
AZURE_SPILLOVER_FROM_HEADER: Final = "x-ms-spillover-from-deployment"
def is_ptu_cost_attribution_enabled() -> bool:
"""Whether PTU flat-cost attribution is turned on for this process."""
"""Whether PTU flat-cost attribution is turned on for this process.
A configured secret manager that ``get_secret`` would ask decides, as it does for every
other setting, read at most once a minute, since the router and the rate limiter ask on
every request. Otherwise the environment decides.
"""
if secret_manager_would_be_consulted(PTU_COST_ATTRIBUTION_ENV_VAR):
return _enabled_in_secret_manager(int(monotonic() // SECRET_MANAGER_FLAG_REFRESH_SECONDS))
return str_to_bool(os.environ.get(PTU_COST_ATTRIBUTION_ENV_VAR)) is True
@lru_cache(maxsize=1)
def _enabled_in_secret_manager(_refresh_window: int) -> bool:
return get_secret_bool(PTU_COST_ATTRIBUTION_ENV_VAR, False) is True
@ -56,14 +74,25 @@ PTU_ZEROED_PRICING: Final[Mapping[str, float | tuple[()] | Mapping[str, float]]]
@dataclass(frozen=True, slots=True)
class PTUTerms:
"""The reservation a deployment declares, once every field has been validated."""
"""The reservation a deployment declares, once every field has been validated.
team_id: str
``shares`` maps every team the capacity is attributed to onto its PTUs and adds up to
``ptu_count``: a deployment declaring a single ``team_id`` holds the whole count under it.
"""
shares: Mapping[str, int]
ptu_count: int
cost_per_ptu_per_hour: float
effective_from: datetime
effective_to: datetime | None
def is_active_on(self, day: date) -> bool:
"""Whether the reservation covers any part of ``day`` (UTC)."""
day_start: Final = datetime.combine(day, time.min, tzinfo=timezone.utc)
return self.effective_from < day_start + timedelta(days=1) and (
self.effective_to is None or self.effective_to > day_start
)
def _to_utc(parsed: datetime) -> datetime:
"""``parsed`` as UTC, reading a naive value as UTC rather than local time."""
@ -127,7 +156,72 @@ def ptu_identity_error(
return None
PTU_MODEL_INFO_FIELDS: Final = ("ptu_count", "cost_per_ptu_per_hour", "ptu_effective_from", "ptu_effective_to")
PTU_MODEL_INFO_FIELDS: Final = (
"ptu_count",
"cost_per_ptu_per_hour",
"ptu_effective_from",
"ptu_effective_to",
"ptu_shares",
)
def _is_mapping(
value: object,
) -> TypeIs[Mapping[object, object]]: # guard-ok: isinstance decides it, keys and values stay object
return isinstance(value, Mapping)
def _whole_ptus(share: object) -> int | None:
"""``share`` as a positive whole number of PTUs, with ``2.0`` read as 2, else None.
``bool`` is excluded because it is an ``int`` subclass and ``True`` would read as one PTU.
"""
if isinstance(share, bool):
return None
whole: Final = int(share) if isinstance(share, float) and share.is_integer() else share
return whole if isinstance(whole, int) and whole > 0 else None
def parsed_ptu_shares(raw: object) -> Mapping[str, int] | None:
"""``ptu_shares`` as team id -> whole PTUs, else None when empty or any entry is unusable."""
if not _is_mapping(raw) or not raw:
return None
candidates: Final = tuple(
(team_id, _whole_ptus(share)) for team_id, share in raw.items() if isinstance(team_id, str) and team_id
)
entries: Final = tuple((team_id, share) for team_id, share in candidates if share is not None)
if len(entries) != len(raw):
return None
return MappingProxyType(dict(entries))
def _parsed_ptu_count(model_info: Mapping[str, object]) -> int | None:
"""``ptu_count`` as the whole number of reserved units within bounds, else None."""
raw: Final = model_info.get("ptu_count")
if isinstance(raw, bool) or not isinstance(raw, (int, float, str)):
return None
if isinstance(raw, float) and not raw.is_integer():
return None
try:
count: Final = int(raw)
except (ValueError, OverflowError):
return None
return count if 0 < count <= ModelInfo.MAX_PTU_COUNT else None
def _declared_shares(model_info: Mapping[str, object], ptu_count: int) -> Mapping[str, int] | None:
"""Who holds the capacity: the single ``team_id`` holding all of it, or the ``ptu_shares``
that add up to it, else None."""
team_id: Final = model_info.get("team_id")
raw_shares: Final = model_info.get("ptu_shares")
if team_id and raw_shares is None:
return MappingProxyType({str(team_id): ptu_count})
if team_id or raw_shares is None:
return None
shares: Final = parsed_ptu_shares(raw_shares)
if shares is None or sum(shares.values()) != ptu_count:
return None
return shares
def declares_ptu(model_info: Mapping[str, object]) -> bool:
@ -135,6 +229,11 @@ def declares_ptu(model_info: Mapping[str, object]) -> bool:
return any(model_info.get(field) is not None for field in PTU_MODEL_INFO_FIELDS)
def declares_ptu_shares(model_info: Mapping[str, object]) -> bool:
"""Whether ``ptu_shares`` is set here, including a split too malformed to honour."""
return model_info.get("ptu_shares") is not None
def ptu_config_error(model_info: Mapping[str, object], *, model_name: str | None = None) -> str | None:
"""Why this PTU configuration cannot be honoured, else None.
@ -153,8 +252,11 @@ def ptu_config_error(model_info: Mapping[str, object], *, model_name: str | None
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:
has_shares: Final = declares_ptu_shares(model_info)
if not has_count and not has_rate and not has_shares:
return None
if not has_count and not has_rate:
return _named("ptu_count and cost_per_ptu_per_hour are required when ptu_shares is set", model_name)
if has_count != has_rate:
return _named("ptu_count and cost_per_ptu_per_hour must be set together", model_name)
if effective_from is None:
@ -164,8 +266,36 @@ def ptu_config_error(model_info: Mapping[str, object], *, model_name: str | None
"today could be billed for days it did not exist",
model_name,
)
if not model_info.get("team_id"):
return _named("team_id is required when PTU fields are set (one model maps to one team)", model_name)
return _ptu_holder_error(model_info, model_name)
def _ptu_holder_error(model_info: Mapping[str, object], model_name: str | None) -> str | None:
"""Why the teams this capacity is attributed to cannot be read, else None.
The shares have to add up to the count exactly: a shortfall would leave PTU-hours the
provider bills attributed to nobody, and a surplus would attribute capacity that was
never reserved.
"""
team_id: Final = model_info.get("team_id")
raw_shares: Final = model_info.get("ptu_shares")
if team_id and raw_shares is not None:
return _named(
"team_id and ptu_shares cannot both be set; ptu_shares lists every team the capacity is split across",
model_name,
)
if not team_id and raw_shares is None:
return _named("team_id or ptu_shares is required when PTU fields are set", model_name)
if raw_shares is None:
return None
shares: Final = parsed_ptu_shares(raw_shares)
if shares is None:
return _named("ptu_shares must map at least one team_id to a positive whole number of PTUs", model_name)
ptu_count: Final = _parsed_ptu_count(model_info)
if ptu_count is None:
return None
allocated: Final = sum(shares.values())
if allocated != ptu_count:
return _named(f"ptu_shares must add up to ptu_count ({allocated} of {ptu_count} allocated)", model_name)
return None
@ -176,17 +306,13 @@ def ptu_terms(model_info: Mapping[str, object]) -> PTUTerms | None:
present but unparseable bound would read as no bound and widen the window to the whole
day, so either one leaves the deployment unpriced until the config is fixed.
"""
ptu_count: Final = model_info.get("ptu_count")
ptu_count_int: Final = _parsed_ptu_count(model_info)
cost_per_hour: Final = model_info.get("cost_per_ptu_per_hour")
team_id: Final = model_info.get("team_id")
if ptu_count is None or cost_per_hour is None or not team_id:
if ptu_count_int is None or isinstance(cost_per_hour, bool) or not isinstance(cost_per_hour, (int, float, str)):
return None
try:
ptu_count_int: Final = int(ptu_count)
cost_per_hour_float: Final = float(cost_per_hour)
except (TypeError, ValueError, OverflowError):
return None
if not 0 < ptu_count_int <= ModelInfo.MAX_PTU_COUNT:
except (ValueError, OverflowError):
return None
if not 0 <= cost_per_hour_float <= ModelInfo.MAX_COST_PER_PTU_PER_HOUR:
return None
@ -199,8 +325,11 @@ def ptu_terms(model_info: Mapping[str, object]) -> PTUTerms | None:
return None
if effective_to is not None and effective_to <= effective_from:
return None
shares: Final = _declared_shares(model_info, ptu_count_int)
if shares is None:
return None
return PTUTerms(
team_id=str(team_id),
shares=shares,
ptu_count=ptu_count_int,
cost_per_ptu_per_hour=cost_per_hour_float,
effective_from=effective_from,

View file

@ -0,0 +1,139 @@
"""How much throughput one Azure OpenAI provisioned throughput unit (PTU) serves per model.
Azure sizes a provisioned deployment in normalized tokens per minute:
``input TPM x (1 - cache hit rate) + output-to-input ratio x output TPM``, divided by the
model's "Input TPM per PTU" to get the PTUs required. The same two numbers turn a team's
PTU share into a per-minute token ceiling and a request's usage back into PTU-hours.
Table read from
https://learn.microsoft.com/en-us/azure/foundry/openai/how-to/provisioned-throughput-sizing#deployment-parameters-and-throughput-values-by-model
on 2026-09-24 (page dated 2026-09-23). Azure deducts cached input tokens in full for every
model except the GPT-6 family, where a cached input token costs a tenth of an uncached one.
"""
import re
from collections.abc import Mapping
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final, Protocol
from typing_extensions import TypeIs # noqa: TID251 # TypeIs reaches typing only on 3.13
from litellm.types.utils import LlmProviders
class NormalizedTokenWeights(Protocol):
@property
def output_to_input_ratio(self) -> float: ...
@property
def cached_input_ratio(self) -> float: ...
@dataclass(frozen=True, slots=True)
class PTUCapacity:
input_tpm_per_ptu: int
output_to_input_ratio: float
cached_input_ratio: float = 0.0
@property
def normalized_tokens_per_ptu_hour(self) -> int:
return self.input_tpm_per_ptu * 60
def input_tpm_for(self, ptus: int) -> int:
return ptus * self.input_tpm_per_ptu
AZURE_PTU_CAPACITY: Final[Mapping[str, PTUCapacity]] = MappingProxyType(
{
"gpt-6-sol": PTUCapacity(3_000, 5.0, cached_input_ratio=0.1),
"gpt-6-astra": PTUCapacity(600, 5.0, cached_input_ratio=0.1),
"gpt-5.6-luna": PTUCapacity(30_000, 6.0),
"gpt-5.6-terra": PTUCapacity(3_000, 6.0),
"gpt-5.6-sol": PTUCapacity(1_200, 6.0),
"gpt-5.5": PTUCapacity(1_200, 6.0),
"gpt-5.4": PTUCapacity(2_400, 6.0),
"gpt-5.4-mini": PTUCapacity(7_900, 6.0),
"gpt-5.3-codex": PTUCapacity(3_400, 8.0),
"gpt-5.2": PTUCapacity(3_400, 8.0),
"gpt-5.2-codex": PTUCapacity(3_400, 8.0),
"gpt-5.1": PTUCapacity(4_750, 8.0),
"gpt-5.1-codex": PTUCapacity(4_750, 8.0),
"gpt-5": PTUCapacity(4_750, 8.0),
"gpt-5-mini": PTUCapacity(23_750, 8.0),
"gpt-4.1": PTUCapacity(3_000, 4.0),
"gpt-4.1-mini": PTUCapacity(14_900, 4.0),
"gpt-4.1-nano": PTUCapacity(59_400, 4.0),
"o3": PTUCapacity(3_000, 4.0),
"o4-mini": PTUCapacity(5_400, 4.0),
"gpt-4o": PTUCapacity(2_500, 4.0),
"gpt-4o-mini": PTUCapacity(37_000, 4.0),
"o3-mini": PTUCapacity(2_500, 4.0),
"o1": PTUCapacity(230, 4.0),
"llama-3.3-70b-instruct": PTUCapacity(8_450, 4.0),
}
)
_VERSION_SUFFIX: Final = re.compile(r"-\d{4}-\d{2}-\d{2}$")
_AZURE_PROVIDERS: Final = frozenset({LlmProviders.AZURE.value, LlmProviders.AZURE_AI.value})
def azure_ptu_capacity(model: str) -> PTUCapacity | None:
"""The sizing row for ``model``, read as its last path segment with a dated version dropped.
``azure/gpt-4.1-2025-04-14`` and ``gpt-4.1`` both resolve to the ``gpt-4.1`` row; a
deployment name that is not a model name resolves to nothing, which is why callers
prefer ``model_info.base_model``.
"""
name: Final = model.rsplit("/", 1)[-1].strip().lower()
return AZURE_PTU_CAPACITY.get(name) or AZURE_PTU_CAPACITY.get(_VERSION_SUFFIX.sub("", name))
def is_str_keyed_mapping(
value: object,
) -> TypeIs[Mapping[str, object]]: # guard-ok: model_info and litellm_params are str-keyed config objects
return isinstance(value, Mapping)
def deployment_ptu_capacity(deployment: Mapping[str, object]) -> PTUCapacity | None:
"""The sizing row a deployment resolves to: ``model_info.base_model`` first, since an
Azure deployment name is arbitrary, then ``litellm_params.model``."""
model_info: Final = deployment.get("model_info")
litellm_params: Final = deployment.get("litellm_params")
candidates: Final = tuple(
value
for value in (
model_info.get("base_model") if is_str_keyed_mapping(model_info) else None,
litellm_params.get("model") if is_str_keyed_mapping(litellm_params) else None,
)
if isinstance(value, str) and value
)
return next((capacity for capacity in map(azure_ptu_capacity, candidates) if capacity is not None), None)
def is_azure_deployment(deployment: Mapping[str, object]) -> bool:
"""Whether ``litellm_params`` route this deployment to Azure OpenAI or Azure AI, by
``custom_llm_provider`` first and the ``model`` prefix otherwise."""
litellm_params: Final = deployment.get("litellm_params")
if not is_str_keyed_mapping(litellm_params):
return False
provider: Final = litellm_params.get("custom_llm_provider")
if isinstance(provider, str):
return provider in _AZURE_PROVIDERS
model: Final = litellm_params.get("model")
return isinstance(model, str) and model.partition("/")[0] in _AZURE_PROVIDERS
def normalized_tokens(
weights: NormalizedTokenWeights, *, prompt_tokens: int, completion_tokens: int, cache_read_tokens: int = 0
) -> float:
"""Azure's normalized token count for one request: uncached input in full, cached input
at the model's cached ratio, output weighted by the output-to-input ratio."""
cached: Final = min(max(cache_read_tokens, 0), max(prompt_tokens, 0))
uncached: Final = max(prompt_tokens, 0) - cached
return uncached + weights.cached_input_ratio * cached + weights.output_to_input_ratio * max(completion_tokens, 0)
def ptu_hours(capacity: PTUCapacity, normalized: float) -> float:
"""PTU-hours ``normalized`` tokens amount to: one PTU serves its input TPM for sixty minutes."""
return normalized / capacity.normalized_tokens_per_ptu_hour

View file

@ -3703,6 +3703,11 @@
"title": "Total Prompt Tokens",
"type": "integer"
},
"total_ptu_hours": {
"default": 0.0,
"title": "Total Ptu Hours",
"type": "number"
},
"total_response_time_ms": {
"default": 0,
"title": "Total Response Time Ms",
@ -4511,6 +4516,11 @@
"title": "Prompt Tokens",
"type": "integer"
},
"ptu_hours": {
"default": 0.0,
"title": "Ptu Hours",
"type": "number"
},
"spend": {
"default": 0.0,
"title": "Spend",

View file

@ -28,7 +28,7 @@ from typing import (
from fastapi import HTTPException
from pydantic import TypeAdapter, ValidationError
from starlette.status import HTTP_503_SERVICE_UNAVAILABLE
from starlette.status import HTTP_400_BAD_REQUEST, HTTP_503_SERVICE_UNAVAILABLE
from typing_extensions import NotRequired, ReadOnly
from litellm import DualCache
@ -46,7 +46,9 @@ from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.prompt_templates.common_utils import (
get_str_from_messages,
)
from litellm.litellm_core_utils.ptu_pricing import is_ptu_cost_attribution_enabled
from litellm.litellm_core_utils.token_counter import offload_token_count
from litellm.llms.azure.ptu_capacity import normalized_tokens
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.auth_utils import (
ESTIMATED_OUTPUT_TOKENS_FIELD,
@ -74,6 +76,7 @@ from litellm.router_utils.add_retry_fallback_headers import (
response_has_hidden_params,
)
from litellm.router_utils.common_utils import resolve_model_group_alias
from litellm.router_utils.ptu_shares import PTUTeamCeiling
from litellm.types.caching import RedisPipelineIncrementOperation
from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject, ResponseAPIUsage
from litellm.types.mcp_server.mcp_server_manager import MCPServer
@ -104,6 +107,17 @@ else:
_REQUEST_RATE_LIMIT_DATA: Final = TypeAdapter(Mapping[str, object])
@dataclass(frozen=True, slots=True)
class _ReconciledUsage:
prompt_tokens: int
completion_tokens: int
cached_tokens: int
@property
def billable_input_tokens(self) -> int:
return max(0, self.prompt_tokens - self.cached_tokens)
@dataclass(frozen=True, slots=True)
class RateLimitedModel:
requested: str
@ -124,6 +138,15 @@ def _resolve_model_group_alias_via_proxy_router(model: str) -> str | None:
return resolve_model_group_alias(llm_router.model_group_alias, model)
def _resolve_ptu_team_ceiling_via_proxy_router(team_id: str, model_group: str) -> PTUTeamCeiling | None:
from litellm.proxy.proxy_server import llm_router
if llm_router is None:
return None
ceiling: Final = llm_router.team_ptu_ceiling(team_id, model_group)
return ceiling if ceiling is not None and is_ptu_cost_attribution_enabled() else None
FAIL_CLOSED_RATE_LIMIT_ENFORCEMENT_SETTING: Final = "fail_closed_rate_limit_enforcement"
RATE_LIMIT_UNVERIFIABLE_MESSAGE: Final = (
"Rate limit enforcement unavailable: request counters could not be verified against Redis, and "
@ -140,6 +163,18 @@ class RateLimitUnverifiableError(HTTPException):
)
class PTUShareExceededError(HTTPException):
def __init__(self, ceiling: PTUTeamCeiling, needed_tokens: int) -> None:
super().__init__(
status_code=HTTP_400_BAD_REQUEST,
detail={
"error": f"This request needs {needed_tokens} normalized tokens, more than your team's PTU share "
f"of {ceiling.model_group} allows in a minute ({ceiling.tpm_limit}), so it can never be admitted. "
"Lower max_tokens or shorten the prompt."
},
)
_FAIL_CLOSED_RATE_LIMIT_ENFORCEMENT_FLAG: Final = TypeAdapter(bool | None)
@ -478,6 +513,7 @@ _AUDIO_BYTES_PER_TOKEN: Final = 1600
# on the same project+model simultaneously without colliding on cache keys.
PROJECT_ITPM_DESCRIPTOR_KEY: Final = "model_per_project_itpm"
PROJECT_OTPM_DESCRIPTOR_KEY: Final = "model_per_project_otpm"
PTU_TEAM_DESCRIPTOR_KEY: Final = "model_per_team_ptu"
# How long an acquired slot counts toward the in-flight total before it is
# considered leaked (worker crashed without any release callback firing) and
# pruned. Also the longest request duration the gauge can track: a request
@ -658,6 +694,9 @@ class RequestRateLimiterStash:
reserved_tokens: int = 0
reserved_model: RateLimitedModel | None = None
reserved_scopes: frozenset[tuple[str, str]] = field(default_factory=frozenset)
ptu_ceiling: PTUTeamCeiling | None = None
ptu_reserved_tokens: int = 0
ptu_reserved_window_identities: frozenset[ReservationWindowIdentity] = field(default_factory=frozenset)
itpm_reserved_tokens: int = 0
itpm_reserved_scopes: frozenset[tuple[str, str]] = field(default_factory=frozenset)
itpm_reserved_window_identities: frozenset[tuple[str, str, Literal["redis", "local"]]] = field(
@ -756,6 +795,21 @@ def get_request_stash_for_call(litellm_call_id: str | None) -> RequestRateLimite
return stash if litellm_call_id == stash.owner_litellm_call_id else None
def _as_window_guarded(operation: RedisPipelineIncrementOperation) -> ReservationAwareIncrementOperation | None:
window_key: Final = operation.get("window_key")
expected_window_start: Final = operation.get("expected_window_start")
if not isinstance(window_key, str) or not isinstance(expected_window_start, str):
return None
return ReservationAwareIncrementOperation(
key=operation["key"],
increment_value=operation["increment_value"],
ttl=operation["ttl"],
window_key=window_key,
expected_window_start=expected_window_start,
reservation_backend="local" if operation.get("reservation_backend") == "local" else "redis",
)
def _call_id_from_callback_kwargs(kwargs: object) -> str | None:
if not isinstance(kwargs, dict):
return None
@ -788,6 +842,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
time_provider: Callable[[], datetime] | None = None,
tag_rate_limit_resolver: TagRateLimitResolver = resolve_tag_rate_limits_from_db,
model_group_resolver: Callable[[str], str | None] = _resolve_model_group_alias_via_proxy_router,
ptu_team_ceiling_resolver: Callable[
[str, str], PTUTeamCeiling | None
] = _resolve_ptu_team_ceiling_via_proxy_router,
fail_closed_resolver: Callable[[], bool] = _fail_closed_rate_limit_enforcement_from_general_settings,
force_hash_tag_grouping_resolver: Callable[[], bool] = _force_redis_hash_tag_grouping_from_litellm_settings,
):
@ -795,6 +852,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
self._time_provider = time_provider or datetime.now
self._tag_rate_limit_resolver = tag_rate_limit_resolver
self._model_group_resolver = model_group_resolver
self._ptu_team_ceiling_resolver = ptu_team_ceiling_resolver
self._fail_closed_resolver = fail_closed_resolver
self._force_hash_tag_grouping_resolver = force_hash_tag_grouping_resolver
if self.internal_usage_cache.dual_cache.redis_cache is not None:
@ -1097,6 +1155,30 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
return total_estimated
def _estimate_ptu_tokens_for_request(
self,
ceiling: PTUTeamCeiling | None,
data: Mapping[str, object],
min_configured_tpm_limit: int | None,
call_type: str | None,
configured_output_tokens: int | None,
raw_estimate: int,
) -> int:
"""The team PTU ceiling counts Azure normalized tokens, so its reservation weighs the
output budget the way the ceiling does instead of the raw sum the other scopes reserve."""
if ceiling is None:
return raw_estimate
estimated_input_tokens, max_tokens_estimate = self._estimate_input_and_output_tokens(
data=data,
min_configured_tpm_limit=min_configured_tpm_limit,
call_type=call_type,
configured_output_tokens=configured_output_tokens,
)
normalized: Final = normalized_tokens(
ceiling, prompt_tokens=estimated_input_tokens, completion_tokens=max_tokens_estimate
)
return max(round(normalized), 1)
def _estimate_input_and_output_tokens(
self,
data: object,
@ -2567,11 +2649,14 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
descriptors: list[RateLimitDescriptor],
estimated_tokens: int,
parent_otel_span: Span | None = None,
scope_estimates: Mapping[str, int] = MappingProxyType({}),
) -> RateLimitResponse:
"""
Reserve ``estimated_tokens`` against every TPM-bearing descriptor
BEFORE the upstream call, so concurrent requests cannot all observe
"under limit" before any of them increments the counter.
``scope_estimates`` replaces that amount per descriptor key for a
scope counted in other units, such as the team PTU ceiling.
Thin wrapper around ``atomic_check_and_increment_by_n``: builds a
TPM-only descriptor/increment list and delegates the all-or-nothing
@ -2598,7 +2683,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
return RateLimitResponse(overall_code="OK", statuses=[])
increments: Final[list[dict[Literal["requests", "tokens"], int]]] = [
{"tokens": estimated_tokens} for _ in tpm_descriptors
{"tokens": scope_estimates.get(d["key"], estimated_tokens)} for d in tpm_descriptors
]
return await self.atomic_check_and_increment_by_n(
descriptors=tpm_descriptors,
@ -3318,6 +3403,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
requested_model=requested_model if isinstance(requested_model, str) else None,
descriptors=descriptors,
)
self._add_team_ptu_rate_limit_descriptor(
user_api_key_dict=user_api_key_dict,
requested_model=requested_model if isinstance(requested_model, str) else None,
descriptors=descriptors,
)
resolved_agent_id: Final = self._get_resolved_agent_id(user_api_key_dict, data)
for agent_id in dict.fromkeys((resolved_agent_id, user_api_key_dict.invoked_agent_id)):
@ -3463,6 +3553,31 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
)
def _add_team_ptu_rate_limit_descriptor(
self,
user_api_key_dict: UserAPIKeyAuth,
requested_model: str | None,
descriptors: _RateLimitDescriptorSink,
) -> None:
model: Final = self._rate_limited_model(requested_model)
if model is None or not user_api_key_dict.team_id:
return
ceiling: Final = self._ptu_team_ceiling_resolver(user_api_key_dict.team_id, model.group)
if ceiling is None:
return
stash: Final = get_request_stash()
if stash is not None:
stash.ptu_ceiling = ceiling
descriptors.append(
RateLimitDescriptor(
key=PTU_TEAM_DESCRIPTOR_KEY,
value=f"{user_api_key_dict.team_id}:{ceiling.model_group}",
rate_limit=RateLimitDescriptorRateLimitObject(
requests_per_unit=None, tokens_per_unit=ceiling.tpm_limit, window_size=self.window_size
),
)
)
def _add_project_model_rate_limit_descriptor_from_metadata(
self,
user_api_key_dict: UserAPIKeyAuth,
@ -3876,10 +3991,16 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
# (still-stashed) reservation and refunds it again.
if tpm_reservation_amount > 0:
await self._refund_reserved_tokens(
scopes=tpm_reservation_scopes,
scopes=tuple(scope for scope in tpm_reservation_scopes if scope[0] != PTU_TEAM_DESCRIPTOR_KEY),
amount=tpm_reservation_amount,
parent_otel_span=user_api_key_dict.parent_otel_span,
)
await self._refund_reserved_tokens(
scopes=tuple(scope for scope in tpm_reservation_scopes if scope[0] == PTU_TEAM_DESCRIPTOR_KEY),
amount=stash.ptu_reserved_tokens,
reservation_windows=stash.ptu_reserved_window_identities,
parent_otel_span=user_api_key_dict.parent_otel_span,
)
stash.reservation_released = True
await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span)
self._handle_rate_limit_error(
@ -4130,13 +4251,17 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
# in-memory check otherwise — single-worker protection still holds
# even without Redis.
# ----------------------------------------------------------------
configured_tpm_limits: Final = [
int(v)
ptu_raw_output_limit: Final = (
stash.ptu_ceiling.raw_output_token_limit() if stash.ptu_ceiling is not None else None
)
configured_tpm_limits: Final = tuple(
ptu_raw_output_limit
if d["key"] == PTU_TEAM_DESCRIPTOR_KEY and ptu_raw_output_limit is not None
else int(v)
for d in descriptors
if d["key"] not in (PROJECT_ITPM_DESCRIPTOR_KEY, PROJECT_OTPM_DESCRIPTOR_KEY)
for v in [(d.get("rate_limit") or {}).get("tokens_per_unit")]
if v is not None
]
and (v := (d.get("rate_limit") or {}).get("tokens_per_unit")) is not None
)
has_tpm_limits: Final = bool(configured_tpm_limits)
# Populated on a successful combined-TPM reservation below, so the
@ -4198,10 +4323,23 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
min_configured_tpm_limit,
)
capped_request: Final = _REQUEST_RATE_LIMIT_DATA.validate_python(data)
ptu_estimated_tokens: Final = self._estimate_ptu_tokens_for_request(
ceiling=stash.ptu_ceiling,
data=capped_request,
min_configured_tpm_limit=min_configured_tpm_limit,
call_type=call_type,
configured_output_tokens=configured_output_tokens,
raw_estimate=estimated_tokens,
)
if stash.ptu_ceiling is not None and ptu_estimated_tokens > stash.ptu_ceiling.tpm_limit:
await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span)
raise PTUShareExceededError(stash.ptu_ceiling, ptu_estimated_tokens)
tpm_response: Final = await self.reserve_tpm_tokens(
descriptors=descriptors,
estimated_tokens=estimated_tokens,
parent_otel_span=user_api_key_dict.parent_otel_span,
scope_estimates=MappingProxyType({PTU_TEAM_DESCRIPTOR_KEY: ptu_estimated_tokens}),
)
if tpm_response["overall_code"] == "OVER_LIMIT":
@ -4217,6 +4355,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
# the (actual - reserved) delta to those — unreserved
# scopes get charged the full actual usage instead.
stash.reserved_tokens = estimated_tokens
stash.ptu_reserved_tokens = ptu_estimated_tokens
stash.ptu_reserved_window_identities = frozenset(
identity
for identity in tpm_response.get("reservation_windows", frozenset())
if identity[0].startswith(f"{{{PTU_TEAM_DESCRIPTOR_KEY}:")
)
stash.reserved_model = self._rate_limited_model(requested_model)
stash.reserved_scopes = frozenset(
(d["key"], d["value"])
@ -4644,20 +4788,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
return max(0, total_tokens), 0, True
return None
def _resolve_io_token_reconcile_usage(
self,
response_obj: object,
) -> tuple[int, int, bool]:
"""
Resolve ``(billable_input_tokens, completion_tokens, usage_resolved)``
for ITPM/OTPM reconciliation. Cache-read tokens are excluded from
billable input -- Bedrock Mantle doesn't count them toward ITPM --
but they're untouched everywhere else (cost/usage logging still sees
the full prompt token count).
"""
def _resolve_reconciled_usage(self, response_obj: object) -> _ReconciledUsage | None:
"""The prompt, completion, and cache-read token counts a response reports, else None
when it reports no usage at all. Cache-read tokens stay inside ``prompt_tokens`` here;
each consumer decides what they cost it."""
rerank_usage: Final = self._resolve_rerank_token_usage(response_obj)
if rerank_usage is not None:
return rerank_usage
return _ReconciledUsage(prompt_tokens=rerank_usage[0], completion_tokens=rerank_usage[1], cached_tokens=0)
usage: Final = self._response_usage(response_obj)
@ -4670,8 +4807,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
else 0
)
if prompt_tokens == 0 and completion_tokens == 0:
return 0, 0, False
return max(0, prompt_tokens - cached_tokens), completion_tokens, True
return None
return _ReconciledUsage(
prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, cached_tokens=cached_tokens
)
if isinstance(usage, ResponseAPIUsage):
response_input_tokens: Final = usage.input_tokens or 0
@ -4680,8 +4819,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
usage.input_tokens_details.cached_tokens or 0 if usage.input_tokens_details is not None else 0
)
if response_input_tokens == 0 and response_output_tokens == 0:
return 0, 0, False
return max(0, response_input_tokens - response_cached_tokens), response_output_tokens, True
return None
return _ReconciledUsage(
prompt_tokens=response_input_tokens,
completion_tokens=response_output_tokens,
cached_tokens=response_cached_tokens,
)
if isinstance(usage, Mapping):
raw_prompt_tokens: Final = usage.get("prompt_tokens") or usage.get("input_tokens") or 0
@ -4696,10 +4839,30 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
mapped_cached_tokens: Final = raw_cached_tokens if isinstance(raw_cached_tokens, int) else 0
if mapped_prompt_tokens == 0 and mapped_completion_tokens == 0:
return 0, 0, False
return max(0, mapped_prompt_tokens - mapped_cached_tokens), mapped_completion_tokens, True
return None
return _ReconciledUsage(
prompt_tokens=mapped_prompt_tokens,
completion_tokens=mapped_completion_tokens,
cached_tokens=mapped_cached_tokens,
)
return 0, 0, False
return None
def _resolve_io_token_reconcile_usage(
self,
response_obj: object,
) -> tuple[int, int, bool]:
"""
Resolve ``(billable_input_tokens, completion_tokens, usage_resolved)``
for ITPM/OTPM reconciliation. Cache-read tokens are excluded from
billable input -- Bedrock Mantle doesn't count them toward ITPM --
but they're untouched everywhere else (cost/usage logging still sees
the full prompt token count).
"""
usage: Final = self._resolve_reconciled_usage(response_obj)
if usage is None:
return 0, 0, False
return usage.billable_input_tokens, usage.completion_tokens, True
def _build_io_token_reservation_ops(
self,
@ -4973,7 +5136,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
# than parsed out of the body) carry their usage in
# ``combined_usage_object`` instead, and would otherwise never charge
# the TPM window.
_usage: Usage | dict | None = None
_usage: Usage | dict[str, object] | None = None
if isinstance(
response_obj,
(
@ -5047,9 +5210,84 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
reserved_tokens=reserved_tokens,
)
)
pipeline_operations.extend(
self._build_team_ptu_tpm_ops(
standard_logging_metadata, # pyright: ignore[reportUnknownArgumentType] # untyped logging metadata
usage_source=response_obj if _usage is None else _usage,
reconcile_model=reconcile_model,
reserved_scopes=reserved_scopes,
reserved_ceiling=stash.ptu_ceiling if stash is not None else None,
reserved_tokens=stash.ptu_reserved_tokens if stash is not None else 0,
reserved_windows=stash.ptu_reserved_window_identities if stash is not None else frozenset(),
total_tokens=total_tokens,
)
)
return pipeline_operations
def _build_team_ptu_tpm_ops(
self,
standard_logging_metadata: Mapping[str, object],
usage_source: object,
reconcile_model: RateLimitedModel | None,
reserved_scopes: Set[tuple[str, str]],
reserved_ceiling: PTUTeamCeiling | None,
reserved_tokens: int,
reserved_windows: frozenset[ReservationWindowIdentity],
total_tokens: int,
) -> Sequence[RedisPipelineIncrementOperation]:
"""Settle the team's PTU counter in Azure normalized tokens: uncached input in full,
cached input at the model's cached ratio, output weighted by its output-to-input
ratio, the way Azure sizes a PTU.
The reservation was taken in those units against the ceiling admission resolved, so that
ceiling settles it even after the share changed or went away or the logging metadata lost
the team; when usage cannot be resolved it charges the raw total the other scopes charge
so the reservation is never left standing.
"""
reserved_ptu_scopes: Final = tuple(scope for scope in reserved_scopes if scope[0] == PTU_TEAM_DESCRIPTOR_KEY)
if reserved_ptu_scopes:
return self._build_project_reservation_ops(
targets=reserved_ptu_scopes,
reserved_scopes=reserved_scopes,
actual_tokens=self._ptu_settlement_tokens(
reserved_ceiling, self._resolve_reconciled_usage(usage_source), total_tokens
),
reserved_tokens=reserved_tokens,
reservation_window_identities=reserved_windows,
)
team_id: Final = standard_logging_metadata.get("user_api_key_team_id")
if reconcile_model is None or not isinstance(team_id, str) or not team_id:
return ()
ceiling: Final = (
reserved_ceiling
if reserved_ceiling is not None
else self._ptu_team_ceiling_resolver(team_id, reconcile_model.group)
)
if ceiling is None:
return ()
return self._build_reservation_aware_tpm_ops(
targets=((PTU_TEAM_DESCRIPTOR_KEY, f"{team_id}:{ceiling.model_group}"),),
reserved_scopes=reserved_scopes,
actual_tokens=self._ptu_settlement_tokens(
ceiling, self._resolve_reconciled_usage(usage_source), total_tokens
),
reserved_tokens=reserved_tokens,
)
@staticmethod
def _ptu_settlement_tokens(ceiling: PTUTeamCeiling | None, usage: _ReconciledUsage | None, raw_total: int) -> int:
if ceiling is None or usage is None:
return raw_total
return round(
normalized_tokens(
ceiling,
prompt_tokens=usage.prompt_tokens,
completion_tokens=usage.completion_tokens,
cache_read_tokens=usage.cached_tokens,
)
)
@with_service_target("rate_limits")
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None:
"""
@ -5073,9 +5311,20 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
response_obj=response_obj,
rate_limit_type=rate_limit_type,
)
if pipeline_operations:
unguarded_operations: Final = [
operation for operation in pipeline_operations if _as_window_guarded(operation) is None
]
window_guarded_operations: Final = tuple(
guarded for operation in pipeline_operations if (guarded := _as_window_guarded(operation)) is not None
)
if unguarded_operations:
await self.async_increment_tokens_with_ttl_preservation(
pipeline_operations=pipeline_operations,
pipeline_operations=unguarded_operations,
parent_otel_span=litellm_parent_otel_span,
)
if window_guarded_operations:
await self.async_increment_reservation_aware_tokens(
pipeline_operations=window_guarded_operations,
parent_otel_span=litellm_parent_otel_span,
)
io_token_operations: Final = self._build_io_token_reservation_ops(
@ -5157,9 +5406,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
statuses=statuses,
)
def _recovered_partial_usage_tokens(self, source: Mapping[str, object]) -> tuple[int, int, int]:
def _recovered_partial_usage(self, source: Mapping[str, object]) -> Usage | None:
usage: Final = source.get("combined_usage_object")
if not isinstance(usage, Usage) or (usage.completion_tokens or 0) <= 0:
return None
return usage
def _recovered_partial_usage_tokens(self, source: Mapping[str, object]) -> tuple[int, int, int]:
usage: Final = self._recovered_partial_usage(source)
if usage is None:
return 0, 0, 0
billable_input, completion_tokens, _ = self._resolve_io_token_reconcile_usage(usage)
return (
@ -5168,6 +5423,27 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
completion_tokens,
)
def _build_ptu_failure_settlement_ops(
self, stash: RequestRateLimiterStash, source: Mapping[str, object], raw_actual_tokens: int
) -> tuple[ReservationAwareIncrementOperation, ...]:
"""Settle the team PTU reservation on failure in the normalized tokens it was taken in:
at the recovered partial usage when there is one, else a full refund."""
ptu_scopes: Final = tuple(scope for scope in stash.reserved_scopes if scope[0] == PTU_TEAM_DESCRIPTOR_KEY)
if not ptu_scopes:
return ()
usage: Final = self._recovered_partial_usage(source)
return self._build_project_reservation_ops(
targets=ptu_scopes,
reserved_scopes=stash.reserved_scopes,
actual_tokens=self._ptu_settlement_tokens(
stash.ptu_ceiling,
self._resolve_reconciled_usage(usage) if usage is not None else None,
raw_actual_tokens,
),
reserved_tokens=stash.ptu_reserved_tokens,
reservation_window_identities=stash.ptu_reserved_window_identities,
)
@with_service_target("rate_limits")
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None:
"""
@ -5210,7 +5486,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
# refund there would drive their counter negative.
pipeline_operations.extend(
self._build_reservation_aware_tpm_ops(
targets=list(stash.reserved_scopes),
targets=tuple(scope for scope in stash.reserved_scopes if scope[0] != PTU_TEAM_DESCRIPTOR_KEY),
reserved_scopes=stash.reserved_scopes,
actual_tokens=tpm_actual,
reserved_tokens=reserved_tokens,
@ -5257,11 +5533,21 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
else ()
)
ptu_operations: Final = (
self._build_ptu_failure_settlement_ops(
stash,
kwargs, # pyright: ignore[reportUnknownArgumentType] # hook kwargs are unannotated
tpm_actual,
)
if stash is not None and reserved_tokens > 0
else ()
)
if pipeline_operations:
await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline_post_call(
pipeline_operations, parent_otel_span=litellm_parent_otel_span
)
for project_operations in (itpm_operations, otpm_operations):
for project_operations in (itpm_operations, otpm_operations, ptu_operations):
if isinstance(project_operations, list):
await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline_post_call(
project_operations, parent_otel_span=litellm_parent_otel_span
@ -5408,7 +5694,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
combined_ops: Final = (
self._build_reservation_aware_tpm_ops(
targets=tuple(stash.reserved_scopes),
targets=tuple(scope for scope in stash.reserved_scopes if scope[0] != PTU_TEAM_DESCRIPTOR_KEY),
reserved_scopes=stash.reserved_scopes,
actual_tokens=tpm_actual,
reserved_tokens=reserved_tokens,
@ -5416,6 +5702,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
if reserved_tokens > 0
else ()
)
ptu_ops: Final = (
self._build_ptu_failure_settlement_ops(
stash,
request_data, # pyright: ignore[reportUnknownArgumentType] # request_data is a bare dict
tpm_actual,
)
if reserved_tokens > 0
else ()
)
itpm_ops: Final = (
self._build_project_reservation_ops(
targets=tuple(stash.itpm_reserved_scopes),
@ -5452,7 +5747,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
if otpm_reserved > 0
else ()
)
if combined_ops or itpm_ops or otpm_ops:
if combined_ops or ptu_ops or itpm_ops or otpm_ops:
verbose_proxy_logger.debug(
"Releasing reserved tokens on proxy-level rejection: tpm=%s, itpm=%s, otpm=%s",
reserved_tokens,
@ -5461,10 +5756,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
if combined_ops:
await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline(
increment_list=combined_ops,
increment_list=list(combined_ops),
litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
)
for project_ops in (itpm_ops, otpm_ops):
for project_ops in (itpm_ops, otpm_ops, ptu_ops):
if isinstance(project_ops, list):
await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline(
increment_list=project_ops,

View file

@ -34,10 +34,12 @@ from litellm.proxy.management_endpoints.daily_activity_scopes import (
EntityScopeResolver,
ResolvedScope,
)
from litellm.proxy.management_endpoints.ptu_consumption import single_team_id, with_ptu_consumption
from litellm.proxy.management_endpoints.team_endpoints import aggregated_date_range_error
from litellm.proxy.management_helpers.utils import management_endpoint_wrapper
from litellm.proxy.utils import PrismaClient, get_prisma_client_or_throw
from litellm.repositories.daily_activity_repository import DailyActivityRepository
from litellm.router import Router
from litellm.types.proxy.management_endpoints.common_daily_activity import (
CacheLeakageKeysResponse,
DailyActivityKeyPageResponse,
@ -72,6 +74,12 @@ def get_daily_activity_repository() -> DailyActivityRepository:
return daily_activity_repository(get_daily_activity_prisma_client())
def get_daily_activity_llm_router() -> Router | None:
from litellm.proxy.proxy_server import llm_router
return llm_router
def _date_range_error(query: EntityQuery, *, user_aggregated: bool) -> InvalidDateRange | None:
if user_aggregated:
date_range: Final = parse_canonical_date_range(query.start_date, query.end_date)
@ -274,6 +282,7 @@ def _register_aggregated_route(router: APIRouter, resolver: EntityScopeResolver,
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
repository: Annotated[DailyActivityRepository, Depends(get_daily_activity_repository)],
prisma_client: Annotated[PrismaClient, Depends(get_daily_activity_prisma_client)],
llm_router: Annotated[Router | None, Depends(get_daily_activity_llm_router)],
api_key_limit: Annotated[
int, Query(ge=1, le=constants.USAGE_TOP_API_KEYS_MAX)
] = constants.USAGE_TOP_API_KEYS_DEFAULT,
@ -286,13 +295,16 @@ def _register_aggregated_route(router: APIRouter, resolver: EntityScopeResolver,
prisma_client,
user_aggregated=resolver.entity == "user",
)
return await get_daily_activity_aggregated(
activity: Final = await get_daily_activity_aggregated(
repository,
resolved.scope,
entity_metadata_field=resolved.entity_metadata,
include_entity_breakdown=resolver.include_entity_breakdown,
api_key_limit=api_key_limit,
)
if resolver.entity != "team":
return activity
return with_ptu_consumption(activity, llm_router, single_team_id(resolved.scope.entity_ids))
except HTTPException:
raise
except Exception as exc:

View file

@ -17,6 +17,7 @@ from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequen
from contextlib import AbstractAsyncContextManager, asynccontextmanager, suppress
from dataclasses import dataclass
from fnmatch import fnmatchcase
from functools import partial
from json import JSONDecodeError
from types import MappingProxyType
from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeAlias, TypeVar, cast, runtime_checkable
@ -37,6 +38,7 @@ from litellm.litellm_core_utils.ptu_pricing import (
PTU_ZEROED_PRICING_FIELDS,
PTU_ZEROED_TABLE_FIELDS,
SEARCH_CONTEXT_SIZES,
parsed_ptu_shares,
ptu_config_error,
)
from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload
@ -105,6 +107,7 @@ from litellm.proxy.spend_tracking.ptu_feature_flag import (
is_ptu_cost_attribution_enabled,
)
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.repositories.chunked_in import find_many_in
from litellm.repositories.credentials_repository import CredentialsRepository
from litellm.repositories.model_repository import ModelRepository
from litellm.repositories.prisma_protocols import TableActions
@ -266,6 +269,9 @@ class _ExistingModelRow(Protocol):
class _TeamRow(Protocol):
@property
def team_id(self) -> str: ...
@property
def models(self) -> Sequence[str]: ...
@ -275,6 +281,8 @@ class _TeamRow(Protocol):
class _TeamLookupTable(Protocol):
def find_unique(self, *, where: Mapping[str, object]) -> Awaitable[_TeamRow | None]: ...
async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_TeamRow]: ...
class _TeamTable(_TeamLookupTable, Protocol):
def update(
@ -840,6 +848,21 @@ def _validate_ptu_model_info(model_info: Mapping[str, object]) -> None:
raise HTTPException(status_code=400, detail=error)
async def _raise_if_ptu_share_teams_missing(
model_info: Mapping[str, object], team_table: Callable[[], _TeamLookupTable]
) -> None:
"""Hold every team named in ``ptu_shares`` to the existence check ``team_id`` already gets."""
shares: Final = parsed_ptu_shares(model_info.get("ptu_shares"))
if shares is None:
return
rows: Final = await find_many_in(team_table(), "team_id", shares.keys())
found: Final = frozenset(row.team_id for row in rows)
missing: Final = tuple(team_id for team_id in shares if team_id not in found)
if not missing:
return
raise HTTPException(status_code=400, detail={"error": f"Team id={', '.join(missing)} does not exist in db"})
# The mirrored per-token pricing fields plus the remaining rates the public cost map or a
# provider default would otherwise supply (the cache back-fills, the Maps grounding rate). An
# unset field falls back to those sources, so a field left out here is one a PTU deployment
@ -1692,8 +1715,10 @@ 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))
incoming_model_info: Final = patch_data.model_info.model_dump(exclude_none=True)
_raise_if_ptu_cost_attribution_disabled(incoming_model_info)
_validate_ptu_model_info(_merged_ptu_model_info(db_model=db_model, patch_data=patch_data))
await _raise_if_ptu_share_teams_missing(incoming_model_info, partial(_repo_team_table, prisma_client))
_raise_if_ptu_deployment_is_priced(
model_info=_merged_ptu_model_info(db_model=db_model, patch_data=patch_data),
supplied=(
@ -2599,6 +2624,7 @@ async def add_new_model(
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)
await _raise_if_ptu_share_teams_missing(incoming_model_info, partial(_repo_team_table, prisma_client))
priced_model_params: Final = _ptu_priced_deployment(model_params)
if store_model_in_db is True:

View file

@ -60,6 +60,7 @@ from litellm.proxy.management_helpers.object_permission_utils import (
prepare_object_permission_upsert,
reject_ambiguous_mcp_tool_permission_keys,
)
from litellm.proxy.management_helpers.ptu_share_guard import raise_if_teams_hold_ptu_shares
from litellm.proxy.management_helpers.utils import (
get_new_internal_user_defaults,
management_endpoint_wrapper,
@ -1025,7 +1026,7 @@ async def delete_organization(
- organization_ids: List[str] - The organization ids to delete.
"""
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
from litellm.proxy.proxy_server import llm_router, prisma_client, proxy_logging_obj, user_api_key_cache
if prisma_client is None:
raise HTTPException(
@ -1039,6 +1040,11 @@ async def delete_organization(
detail={"error": "Only proxy admins can delete organizations"},
)
teams_in_deleted_orgs: Final = await find_many_in(
_table(TeamRepository(prisma_client)), "organization_id", data.organization_ids
)
raise_if_teams_hold_ptu_shares(tuple(team.team_id for team in teams_in_deleted_orgs), llm_router)
deleted_orgs: Final = []
for organization_id in data.organization_ids:
# delete all teams in the organization

View file

@ -0,0 +1,145 @@
"""PTU-hours a team's tokens amount to, attached to a daily activity response.
Azure sizes a provisioned deployment in normalized tokens per minute per PTU, so the prompt,
cached, and completion tokens a team sent to a PTU model group convert back to the share of
a PTU-hour it consumed, reported next to the raw token counts.
"""
from collections.abc import Callable, Sequence
from datetime import date
from functools import cache
from types import MappingProxyType
from typing import Final
from litellm.litellm_core_utils.ptu_pricing import is_ptu_cost_attribution_enabled
from litellm.llms.azure.ptu_capacity import PTUCapacity, normalized_tokens, ptu_hours
from litellm.router import Router
from litellm.router_utils.common_utils import resolve_model_group_alias
from litellm.router_utils.ptu_shares import (
model_group_ptu_capacity,
routed_deployments,
team_servable_deployments,
)
from litellm.types.proxy.management_endpoints.common_daily_activity import (
DailySpendData,
MetricWithMetadata,
SpendAnalyticsPaginatedResponse,
SpendMetrics,
)
def _with_ptu_hours(metrics: SpendMetrics, capacity: PTUCapacity) -> SpendMetrics:
consumed: Final = ptu_hours(
capacity,
normalized_tokens(
capacity,
prompt_tokens=metrics.prompt_tokens,
completion_tokens=metrics.completion_tokens,
cache_read_tokens=metrics.cache_read_input_tokens,
),
)
return metrics.model_copy(update=MappingProxyType({"ptu_hours": consumed}))
def _model_group_with_ptu_hours(bucket: MetricWithMetadata, capacity: PTUCapacity) -> MetricWithMetadata:
api_key_breakdown: Final = {
api_key: key_bucket.model_copy(
update=MappingProxyType({"metrics": _with_ptu_hours(key_bucket.metrics, capacity)})
)
for api_key, key_bucket in bucket.api_key_breakdown.items()
}
return bucket.model_copy(
update=MappingProxyType(
{"metrics": _with_ptu_hours(bucket.metrics, capacity), "api_key_breakdown": api_key_breakdown}
)
)
def _day_with_ptu_hours(
day: DailySpendData, capacity_for_model_group: Callable[[str, date], PTUCapacity | None]
) -> DailySpendData:
priced: Final = MappingProxyType(
{
model_group: _model_group_with_ptu_hours(bucket, capacity)
for model_group, bucket in day.breakdown.model_groups.items()
if (capacity := capacity_for_model_group(model_group, day.date)) is not None
}
)
if not priced:
return day
model_groups: Final = {
**day.breakdown.model_groups,
**priced,
}
return day.model_copy(
update=MappingProxyType(
{
"metrics": day.metrics.model_copy(
update=MappingProxyType({"ptu_hours": sum(bucket.metrics.ptu_hours for bucket in priced.values())})
),
"breakdown": day.breakdown.model_copy(update=MappingProxyType({"model_groups": model_groups})),
}
)
)
def attach_ptu_hours(
response: SpendAnalyticsPaginatedResponse, capacity_for_model_group: Callable[[str, date], PTUCapacity | None]
) -> SpendAnalyticsPaginatedResponse:
"""The response with ``ptu_hours`` filled in on every PTU model group and its api keys,
on each day, and on the total, from the tokens already on the page.
A model group the resolver has no sizing row for on that day keeps ``ptu_hours`` at zero.
"""
days: Final = tuple(_day_with_ptu_hours(day, capacity_for_model_group) for day in response.results)
results: Final = list(days)
return response.model_copy(
update=MappingProxyType(
{
"results": results,
"metadata": response.metadata.model_copy(
update=MappingProxyType({"total_ptu_hours": sum(day.metrics.ptu_hours for day in days)})
),
}
)
)
def capacity_by_requested_name(llm_router: Router, team_id: str | None) -> Callable[[str, date], PTUCapacity | None]:
"""The sizing row behind the name a usage row is keyed by, resolved the way the ceiling
resolves a request: an alias to its group, then a group, a routing group, a deployment id,
the wildcard route the router picks for the name, or a provider model to the deployments it
is served from. For one team those are narrowed to the deployments it can be served from,
its shared one first, so a team with no share on a name reads no PTU-hours for it; a page
spanning teams keeps every deployment behind the name. Only a reservation covering the row's
day sizes it, so tokens before a reservation starts or after it ends read no PTU-hours."""
listed_rows: Final = llm_router.get_model_list() or ()
aliases: Final = llm_router.model_group_alias
def capacity_for(requested_model: str, day: date) -> PTUCapacity | None:
model_group: Final = resolve_model_group_alias(aliases, requested_model) or requested_model
routed: Final = routed_deployments(
listed_rows,
llm_router.model_list, # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType] # Router.model_list is a bare list
model_group,
lambda name: llm_router.wildcard_route_deployments(name, team_id),
)
return model_group_ptu_capacity(routed if team_id is None else team_servable_deployments(routed, team_id), day)
return cache(capacity_for)
def single_team_id(team_ids: Sequence[str] | None) -> str | None:
"""The one team a usage page is scoped to, else None for a page spanning several or all teams."""
return team_ids[0] if team_ids is not None and len(team_ids) == 1 else None
def with_ptu_consumption(
activity: SpendAnalyticsPaginatedResponse, llm_router: Router | None, team_id: str | None
) -> SpendAnalyticsPaginatedResponse:
"""``activity`` with PTU-hours attached from the deployments ``team_id`` is served from (every
sized deployment when the page spans teams), untouched while PTU cost attribution is off or no
router is loaded."""
if llm_router is None or not is_ptu_cost_attribution_enabled():
return activity
return attach_ptu_hours(activity, capacity_by_requested_name(llm_router, team_id))

View file

@ -73,6 +73,7 @@ from litellm.proxy.management_helpers.bulk_user_creation import (
_bounded, # pyright: ignore[reportPrivateUsage] # the bounded gather /user/bulk_new runs its row writes through
)
from litellm.proxy.management_helpers.bulk_user_deletion import delete_users_from_current_teams
from litellm.proxy.management_helpers.ptu_share_guard import raise_if_teams_hold_ptu_shares
from litellm.proxy.management_helpers.team_roster_sync import (
ROSTER_TX_TIMEOUT,
MembersMissing,
@ -2818,8 +2819,11 @@ async def delete_group(
"""
verbose_proxy_logger.debug("SCIM DELETE GROUP request for group_id=%s", group_id)
try:
from litellm.proxy.proxy_server import llm_router
prisma_client: Final = await _get_prisma_client_or_raise_exception()
existing_team: Final = await _check_team_exists(group_id)
raise_if_teams_hold_ptu_shares((group_id,), llm_router)
member_ids: Final = await _get_team_member_user_ids_from_team(existing_team)

View file

@ -159,6 +159,7 @@ from litellm.proxy.management_endpoints.common_utils import ( # noqa: F401 # l
from litellm.proxy.management_endpoints.organization_endpoints import (
add_member_to_organization,
)
from litellm.proxy.management_endpoints.ptu_consumption import single_team_id, with_ptu_consumption
from litellm.proxy.management_endpoints.router_weights import validate_router_settings_weights
from litellm.proxy.management_endpoints.tag_management_endpoints import (
get_daily_activity,
@ -190,6 +191,7 @@ from litellm.proxy.management_helpers.object_permission_utils import ( # noqa:
invalidate_cached_object_permissions,
set_object_permission,
)
from litellm.proxy.management_helpers.ptu_share_guard import raise_if_teams_hold_ptu_shares
from litellm.proxy.management_helpers.team_member_permission_checks import (
TeamMemberPermissionChecks,
)
@ -4530,6 +4532,7 @@ async def delete_team(
from litellm.proxy.proxy_server import (
create_audit_log_for_update,
litellm_proxy_admin_name,
llm_router,
prisma_client,
proxy_logging_obj,
user_api_key_cache,
@ -4561,6 +4564,8 @@ async def delete_team(
team_rows.append(team_row_pydantic)
raise_if_teams_hold_ptu_shares(data.team_ids, llm_router)
await _persist_deleted_team_records(
teams=team_rows,
prisma_client=prisma_client,
@ -6873,6 +6878,7 @@ async def get_team_daily_activity(
SpendAnalyticsPaginatedResponse: Paginated response containing daily activity data.
"""
from litellm.proxy.proxy_server import (
llm_router,
prisma_client,
proxy_logging_obj,
user_api_key_cache,
@ -6891,7 +6897,7 @@ async def get_team_daily_activity(
proxy_logging_obj=proxy_logging_obj,
)
return await get_daily_activity(
activity: Final = await get_daily_activity(
prisma_client=prisma_client,
table_name="litellm_dailyteamspend",
entity_id_field="team_id",
@ -6905,6 +6911,7 @@ async def get_team_daily_activity(
page=page,
page_size=page_size,
)
return with_ptu_consumption(activity, llm_router, single_team_id(scope.team_ids))
_MAX_AGGREGATED_RANGE_DAYS: Final = 400

View file

@ -0,0 +1,30 @@
"""Refuses deleting a team that still holds a PTU share, on every route that deletes teams."""
from collections.abc import Sequence
from typing import Final
from fastapi import HTTPException
from litellm.router import Router
from litellm.router_utils.ptu_shares import deployments_shared_with
def raise_if_teams_hold_ptu_shares(team_ids: Sequence[str], llm_router: Router | None) -> None:
"""A team still named in a deployment's ``ptu_shares`` would keep being billed for PTUs nobody can use."""
if llm_router is None:
return
held: Final = tuple(
(team_id, deployment_ids)
for team_id in team_ids
if (deployment_ids := deployments_shared_with(llm_router.get_model_list() or (), team_id))
)
if not held:
return
team_id, deployment_ids = held[0]
raise HTTPException(
status_code=400,
detail={
"error": f"Team {team_id} holds a PTU share on deployment {', '.join(deployment_ids)}. "
"Move its share to another team in model_info.ptu_shares before deleting the team."
},
)

View file

@ -14,11 +14,12 @@ and share the existing unique constraint.
import asyncio
import json
from collections.abc import Awaitable, Callable, Mapping
from collections.abc import Awaitable, Callable, Iterator, Mapping
from dataclasses import dataclass
from datetime import date, datetime, time, timedelta, timezone
from itertools import chain, groupby
from types import MappingProxyType
from typing import TYPE_CHECKING, Final
from typing import TYPE_CHECKING, Final, Protocol
from litellm._internal_context import with_service_target
from litellm._logging import verbose_proxy_logger
@ -59,6 +60,27 @@ def _proxy_model_table(prisma_client: "PrismaClient") -> "TableActions[prisma_mo
return ModelRepository(prisma_client).table
class _SentinelBatchTable(Protocol):
def upsert(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> None: ...
def delete_many(self, *, where: Mapping[str, object]) -> None: ...
class _SentinelBatch(Protocol):
litellm_dailyteamspend: _SentinelBatchTable
class _SentinelBatchManager(Protocol):
async def __aenter__(self) -> _SentinelBatch: ...
async def __aexit__(self, exc_type: object, exc_value: object, traceback: object) -> bool | None: ...
def _sentinel_batch(prisma_client: "PrismaClient") -> _SentinelBatchManager:
batch: Final[_SentinelBatchManager] = prisma_client.db.batch_()
return batch
_HOURS_PER_DAY: Final = 24
_PRUNE_ID_CHUNK_SIZE: Final = 5_000
_UPSERT_ATTEMPTS: Final = 3
@ -85,7 +107,8 @@ class BackfillResult:
@dataclass(frozen=True, slots=True)
class PTUModel:
"""A model deployment carrying valid manual PTU config."""
"""One team's holding on a deployment carrying valid manual PTU config: the whole count
under a single team_id, or that team's share of a split deployment."""
model_id: str
model_name: str
@ -101,10 +124,10 @@ def _public_model_name(row: object, model_info: Mapping[str, object]) -> str:
Creating a team-scoped deployment rewrites model_name to a synthetic routing key
(``model_name_<team_id>_<uuid4>``) and keeps the chosen name in
``model_info.team_public_model_name``. PTU config is only accepted alongside a
team_id, so every PTU deployment carries that synthetic name; keying the sentinel
row on it would file each charge under a UUID that no usage view can resolve and
that never lines up with the same model's request rows.
``model_info.team_public_model_name``. A PTU deployment held by one team_id carries
that synthetic name; keying the sentinel row on it would file each charge under a
UUID that no usage view can resolve and that never lines up with the same model's
request rows. A deployment split into ptu_shares keeps the name it was declared with.
"""
public_name: Final = model_info.get("team_public_model_name")
if isinstance(public_name, str) and public_name:
@ -131,7 +154,7 @@ def _decode_model_info(raw: object) -> "Mapping[str, object] | None":
@dataclass(frozen=True, slots=True)
class _PTUDeployment:
"""A deployment in the shape ``_parse_ptu_model`` reads, whatever declared it.
"""A deployment in the shape ``_parse_ptu_models`` reads, whatever declared it.
A ``LiteLLM_ProxyModelTable`` row already has it. A router entry does not: its id
lives in ``model_info.id`` rather than on the entry itself.
@ -163,26 +186,30 @@ def _router_deployment(deployment: Mapping[str, object]) -> _PTUDeployment | Non
)
def _parse_ptu_model(row: object) -> PTUModel | None:
"""Return a PTUModel when the deployment carries valid manual PTU config, else None.
def _parse_ptu_models(row: object) -> tuple[PTUModel, ...]:
"""One PTUModel per team the deployment's valid PTU config attributes capacity to, else empty.
Valid means model_info has a positive ptu_count, a non-negative
cost_per_ptu_per_hour, and a team_id (1 model -> 1 team).
A deployment held by one team_id yields one entry carrying the whole count; a deployment
split into ptu_shares yields one per team carrying that team's share, so each team's row
accrues its share of the flat cost.
"""
model_info: Final = _decode_model_info(getattr(row, "model_info", None))
if model_info is None:
return None
return ()
terms: Final = ptu_terms(model_info)
if terms is None:
return None
return PTUModel(
model_id=str(getattr(row, "model_id", "") or ""),
model_name=_public_model_name(row, model_info),
team_id=terms.team_id,
ptu_count=terms.ptu_count,
cost_per_ptu_per_hour=terms.cost_per_ptu_per_hour,
effective_from=terms.effective_from,
effective_to=terms.effective_to,
return ()
return tuple(
PTUModel(
model_id=str(getattr(row, "model_id", "") or ""),
model_name=_public_model_name(row, model_info),
team_id=team_id,
ptu_count=share,
cost_per_ptu_per_hour=terms.cost_per_ptu_per_hour,
effective_from=terms.effective_from,
effective_to=terms.effective_to,
)
for team_id, share in sorted(terms.shares.items())
)
@ -236,104 +263,136 @@ def _aggregate_charges(ptu_models: tuple[PTUModel, ...], day: date) -> tuple[_PT
)
async def _upsert_ptu_daily_row(
prisma_client: "PrismaClient",
*,
team_id: str,
model_id: str,
model_name: str,
date_str: str,
flat_cost: float,
) -> None:
"""Idempotent upsert of a sentinel-api_key row on LiteLLM_DailyTeamSpend.
@dataclass(frozen=True, slots=True)
class _DeploymentDay:
"""Every team's charge for one deployment on one UTC day.
The backfill counts a deployment day as priced once any team's row exists for it, so
the charges are written together or not at all: a day where only some teams' rows
landed would read as priced and the missing teams would never be billed.
"""
date_str: str
charges: tuple[_PTUCharge, ...]
@property
def model_id(self) -> str:
return self.charges[0].model_id
def _deployment_days(charges: tuple[_PTUCharge, ...], date_str: str) -> tuple[_DeploymentDay, ...]:
by_deployment: Final = sorted(charges, key=lambda charge: (charge.model_id, charge.team_id))
return tuple(
_DeploymentDay(date_str=date_str, charges=tuple(group))
for _, group in groupby(by_deployment, key=lambda charge: charge.model_id)
)
def _deployment_day_rows_filter(deployment_day: _DeploymentDay) -> "Mapping[str, object]":
"""Every sentinel row of the deployment day, whichever team it bills, as a plain dict for the
same reason as ``_prune_filter``."""
return {
"date": deployment_day.date_str,
"api_key": PTU_SENTINEL_API_KEY,
"model": deployment_day.model_id,
}
def _queue_deployment_day(table: _SentinelBatchTable, deployment_day: _DeploymentDay) -> None:
"""Queue the replacement of one deployment day's sentinel-api_key rows on LiteLLM_DailyTeamSpend:
the day's rows go, then one row per charged team is upserted.
Both land in the day's one transaction, so rerunning a priced day after its shares moved
bills it to the new holders alone, whether or not the run may prune. The delete is scoped
by the day and deployment rather than by the teams no longer charged, since a share map
names any number of teams and a negated team list would be an unbounded ``NOT IN``.
``model`` holds the deployment id because it is part of the table's unique key and a
rename must not move the row. ``model_group`` carries the operator-facing name, which
is outside the key and is what the usage views display.
"""
where: Final = {
"team_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint": {
"team_id": team_id,
"date": date_str,
"api_key": PTU_SENTINEL_API_KEY,
"model": model_id,
"custom_llm_provider": "",
"mcp_namespaced_tool_name": "",
"endpoint": "",
}
}
now: Final = datetime.now(timezone.utc)
await _daily_team_spend_table(prisma_client).upsert(
where=where,
data={
"create": {
"team_id": team_id,
"date": date_str,
"api_key": PTU_SENTINEL_API_KEY,
"model": model_id,
"model_group": model_name,
"custom_llm_provider": "",
"mcp_namespaced_tool_name": "",
"endpoint": "",
"ptu_flat_cost": flat_cost,
table.delete_many(where=_deployment_day_rows_filter(deployment_day))
for charge in deployment_day.charges:
table.upsert(
where={
"team_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint": {
"team_id": charge.team_id,
"date": deployment_day.date_str,
"api_key": PTU_SENTINEL_API_KEY,
"model": charge.model_id,
"custom_llm_provider": "",
"mcp_namespaced_tool_name": "",
"endpoint": "",
}
},
"update": {
"model_group": model_name,
"ptu_flat_cost": flat_cost,
"updated_at": now,
data={
"create": {
"team_id": charge.team_id,
"date": deployment_day.date_str,
"api_key": PTU_SENTINEL_API_KEY,
"model": charge.model_id,
"model_group": charge.model_name,
"custom_llm_provider": "",
"mcp_namespaced_tool_name": "",
"endpoint": "",
"ptu_flat_cost": charge.flat_cost,
},
"update": {
"model_group": charge.model_name,
"ptu_flat_cost": charge.flat_cost,
"updated_at": now,
},
},
},
)
)
async def _upsert_charge_with_retry(
prisma_client: "PrismaClient",
*,
charge: _PTUCharge,
date_str: str,
) -> bool:
"""Write one charge, retrying transient failures. Returns False once attempts are spent.
async def _write_deployment_day_with_retry(prisma_client: "PrismaClient", deployment_day: _DeploymentDay) -> bool:
"""Write one deployment day's charges in one transaction, retrying transient failures.
Returns False once attempts are spent, with none of that day's rows written.
The upsert is idempotent on the sentinel unique key, so a retry can only rewrite the
same amount for the same day. Retrying in-run matters because the scheduled job moves
on to the next date: a write lost here is a day of PTU cost that no later run replays.
The upserts are idempotent on the sentinel unique key, so a retry can only rewrite the
same amounts for the same day. Retrying in-run matters because the scheduled job moves
on to the next date: a write lost here waits for the next night's backfill.
"""
teams: Final = ",".join(charge.team_id for charge in deployment_day.charges)
for attempt in range(1, _UPSERT_ATTEMPTS + 1):
try:
await _upsert_ptu_daily_row(
prisma_client,
team_id=charge.team_id,
model_id=charge.model_id,
model_name=charge.model_name,
date_str=date_str,
flat_cost=charge.flat_cost,
)
async with _sentinel_batch(prisma_client) as batch:
_queue_deployment_day(batch.litellm_dailyteamspend, deployment_day)
return True
except Exception as exc: # noqa: BLE001 # one bad row must not stop the batch
except Exception as exc: # noqa: BLE001 # one bad deployment day must not stop the batch
if attempt < _UPSERT_ATTEMPTS:
verbose_proxy_logger.warning(
"PTU rollup: upsert attempt %d/%d failed for team=%s model=%s day=%s: %s",
"PTU rollup: write attempt %d/%d failed for model=%s teams=%s day=%s: %s",
attempt,
_UPSERT_ATTEMPTS,
charge.team_id,
charge.model_name,
date_str,
deployment_day.charges[0].model_name,
teams,
deployment_day.date_str,
exc,
)
await asyncio.sleep(_UPSERT_RETRY_BACKOFF_SECONDS * attempt)
continue
verbose_proxy_logger.error(
"PTU rollup: upsert failed after %d attempts for team=%s model=%s day=%s "
"PTU rollup: write failed after %d attempts for model=%s teams=%s day=%s "
"(rerun the rollup for that date to recover): %s",
_UPSERT_ATTEMPTS,
charge.team_id,
charge.model_name,
date_str,
deployment_day.charges[0].model_name,
teams,
deployment_day.date_str,
exc,
)
return False
async def _write_deployment_days(prisma_client: "PrismaClient", deployment_days: tuple[_DeploymentDay, ...]) -> int:
landed: Final = tuple(
[await _write_deployment_day_with_retry(prisma_client, deployment_day) for deployment_day in deployment_days]
)
return sum(len(day.charges) for day, ok in zip(deployment_days, landed, strict=True) if ok)
@dataclass(frozen=True, slots=True)
class _LoadedDeployments:
"""The deployments a run will price, and every deployment id it looked at.
@ -379,9 +438,7 @@ async def _load_ptu_models(prisma_client: "PrismaClient", *, router: object | No
rows: Final = await _proxy_model_table(prisma_client).find_many()
db_ids: Final = frozenset(model_id for row in rows if (model_id := str(getattr(row, "model_id", "") or "")))
config_records: Final = _config_deployments(router, owned_by_db=db_ids)
models: Final = tuple(
parsed for parsed in (_parse_ptu_model(row) for row in (*rows, *config_records)) if parsed is not None
)
models: Final = tuple(chain.from_iterable(_parse_ptu_models(row) for row in (*rows, *config_records)))
return _LoadedDeployments(
models=models,
scanned_ids=db_ids
@ -425,10 +482,7 @@ async def run_ptu_flat_cost_rollup(
ptu_models: Final = loaded.models
charges: Final = _aggregate_charges(ptu_models, day)
landed: Final = tuple(
[await _upsert_charge_with_retry(prisma_client, charge=charge, date_str=date_str) for charge in charges]
)
rows_written: Final = sum(landed)
rows_written: Final = await _write_deployment_days(prisma_client, _deployment_days(charges, date_str))
rows_failed: Final = len(charges) - rows_written
if not may_prune:
@ -454,16 +508,17 @@ async def run_ptu_flat_cost_rollup(
scanned_ids=loaded.scanned_ids,
)
models_processed: Final = len(frozenset(model.model_id for model in ptu_models))
verbose_proxy_logger.info(
"PTU rollup for %s: %d PTU models processed, %d rows written, %d rows failed",
date_str,
len(ptu_models),
models_processed,
rows_written,
rows_failed,
)
return RollupResult(
day=day,
models_processed=len(ptu_models),
models_processed=models_processed,
rows_written=rows_written,
rows_failed=rows_failed,
lapsed=_lapsed_models(ptu_models, run_started),
@ -488,11 +543,13 @@ def _lapsed_models(ptu_models: tuple[PTUModel, ...], now: datetime) -> tuple[str
charge the provider does not make for reserved capacity.
"""
return tuple(
_slack_safe(model.model_name)
for model in sorted(
(m for m in ptu_models if m.effective_to is not None and m.effective_to <= now),
key=lambda m: m.effective_to,
reverse=True,
dict.fromkeys(
_slack_safe(model.model_name)
for model in sorted(
(m for m in ptu_models if m.effective_to is not None and m.effective_to <= now),
key=lambda m: m.effective_to or now,
reverse=True,
)
)
)
@ -514,16 +571,26 @@ def _backfill_window(ptu_models: tuple[PTUModel, ...], end: date) -> tuple[date,
return tuple(start + timedelta(days=offset) for offset in range((end - start).days + 1))
def _unpriced_deployment_days(
ptu_models: tuple[PTUModel, ...], days: tuple[date, ...], priced: frozenset[tuple[str, str]]
) -> Iterator[_DeploymentDay]:
for day in days:
for deployment_day in _deployment_days(_aggregate_charges(ptu_models, day), day.isoformat()):
if (deployment_day.model_id, deployment_day.date_str) not in priced:
yield deployment_day
async def _existing_sentinel_keys(
prisma_client: "PrismaClient",
*,
start: date,
end: date,
) -> frozenset[tuple[str, str, str]]:
"""``(team_id, deployment id, date)`` of every PTU sentinel row within ``[start, end]``.
) -> frozenset[tuple[str, str]]:
"""``(deployment id, date)`` of every PTU sentinel row within ``[start, end]``.
The row's ``model`` column holds the deployment id, so this is an exact identity and
survives a rename. Nothing here reads the display name.
survives a rename. The team is left out on purpose: shares reassigned after a day was
billed must not bill that day again to the team that holds the share now.
"""
date_range: Final = {"gte": start.isoformat(), "lte": end.isoformat()}
rows: Final = await _daily_team_spend_table(prisma_client).find_many(
@ -531,7 +598,6 @@ async def _existing_sentinel_keys(
)
return frozenset(
(
str(getattr(row, "team_id", "") or ""),
str(getattr(row, "model", "") or ""),
str(getattr(row, "date", "") or ""),
)
@ -548,10 +614,12 @@ async def run_ptu_flat_cost_backfill(
Writes only the charges that are missing and never rewrites or deletes an existing
row, so a day already priced keeps the amount it was billed, whatever the config says
now. A day counts as priced when a sentinel row exists for that deployment id, so
renaming a deployment neither re-prices its history nor files a second charge beside
the row already there. Zero-cost days write nothing, which leaves a day
outside a window reconsidered on each run rather than recorded as done.
now. A day counts as priced when a sentinel row exists for that deployment id, for any
team, so neither renaming a deployment nor moving its shares to another team re-prices
its history or files a second charge beside the rows already there. That is sound
because a deployment day's team rows are written in one transaction. Zero-cost days
write nothing, which leaves a day outside a window reconsidered on each run rather
than recorded as done.
It deletes nothing. Removing a deployment stops it accruing new charges and leaves the
days it was billed for standing, since those days were incurred.
@ -569,36 +637,26 @@ async def run_ptu_flat_cost_backfill(
return BackfillResult(start=end, end=end, days_scanned=0, rows_written=0)
priced: Final = await _existing_sentinel_keys(prisma_client, start=days[0], end=days[-1])
missing: Final = tuple(
(day.isoformat(), charge)
for day in days
for charge in _aggregate_charges(ptu_models, day)
if (charge.team_id, charge.model_id, day.isoformat()) not in priced
)
missing: Final = tuple(_unpriced_deployment_days(ptu_models, days, priced))
if not missing:
return BackfillResult(start=days[0], end=days[-1], days_scanned=len(days), rows_written=0)
landed: Final = tuple(
[
await _upsert_charge_with_retry(prisma_client, charge=charge, date_str=date_str)
for date_str, charge in missing
]
)
rows_written: Final = sum(landed)
charge_count: Final = sum(len(deployment_day.charges) for deployment_day in missing)
rows_written: Final = await _write_deployment_days(prisma_client, missing)
verbose_proxy_logger.info(
"PTU backfill for %s to %s: %d unpriced charge(s) found, %d written, %d failed",
days[0].isoformat(),
days[-1].isoformat(),
len(missing),
charge_count,
rows_written,
len(missing) - rows_written,
charge_count - rows_written,
)
return BackfillResult(
start=days[0],
end=days[-1],
days_scanned=len(days),
rows_written=rows_written,
rows_failed=len(missing) - rows_written,
rows_failed=charge_count - rows_written,
)

View file

@ -106,6 +106,7 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import SERVICE_TIER_COST_KEY
from litellm.litellm_core_utils.ptu_pricing import (
PTU_COST_ATTRIBUTION_ENV_VAR,
declares_ptu,
declares_ptu_shares,
is_ptu_cost_attribution_enabled,
ptu_config_error,
ptu_identity_error,
@ -199,7 +200,9 @@ from litellm.router_utils.common_utils import (
get_request_team_id,
is_proxy_admin_request,
provider_for_generic_call,
request_user_api_key_auth,
resolve_model_group_alias,
team_may_use_deployment,
truncate_fallback_error_detail,
warn_on_ignored_model_capability_flags,
warn_on_provider_credential_mismatch,
@ -273,6 +276,14 @@ from litellm.router_utils.pre_call_checks.model_rate_limit_check import (
from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import (
PromptCachingDeploymentCheck,
)
from litellm.router_utils.ptu_shares import (
PTUTeamCeiling,
filter_ptu_shared_deployments,
ptu_capacity_warning,
ptu_share_teams,
share_holder_request_kwargs,
team_ptu_ceiling,
)
from litellm.router_utils.reasoning_effort_capability import (
deployment_is_catalog_mapped,
intersect_supported_reasoning_efforts,
@ -1258,6 +1269,8 @@ class Router:
self.cached_deployment_model_info = lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)(
self.get_deployment_model_info
)
self._ptu_share_teams = lru_cache(maxsize=1)(self._collect_ptu_share_teams)
self._cached_team_ptu_ceiling = lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)(self._resolve_team_ptu_ceiling)
self._discovered_model_info_cache: InMemoryCache = InMemoryCache(
max_size_in_memory=max(len(model_list or ()), 1),
default_ttl=2 * MODEL_INFO_REFRESH_SECONDS,
@ -4703,7 +4716,10 @@ class Router:
**kwargs,
):
await self._wait_for_scheduler_turn(
model=model, priority=priority, parent_otel_span=get_parent_otel_span_from_kwargs(kwargs)
model=model,
priority=priority,
parent_otel_span=get_parent_otel_span_from_kwargs(kwargs),
request_kwargs=cast(Mapping[str, object], kwargs), # cast-ok: kwargs keys are param names
)
try:
_response: Final = await self.acompletion(model=model, messages=messages, stream=stream, **kwargs)
@ -4727,7 +4743,10 @@ class Router:
kwargs: dict[str, object],
):
await self._wait_for_scheduler_turn(
model=model, priority=priority, parent_otel_span=get_parent_otel_span_from_kwargs(kwargs)
model=model,
priority=priority,
parent_otel_span=get_parent_otel_span_from_kwargs(kwargs),
request_kwargs=kwargs,
)
try:
_response: Final = await original_function(*args, **kwargs)
@ -4742,9 +4761,15 @@ class Router:
setattr(e, "priority", priority)
raise e
async def _wait_for_scheduler_turn(self, model: str, priority: int, parent_otel_span: Span | None) -> None:
async def _wait_for_scheduler_turn(
self, model: str, priority: int, parent_otel_span: Span | None, request_kwargs: Mapping[str, object]
) -> None:
async def healthy_deployments() -> Sequence[object]:
deployments, _ = await self._async_get_healthy_deployments(model=model, parent_otel_span=parent_otel_span)
deployments, _ = await self._async_get_healthy_deployments(
model=model,
parent_otel_span=parent_otel_span,
request_kwargs=share_holder_request_kwargs(request_kwargs, self._ptu_share_teams()),
)
return deployments
await self.scheduler.wait_for_turn(
@ -6051,7 +6076,9 @@ class Router:
model_group: Final = cast(str, initial_kwargs.get("model")) # cast-ok: kwargs always carries the model group
retry_kwargs: Final = mid_stream_retry_kwargs(initial_kwargs)
healthy_deployments, all_deployments = await self._async_get_healthy_deployments(
model=model_group, parent_otel_span=get_parent_otel_span_from_kwargs(retry_kwargs)
model=model_group,
parent_otel_span=get_parent_otel_span_from_kwargs(retry_kwargs),
request_kwargs=share_holder_request_kwargs(initial_kwargs, self._ptu_share_teams()),
)
budget, policy_applies = self._stream_retry_budget(_mid_stream_retry_trigger(e), initial_kwargs)
last_error = e # rebind-ok: the newest failure is what the fallback chain and the caller see
@ -8205,6 +8232,10 @@ class Router:
) = await self._async_get_healthy_deployments(
model=kwargs.get("model") or "",
parent_otel_span=parent_otel_span,
request_kwargs=share_holder_request_kwargs(
cast(Mapping[str, object], kwargs), # cast-ok: kwargs keys are param names
self._ptu_share_teams(),
),
)
# Check retry policy FIRST, before should_retry_this_error
@ -8297,6 +8328,10 @@ class Router:
) = await self._async_get_healthy_deployments(
model=_model,
parent_otel_span=parent_otel_span,
request_kwargs=share_holder_request_kwargs(
cast(Mapping[str, object], kwargs), # cast-ok: kwargs keys are param names
self._ptu_share_teams(),
),
)
else:
_healthy_deployments = []
@ -9110,18 +9145,19 @@ class Router:
return healthy_deployments, _all_deployments
async def _async_get_healthy_deployments(
self, model: str, parent_otel_span: Span | None
self, model: str, parent_otel_span: Span | None, request_kwargs: Mapping[str, object] | None = None
) -> tuple[list[dict], list[dict]]:
"""
Returns Tuple of:
- Tuple[List[Dict], List[Dict]]:
1. healthy_deployments: list of healthy deployments
2. all_deployments: list of all deployments
1. healthy_deployments: list of healthy deployments the caller of ``request_kwargs`` may use
2. all_deployments: list of all deployments that caller may use
"""
_all_deployments: list = []
try:
_, _all_deployments = self._common_checks_available_deployment(
model=model,
request_kwargs=request_kwargs,
)
if isinstance(_all_deployments, dict):
return [], _all_deployments
@ -9484,11 +9520,27 @@ class Router:
ptu_error: Final = (
(ptu_config_error(_model_info, model_name=_model_name) or identity_error) if config_sourced else None
)
if ptu_error is not None and is_ptu_cost_attribution_enabled():
declares_split: Final = declares_ptu_shares(_model_info)
if ptu_error is not None and (declares_split or is_ptu_cost_attribution_enabled()):
raise ValueError(ptu_error)
access_windows_error: Final = access_windows_config_error(_model_info, model_name=_model_name)
if access_windows_error is not None:
raise ValueError(access_windows_error)
capacity_warning: Final = (
ptu_capacity_warning(
_model_name,
MappingProxyType(
{ # pyright: ignore[reportUnknownArgumentType] # router deployment dicts are untyped
"model_info": _model_info,
"litellm_params": _litellm_params,
}
),
)
if is_ptu_cost_attribution_enabled()
else None
)
if capacity_warning is not None:
verbose_router_logger.warning(capacity_warning)
zeroed_pricing: Final = zeroed_ptu_pricing(_model_info, _litellm_params) if config_sourced else None
merged_params: Final[Mapping[str, object]] = (
_litellm_params if zeroed_pricing is None else MappingProxyType({**_litellm_params, **zeroed_pricing})
@ -10871,7 +10923,7 @@ class Router:
"""
model_info: Final = model.get("model_info") if isinstance(model, dict) else model.model_info
owner_team_id: Final = cast(Mapping[str, object], model_info).get("team_id") if model_info is not None else None
return owner_team_id is None or owner_team_id == team_id
return team_may_use_deployment(owner_team_id, team_id)
_deployment_usable_by_team = deployment_usable_by_team
@ -12701,6 +12753,34 @@ class Router:
]
return [{**DeploymentTypedDict(**m), "model_name": model_name} for m in matches]
def wildcard_route_deployments(self, model: str, team_id: str | None) -> Sequence[Mapping[str, object]]:
"""The deployments routing serves ``model`` from through a wildcard: the most specific
proxy-wide pattern matching it, else ``team_id``'s own pattern, never both."""
proxy_wide: Final[Sequence[Mapping[str, object]]] = self.pattern_router.get_deployments_by_pattern(model=model)
team_router: Final = self.team_pattern_routers.get(team_id) if team_id is not None else None
if proxy_wide or team_router is None:
return proxy_wide
return team_router.get_deployments_by_pattern(model=model)
def team_ptu_ceiling(self, team_id: str, model_group: str) -> PTUTeamCeiling | None:
"""The per-minute ceiling ``team_id``'s PTU shares set on the group serving ``model_group``,
else None, kept until the model list next changes."""
if team_id not in self._ptu_share_teams():
return None
return self._cached_team_ptu_ceiling(team_id, model_group)
def _collect_ptu_share_teams(self) -> frozenset[str]:
return ptu_share_teams(self.get_model_list() or ())
def _resolve_team_ptu_ceiling(self, team_id: str, model_group: str) -> PTUTeamCeiling | None:
return team_ptu_ceiling(
self.get_model_list() or (),
self.model_list,
team_id,
model_group,
lambda name: self.wildcard_route_deployments(name, team_id),
)
def get_model_list_of_routed_group(self, model_group: str) -> list[DeploymentTypedDict]:
"""
The deployments a request the router has already resolved to model_group is
@ -12748,6 +12828,8 @@ class Router:
"""
self._cached_get_model_group_info.cache_clear()
self.cached_deployment_model_info.cache_clear()
self._ptu_share_teams.cache_clear()
self._cached_team_ptu_ceiling.cache_clear()
self._zero_cost_cache.clear()
self._routing_group_rows = None
@ -13311,19 +13393,9 @@ class Router:
if team_deployments:
return model, team_deployments
pattern_deployments = self.pattern_router.get_deployments_by_pattern(
model=model,
)
if pattern_deployments:
return model, pattern_deployments
if request_team_id is not None and request_team_id in self.team_pattern_routers:
pattern_deployments = self.team_pattern_routers[request_team_id].get_deployments_by_pattern(
model=model,
)
if pattern_deployments:
return model, pattern_deployments
wildcard_deployments: Final = self.wildcard_route_deployments(model, request_team_id)
if wildcard_deployments:
return model, list(wildcard_deployments)
if self.default_deployment is not None:
# Shallow copy with nested litellm_params copy (100x+ faster than deepcopy)
@ -13387,7 +13459,7 @@ class Router:
messages: list[dict[str, str]] | None = None,
input: str | list | None = None,
specific_deployment: bool | None = False,
request_kwargs: dict | None = None,
request_kwargs: Mapping[str, object] | None = None,
) -> tuple[str, list | dict]:
"""
Common checks for 'get_available_deployment' across sync + async call.
@ -13583,13 +13655,20 @@ class Router:
model=model,
llm_provider="",
)
return result.deployments
shared: Final = filter_ptu_shared_deployments(result.deployments, request_team_id)
if shared.withheld and len(shared.deployments) == 0:
raise litellm.BadRequestError(
message=f"Deployment {model} is reserved for the teams holding a PTU share of it",
model=model,
llm_provider="",
)
return shared.deployments
def _filter_deployments_by_model_access_groups(
self,
model: str,
healthy_deployments: list,
request_kwargs: dict | None,
request_kwargs: Mapping[str, object] | None,
request_team_id: str | None,
) -> list:
"""
@ -13603,9 +13682,7 @@ class Router:
if not healthy_deployments or request_kwargs is None:
return healthy_deployments
metadata: Final = request_kwargs.get("metadata") or {}
litellm_metadata: Final = request_kwargs.get("litellm_metadata") or {}
user_api_key_auth: Final = metadata.get("user_api_key_auth") or litellm_metadata.get("user_api_key_auth")
user_api_key_auth: Final = request_user_api_key_auth(request_kwargs)
if user_api_key_auth is None:
return healthy_deployments

View file

@ -35,20 +35,29 @@ def without_router_only_kwargs(kwargs: Mapping[str, _V]) -> dict[str, _V]:
return {key: value for key, value in kwargs.items() if key not in ROUTER_ONLY_CALL_KWARGS}
def is_proxy_admin_request(request_kwargs: Mapping[str, object] | None) -> bool:
if request_kwargs is None:
return False
def request_user_api_key_auth(request_kwargs: Mapping[str, object]) -> object:
"""The caller's auth object, from whichever metadata bucket this surface writes to."""
metadata_value: Final = request_kwargs.get("metadata")
litellm_metadata_value: Final = request_kwargs.get("litellm_metadata")
metadata: Final = metadata_value if isinstance(metadata_value, Mapping) else {}
litellm_metadata: Final = litellm_metadata_value if isinstance(litellm_metadata_value, Mapping) else {}
user_api_key_auth: Final = metadata.get("user_api_key_auth") or litellm_metadata.get("user_api_key_auth")
return getattr(user_api_key_auth, "user_role", None) == "proxy_admin"
return metadata.get("user_api_key_auth") or litellm_metadata.get("user_api_key_auth")
def is_proxy_admin_request(request_kwargs: Mapping[str, object] | None) -> bool:
if request_kwargs is None:
return False
return getattr(request_user_api_key_auth(request_kwargs), "user_role", None) == "proxy_admin"
_is_proxy_admin_request = is_proxy_admin_request
def team_may_use_deployment(owner_team_id: object, team_id: str | None) -> bool:
"""A deployment owned by a team (``model_info.team_id``) serves only that team; one with no owner serves every team."""
return owner_team_id is None or owner_team_id == team_id
def get_request_team_id(request_kwargs: Mapping[str, object] | None) -> str | None:
"""The caller's team id, from whichever metadata bucket this surface writes to."""
if request_kwargs is None:

View file

@ -0,0 +1,275 @@
"""Per-team PTU shares on a shared Azure provisioned deployment.
A deployment declaring ``model_info.ptu_shares`` is served only to the teams named in it,
and each team's share converts to a normalized-tokens-per-minute ceiling through the
model's Azure sizing row, so the proxy enforces the split instead of an operator hand
converting PTUs to ``model_tpm_limit``.
"""
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from datetime import date
from typing import Final, Generic, TypeAlias, TypeVar
from litellm.litellm_core_utils.ptu_pricing import parsed_ptu_shares, ptu_terms
from litellm.llms.azure.ptu_capacity import (
PTUCapacity,
deployment_ptu_capacity,
is_azure_deployment,
is_str_keyed_mapping,
)
from litellm.router_utils.common_utils import get_request_team_id, team_may_use_deployment
_DeploymentT: Final = TypeVar("_DeploymentT", bound=Mapping[str, object])
WildcardRoute: TypeAlias = Callable[[str], Sequence[Mapping[str, object]]]
@dataclass(frozen=True, slots=True)
class PTUTeamCeiling:
model_group: str
tpm_limit: int
output_to_input_ratio: float
cached_input_ratio: float
def raw_output_token_limit(self) -> int:
"""The ceiling expressed in unweighted tokens: what fits under it when every token is output."""
return max(1, int(self.tpm_limit / max(self.output_to_input_ratio, 1.0)))
@dataclass(frozen=True, slots=True)
class PTUShareFilterResult(Generic[_DeploymentT]):
deployments: tuple[_DeploymentT, ...]
withheld: bool
def _deployment_shares(deployment: Mapping[str, object]) -> Mapping[str, int] | None:
"""The teams a deployment is declared split across, else None.
The declared map is the access rule on its own; the pricing terms are read separately, so a
split whose terms are missing still serves only the teams it names rather than everyone.
"""
model_info: Final = deployment.get("model_info")
if not is_str_keyed_mapping(model_info):
return None
return parsed_ptu_shares(model_info.get("ptu_shares"))
def filter_ptu_shared_deployments(
healthy_deployments: Sequence[_DeploymentT], request_team_id: str | None
) -> PTUShareFilterResult[_DeploymentT]:
"""Keep only the deployments ``request_team_id`` holds a share on when it holds one, else only
the ones split into no shares, so a team's ceiling and PTU-hours count just the traffic its
reservation serves, and an open deployment beside a shared one serves the teams holding none.
A caller with no team, the master key included, holds no share, the same way an access
window reserves a deployment away from it.
"""
team_id: Final = request_team_id or ""
checks: Final = tuple((deployment, _deployment_shares(deployment)) for deployment in healthy_deployments)
holds_share: Final = any(shares is not None and team_id in shares for _, shares in checks)
kept: Final = tuple(
deployment for deployment, shares in checks if (team_id in shares if shares is not None else not holds_share)
)
return PTUShareFilterResult(deployments=kept, withheld=len(kept) < len(checks))
def ptu_share_teams(deployments: Sequence[Mapping[str, object]]) -> frozenset[str]:
"""Every team holding a share on any of ``deployments``."""
return frozenset[str]().union(
*(shares for deployment in deployments if (shares := _deployment_shares(deployment)) is not None)
)
def share_holder_request_kwargs(
request_kwargs: Mapping[str, object], share_teams: frozenset[str]
) -> Mapping[str, object] | None:
"""``request_kwargs`` when the caller's team is in ``share_teams``, else None."""
return request_kwargs if get_request_team_id(request_kwargs) in share_teams else None
def deployments_shared_with(deployments: Sequence[Mapping[str, object]], team_id: str) -> tuple[str, ...]:
"""The ids of the deployments split into PTU shares ``team_id`` holds one of, each once."""
return tuple(
sorted(
frozenset(
str(_deployment_id(deployment))
for deployment in deployments
if (shares := _deployment_shares(deployment)) is not None and team_id in shares
)
)
)
def team_ptu_ceiling(
listed_rows: Sequence[Mapping[str, object]],
deployments: Sequence[Mapping[str, object]],
team_id: str,
requested_model: str,
wildcard_route: WildcardRoute,
) -> PTUTeamCeiling | None:
"""The per-minute normalized-token ceiling ``team_id``'s shares on the group serving
``requested_model`` add up to, else None when the team holds no share on a deployment with
a known sizing row.
``listed_rows`` is every row the router lists a name under, alias and routing-group copies
included, ``deployments`` is the router's own deployments, and ``wildcard_route`` gives the
deployments the wildcard route serving a name routes it to. A name resolves to the
deployments behind it, so a group, a routing group, a deployment id, a wildcard-matched
name, and a provider model all count against the one ceiling of the group whose shared
deployment the team can be served from.
Two shared deployments of different models in one group are weighted by the larger
output and cached-input ratios, which over-counts those tokens on the cheaper one rather
than under-counting them on the dearer one.
"""
model_group: Final = _model_group_of(listed_rows, deployments, team_id, requested_model, wildcard_route)
priced: Final = tuple(
(shares[team_id], capacity)
for deployment in model_group_deployments(deployments, model_group)
if (shares := _deployment_shares(deployment)) is not None
and team_id in shares
and (capacity := deployment_ptu_capacity(deployment)) is not None
)
if not priced:
return None
return PTUTeamCeiling(
model_group=model_group,
tpm_limit=sum(capacity.input_tpm_for(share) for share, capacity in priced),
output_to_input_ratio=max(capacity.output_to_input_ratio for _, capacity in priced),
cached_input_ratio=max(capacity.cached_input_ratio for _, capacity in priced),
)
def model_group_deployments(deployments: Sequence[_DeploymentT], model_group: str) -> tuple[_DeploymentT, ...]:
"""Every deployment serving ``model_group``: by its own name, or by the public name a
team-scoped deployment keeps in ``model_info.team_public_model_name`` after the router
renames it to a unique internal one."""
return tuple(
deployment
for deployment in deployments
if deployment.get("model_name") == model_group
or (
is_str_keyed_mapping(model_info := deployment.get("model_info"))
and model_info.get("team_public_model_name") == model_group
)
)
def _deployment_id(deployment: Mapping[str, object]) -> object:
model_info: Final = deployment.get("model_info")
return model_info.get("id") if is_str_keyed_mapping(model_info) else None
def _provider_model(deployment: Mapping[str, object]) -> object:
litellm_params: Final = deployment.get("litellm_params")
return litellm_params.get("model") if is_str_keyed_mapping(litellm_params) else None
def _deployment_model_group(deployment: Mapping[str, object]) -> str | None:
model_info: Final = deployment.get("model_info")
public_name: Final = model_info.get("team_public_model_name") if is_str_keyed_mapping(model_info) else None
if isinstance(public_name, str):
return public_name
model_name: Final = deployment.get("model_name")
return model_name if isinstance(model_name, str) else None
def _deployments_with_ids(
deployments: Sequence[Mapping[str, object]], rows: Sequence[Mapping[str, object]]
) -> tuple[Mapping[str, object], ...]:
ids: Final = frozenset(_deployment_id(row) for row in rows)
return tuple(deployment for deployment in deployments if _deployment_id(deployment) in ids)
def routed_deployments(
listed_rows: Sequence[Mapping[str, object]],
deployments: Sequence[Mapping[str, object]],
requested_model: str,
wildcard_route: WildcardRoute,
) -> tuple[Mapping[str, object], ...]:
"""The router's own deployments behind ``requested_model``, in the order the router resolves a
name: those of the rows listed under it when it names a group, else the one it names by id,
else those of the one wildcard route that serves it, else the ones it names by provider model."""
listed: Final = _deployments_with_ids(deployments, model_group_deployments(listed_rows, requested_model))
if listed:
return listed
by_id: Final = tuple(deployment for deployment in deployments if _deployment_id(deployment) == requested_model)
if by_id:
return by_id
wildcard: Final = _deployments_with_ids(deployments, wildcard_route(requested_model))
if wildcard:
return wildcard
return tuple(deployment for deployment in deployments if _provider_model(deployment) == requested_model)
def _model_group_of(
listed_rows: Sequence[Mapping[str, object]],
deployments: Sequence[Mapping[str, object]],
team_id: str,
requested_model: str,
wildcard_route: WildcardRoute,
) -> str:
"""The group of the deployment behind ``requested_model`` this team can be served from."""
servable: Final = team_servable_deployments(
routed_deployments(listed_rows, deployments, requested_model, wildcard_route), team_id
)
return next(
(group for deployment in servable if (group := _deployment_model_group(deployment)) is not None),
requested_model,
)
def _deployment_owner(deployment: Mapping[str, object]) -> object:
model_info: Final = deployment.get("model_info")
return model_info.get("team_id") if is_str_keyed_mapping(model_info) else None
def team_servable_deployments(deployments: Sequence[_DeploymentT], team_id: str) -> tuple[_DeploymentT, ...]:
"""The deployments ``team_id`` can be served from: never one reserved for another team, and only
the ones it holds a share on when it holds any."""
shareable: Final = filter_ptu_shared_deployments(deployments, team_id).deployments
return tuple(
deployment for deployment in shareable if team_may_use_deployment(_deployment_owner(deployment), team_id)
)
def model_group_ptu_capacity(deployments: Sequence[Mapping[str, object]], day: date) -> PTUCapacity | None:
"""The sizing row of the group's first deployment reserved on ``day``, single-team or shared,
so a team's tokens on that group that day convert to PTU-hours."""
return next(
(
capacity
for deployment in deployments
if is_str_keyed_mapping(model_info := deployment.get("model_info"))
and (terms := ptu_terms(model_info)) is not None
and terms.is_active_on(day)
and (capacity := deployment_ptu_capacity(deployment)) is not None
),
None,
)
def ptu_capacity_warning(model_name: str, deployment: Mapping[str, object]) -> str | None:
"""Why this reserved deployment's tokens cannot be converted to PTUs, else None.
Only a deployment that declares shares (whose ceilings need a sizing row) or one served by
Azure (whose PTU hours need one) is worth warning about; a single-team reservation on another
provider only ever used the flat-cost rollup, which needs no sizing.
"""
model_info: Final = deployment.get("model_info")
if not is_str_keyed_mapping(model_info) or ptu_terms(model_info) is None:
return None
if deployment_ptu_capacity(deployment) is not None:
return None
has_shares: Final = model_info.get("ptu_shares") is not None
if has_shares:
return (
f"PTU deployment '{model_name}' has no Azure sizing row for its model, so its PTU shares set no "
"team ceiling and its usage reports no PTU hours; set model_info.base_model to the Azure model name"
)
if is_azure_deployment(deployment):
return (
f"PTU deployment '{model_name}' has no Azure sizing row for its model, so its usage reports no "
"PTU hours; set model_info.base_model to the Azure model name"
)
return None

View file

@ -22,6 +22,7 @@ class GroupByDimension(str, Enum):
class SpendMetrics(LiteLLMBaseModel):
spend: float = Field(default=0.0)
flat_cost: float = Field(default=0.0)
ptu_hours: float = Field(default=0.0)
prompt_tokens: int = Field(default=0)
completion_tokens: int = Field(default=0)
cache_read_input_tokens: int = Field(default=0)
@ -86,6 +87,7 @@ class DailySpendData(LiteLLMBaseModel):
class DailySpendMetadata(LiteLLMBaseModel):
total_spend: float = Field(default=0.0)
total_flat_cost: float = Field(default=0.0)
total_ptu_hours: float = Field(default=0.0)
total_prompt_tokens: int = Field(default=0)
total_completion_tokens: int = Field(default=0)
total_tokens: int = Field(default=0)

View file

@ -22,7 +22,7 @@ from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
import httpx
from openai import Timeout as SDKTimeout
from pydantic import ConfigDict, Field, JsonValue, field_validator, model_validator
from pydantic import BeforeValidator, ConfigDict, Field, JsonValue, StrictInt, field_validator, model_validator
from typing_extensions import Protocol, ReadOnly, Required, TypedDict, runtime_checkable
from litellm._logging import verbose_logger
@ -250,6 +250,10 @@ class ModelAccessWindow(LiteLLMBaseModel):
return self
def _whole_float_as_int(value: object) -> object:
return int(value) if isinstance(value, float) and value.is_integer() else value
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.
@ -288,6 +292,7 @@ class ModelInfo(MirroredPricingParams):
cost_per_ptu_per_hour: float | None = None
ptu_effective_from: datetime.datetime | None = None
ptu_effective_to: datetime.datetime | None = None
ptu_shares: Mapping[str, Annotated[StrictInt, BeforeValidator(_whole_float_as_int)]] | None = None
# when tag-based routing's "!" or "&" constraints eliminate every deployment
# in this model group, fall back to the default-tagged pool instead of

View file

@ -112,6 +112,8 @@ ignored_function_names = [
"_get_wildcard_deployments", # Tested through the get_model_list_of_routed_group wildcard test in test_router.py
"_is_fallback_hop", # Tested through the order fallback hop tests in test_router_order_fallback.py
"_deployment_that_just_failed", # Tested through the same-boundary hop test in test_router_order_fallback.py
"_collect_ptu_share_teams", # Cached behind team_ptu_ceiling, tested through the live model list test in test_router.py
"_resolve_team_ptu_ceiling", # Cached behind team_ptu_ceiling, tested through the live model list test in test_router.py
]

View file

@ -0,0 +1,412 @@
import json
import uuid
from collections.abc import Generator, Iterator, Mapping
from concurrent.futures import ThreadPoolExecutor
from contextlib import ExitStack, contextmanager
from dataclasses import dataclass
from functools import partial
from types import MappingProxyType
from typing import Final
import httpx
import pytest
from integration._support.client import Gateway, eventually, gateway_from_environment, object_value, string_value
from integration._support.database import read_rows
from integration._support.openai_wire import chat_reply
from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue, TypeAdapter
pytestmark: Final = pytest.mark.timeout(4 * graceful_stop_seconds() + 240)
_RELOAD_SECONDS: Final = 5
_PROPAGATION_SECONDS: Final = 6 * _RELOAD_SECONDS
_ATTRIBUTION_ON: Final = MappingProxyType(
{
"LITELLM_ENABLE_PTU_COST_ATTRIBUTION": "True",
"PROXY_CONFIG_RELOAD_INTERVAL_SECONDS": str(_RELOAD_SECONDS),
}
)
_SCIM_HEADERS: Final = MappingProxyType({"Content-Type": "application/scim+json"})
_SHARES_REQUIRED: Final = "ptu_shares must map at least one team_id to a positive whole number of PTUs"
_DEPLOYMENT_IDS: Final = TypeAdapter(dict[str, list[str]])
_SAMPLES: Final = 16
def _reply(_request: Request) -> Reply:
return chat_reply(f"chatcmpl-{uuid.uuid4().hex}", "gpt-4.1", "served", stream=False)
@dataclass(frozen=True, slots=True)
class _Guarded:
owned: OwnedProxy
wire_url: str
@property
def proxy(self) -> Gateway:
return self.owned.gateway
@pytest.fixture(scope="module")
def guarded(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Guarded]:
directory: Final = tmp_path_factory.mktemp("ptu-share-guards")
with (
gateway_from_environment() as gateway,
wire_server(_reply) as wire,
owned_proxy_process(gateway, directory, _ATTRIBUTION_ON, workers=2) as owned,
):
yield _Guarded(owned, wire.url)
@pytest.fixture
def cleanups() -> Iterator[ExitStack]:
with ExitStack() as stack:
yield stack
def _ptu_terms(shares: Mapping[str, JsonValue] | None, count: int) -> dict[str, JsonValue]:
return {
"base_model": "azure/gpt-4.1",
"ptu_count": count,
"cost_per_ptu_per_hour": 2.0,
"ptu_effective_from": "2026-01-01T00:00:00Z",
"ptu_shares": dict(shares) if shares is not None else None,
}
def _deployment(wire_url: str, name: str, model_info: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
return {
"model_name": name,
"litellm_params": {
"model": "azure/ptu-guard",
"api_base": wire_url,
"api_key": "synthetic-azure-key",
"api_version": "2025-03-01-preview",
},
"model_info": dict(model_info),
}
def _model_name() -> str:
return f"ptu-guard-{uuid.uuid4().hex}"
def _stored_models(name: str) -> list[dict[str, JsonValue]]:
return read_rows('SELECT model_id, model_info FROM "LiteLLM_ProxyModelTable" WHERE model_name = %s', (name,))
def _team_rows(team_id: str) -> list[dict[str, JsonValue]]:
return read_rows('SELECT team_id FROM "LiteLLM_TeamTable" WHERE team_id = %s', (team_id,))
def _organization_rows(organization_id: str) -> list[dict[str, JsonValue]]:
return read_rows(
'SELECT organization_id FROM "LiteLLM_OrganizationTable" WHERE organization_id = %s', (organization_id,)
)
def _delete_team_if_present(proxy: Gateway, team_id: str) -> None:
if _team_rows(team_id):
proxy.post("/team/delete", {"team_ids": [team_id]})
assert _team_rows(team_id) == []
def _team(proxy: Gateway, cleanups: ExitStack, **fields: JsonValue) -> str:
team_id: Final = f"ptu-guard-team-{uuid.uuid4().hex}"
proxy.post("/team/new", {"team_id": team_id, "team_alias": team_id, **fields})
cleanups.callback(_delete_team_if_present, proxy, team_id)
return team_id
def _team_key(proxy: Gateway, team_id: str) -> str:
return string_value(proxy.post("/key/generate", {"team_id": team_id})["key"])
def _delete_organization_if_present(proxy: Gateway, organization_id: str, budget_id: str) -> None:
if _organization_rows(organization_id):
deleted: Final = proxy.request("DELETE", "/organization/delete", {"organization_ids": [organization_id]})
assert deleted.status_code == 200, deleted.text
assert _organization_rows(organization_id) == []
proxy.post("/budget/delete", {"id": budget_id})
def _organization(proxy: Gateway, cleanups: ExitStack) -> str:
created: Final = proxy.post("/organization/new", {"organization_alias": f"ptu-guard-org-{uuid.uuid4().hex}"})
organization_id: Final = string_value(created["organization_id"])
cleanups.callback(_delete_organization_if_present, proxy, organization_id, string_value(created["budget_id"]))
return organization_id
def _worker_deployment_ids(owned: OwnedProxy, model_name: str) -> tuple[int, frozenset[str]]:
with httpx.Client(
base_url=str(owned.gateway.client.base_url),
timeout=10,
trust_env=False,
limits=httpx.Limits(max_keepalive_connections=0),
) as client:
response: Final = client.get("/health/integration-deployment-ids", params={"model_name": model_name})
assert response.status_code == 200, response.text
return (
int(response.headers["x-integration-worker-pid"]),
frozenset(_DEPLOYMENT_IDS.validate_python(response.json())["model_ids"]),
)
def _served_on_no_worker(owned: OwnedProxy, model_name: str) -> bool:
with ThreadPoolExecutor(max_workers=8) as pool:
samples: Final = tuple(pool.map(partial(_worker_deployment_ids, owned), (model_name,) * _SAMPLES))
workers: Final = owned.worker_pids()
return (
bool(workers)
and workers <= frozenset(pid for pid, _ in samples)
and all(not model_ids for _, model_ids in samples)
)
def _delete_model(guarded: _Guarded, name: str, identity: str) -> None:
guarded.proxy.post("/model/delete", {"id": identity})
eventually(lambda: _served_on_no_worker(guarded.owned, name), bool, seconds=_PROPAGATION_SECONDS)
@contextmanager
def _shared_model(guarded: _Guarded, model_info: Mapping[str, JsonValue]) -> Generator[tuple[str, str]]:
name: Final = _model_name()
created: Final = guarded.proxy.post("/model/new", _deployment(guarded.wire_url, name, model_info))
identity: Final = string_value(object_value(created["model_info"])["id"])
try:
eventually(
lambda: guarded.owned.serves_only_on_every_worker(name, identity), bool, seconds=_PROPAGATION_SECONDS
)
yield name, identity
finally:
_delete_model(guarded, name, identity)
def _create(guarded: _Guarded, name: str, model_info: Mapping[str, JsonValue]) -> httpx.Response:
return guarded.proxy.request("POST", "/model/new", _deployment(guarded.wire_url, name, model_info))
def _create_raw(guarded: _Guarded, name: str, raw_shares: str, count: int) -> httpx.Response:
terms: Final = {key: value for key, value in _ptu_terms(None, count).items() if key != "ptu_shares"}
body: Final = json.dumps(_deployment(guarded.wire_url, name, terms))
spliced: Final = f'{body[:-2]}, "ptu_shares": {raw_shares}}}}}'
return guarded.proxy.client.post(
"/model/new",
content=spliced,
headers={"Authorization": f"Bearer {guarded.proxy.key}", "Content-Type": "application/json"},
)
def _holds(team_id: str, identity: str) -> str:
return f"Team {team_id} holds a PTU share on deployment {identity}"
def _chat_statuses(proxy: Gateway, key: str, model: str) -> tuple[int, ...]:
def status(_index: int) -> int:
with httpx.Client(
base_url=str(proxy.client.base_url),
timeout=30,
trust_env=False,
limits=httpx.Limits(max_keepalive_connections=0),
) as client:
return client.post(
"/v1/chat/completions",
json={"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": uuid.uuid4().hex}]},
headers={"Authorization": f"Bearer {key}"},
).status_code
with ThreadPoolExecutor(max_workers=8) as pool:
return tuple(pool.map(status, range(_SAMPLES)))
def test_a_share_for_a_team_that_does_not_exist_is_refused(guarded: _Guarded, cleanups: ExitStack) -> None:
holder: Final = _team(guarded.proxy, cleanups)
missing: Final = f"ptu-guard-missing-{uuid.uuid4().hex}"
name: Final = _model_name()
response: Final = _create(guarded, name, _ptu_terms({holder: 1, missing: 1}, 2))
assert response.status_code == 400, response.text
assert f"Team id={missing} does not exist in db" in response.text
assert _stored_models(name) == []
def test_shares_that_do_not_add_up_to_the_count_are_refused(guarded: _Guarded, cleanups: ExitStack) -> None:
holder: Final = _team(guarded.proxy, cleanups)
name: Final = _model_name()
response: Final = _create(guarded, name, _ptu_terms({holder: 1}, 2))
assert response.status_code == 400, response.text
assert "ptu_shares must add up to ptu_count (1 of 2 allocated)" in response.text
assert _stored_models(name) == []
def test_team_id_and_shares_together_are_refused(guarded: _Guarded, cleanups: ExitStack) -> None:
holder: Final = _team(guarded.proxy, cleanups)
name: Final = _model_name()
response: Final = _create(guarded, name, {**_ptu_terms({holder: 1}, 1), "team_id": holder})
assert response.status_code == 400, response.text
assert "team_id and ptu_shares cannot both be set" in response.text
assert _stored_models(name) == []
@pytest.mark.parametrize(
("raw_shares", "count", "status", "reason"),
[
pytest.param('["{holder}"]', 1, 422, "ptu_shares", id="list"),
pytest.param('"{holder}"', 1, 422, "ptu_shares", id="string"),
pytest.param("{{}}", 1, 400, _SHARES_REQUIRED, id="empty"),
pytest.param('{{"{holder}": -1}}', 1, 400, _SHARES_REQUIRED, id="negative"),
pytest.param('{{"{holder}": 0}}', 1, 400, _SHARES_REQUIRED, id="zero"),
pytest.param('{{"{holder}": 1.5}}', 2, 422, "ptu_shares", id="fraction"),
pytest.param('{{"{holder}": true}}', 1, 422, "ptu_shares", id="boolean"),
pytest.param('{{"{holder}": "1"}}', 1, 422, "ptu_shares", id="numeric-string"),
pytest.param("null", 1, 400, "team_id or ptu_shares is required when PTU fields are set", id="null"),
pytest.param(
'{{"{holder}": 1, "{holder}": 1}}',
2,
400,
"ptu_shares must add up to ptu_count (1 of 2 allocated)",
id="dup",
),
pytest.param('{{"{long_team}": 1}}', 1, 400, "does not exist in db", id="5kb-team-id"),
],
)
def test_malformed_shares_are_refused_with_the_reason(
guarded: _Guarded, cleanups: ExitStack, raw_shares: str, count: int, status: int, reason: str
) -> None:
holder: Final = _team(guarded.proxy, cleanups)
name: Final = _model_name()
response: Final = _create_raw(guarded, name, raw_shares.format(holder=holder, long_team="t" * 5000), count)
assert response.status_code == status, response.text
assert reason in response.text
assert _stored_models(name) == []
assert guarded.proxy.get("/model/info")["data"]
def test_with_attribution_off_shares_cannot_be_stored() -> None:
name: Final = _model_name()
with gateway_from_environment() as gateway:
try:
response: Final = gateway.request(
"POST",
"/model/new",
_deployment(gateway.upstream_url, name, {"ptu_shares": {f"ptu-guard-team-{uuid.uuid4().hex}": 1}}),
)
assert response.status_code == 400, response.text
assert "PTU cost attribution is disabled, so ptu_shares cannot be set" in response.text
assert _stored_models(name) == []
finally:
for row in _stored_models(name):
gateway.post("/model/delete", {"id": string_value(row["model_id"])})
def test_moving_a_share_to_a_team_that_does_not_exist_keeps_the_stored_shares(
guarded: _Guarded, cleanups: ExitStack
) -> None:
holder: Final = _team(guarded.proxy, cleanups)
missing: Final = f"ptu-guard-missing-{uuid.uuid4().hex}"
name, identity = cleanups.enter_context(_shared_model(guarded, _ptu_terms({holder: 1}, 1)))
response: Final = guarded.proxy.request(
"PATCH", f"/model/{identity}/update", {"model_info": {"ptu_shares": {missing: 1}}}
)
assert response.status_code == 400, response.text
assert f"Team id={missing} does not exist in db" in response.text
assert [object_value(row["model_info"])["ptu_shares"] for row in _stored_models(name)] == [{holder: 1}]
def test_deleting_a_team_that_holds_a_share_is_refused(guarded: _Guarded, cleanups: ExitStack) -> None:
holder: Final = _team(guarded.proxy, cleanups)
_name, identity = cleanups.enter_context(_shared_model(guarded, _ptu_terms({holder: 1}, 1)))
response: Final = guarded.proxy.request("POST", "/team/delete", {"team_ids": [holder]})
assert response.status_code == 400, response.text
assert _holds(holder, identity) in response.text
assert _team_rows(holder) == [{"team_id": holder}]
def test_deleting_an_organization_whose_team_holds_a_share_is_refused(guarded: _Guarded, cleanups: ExitStack) -> None:
organization_id: Final = _organization(guarded.proxy, cleanups)
holder: Final = _team(guarded.proxy, cleanups, organization_id=organization_id)
_name, identity = cleanups.enter_context(_shared_model(guarded, _ptu_terms({holder: 1}, 1)))
response: Final = guarded.proxy.request("DELETE", "/organization/delete", {"organization_ids": [organization_id]})
assert response.status_code == 400, response.text
assert _holds(holder, identity) in response.text
assert _organization_rows(organization_id) == [{"organization_id": organization_id}]
assert _team_rows(holder) == [{"team_id": holder}]
def test_deleting_a_scim_group_that_holds_a_share_is_refused(guarded: _Guarded, cleanups: ExitStack) -> None:
holder: Final = _team(guarded.proxy, cleanups)
_name, identity = cleanups.enter_context(_shared_model(guarded, _ptu_terms({holder: 1}, 1)))
response: Final = guarded.proxy.request("DELETE", f"/scim/v2/Groups/{holder}", headers=_SCIM_HEADERS)
assert response.status_code == 400, response.text
assert _holds(holder, identity) in response.text
assert _team_rows(holder) == [{"team_id": holder}]
def test_a_team_without_a_share_deletes_while_another_team_holds_one(guarded: _Guarded, cleanups: ExitStack) -> None:
holder: Final = _team(guarded.proxy, cleanups)
bystander: Final = _team(guarded.proxy, cleanups)
cleanups.enter_context(_shared_model(guarded, _ptu_terms({holder: 1}, 1)))
response: Final = guarded.proxy.request("POST", "/team/delete", {"team_ids": [bystander]})
assert response.status_code == 200, response.text
assert _team_rows(bystander) == []
def test_after_a_share_moves_the_former_holder_deletes_and_the_new_holder_is_served(
guarded: _Guarded, cleanups: ExitStack
) -> None:
former: Final = _team(guarded.proxy, cleanups)
successor: Final = _team(guarded.proxy, cleanups)
successor_key: Final = _team_key(guarded.proxy, successor)
name, identity = cleanups.enter_context(_shared_model(guarded, _ptu_terms({former: 1}, 1)))
moved: Final = guarded.proxy.request(
"PATCH", f"/model/{identity}/update", {"model_info": {"ptu_shares": {successor: 1}}}
)
assert moved.status_code == 200, moved.text
deleted: Final = eventually(
lambda: guarded.proxy.request("POST", "/team/delete", {"team_ids": [former]}),
lambda response: response.status_code != 400,
seconds=_PROPAGATION_SECONDS,
)
assert deleted.status_code == 200, deleted.text
assert _team_rows(former) == []
assert [object_value(row["model_info"])["ptu_shares"] for row in _stored_models(name)] == [{successor: 1}]
served: Final = eventually(
lambda: _chat_statuses(guarded.proxy, successor_key, name),
lambda statuses: set(statuses) == {200},
seconds=_PROPAGATION_SECONDS,
return_last_on_timeout=True,
)
assert set(served) == {200}, served
def test_shares_are_stored_as_whole_ptus_and_served_on_every_worker(guarded: _Guarded, cleanups: ExitStack) -> None:
first: Final = _team(guarded.proxy, cleanups)
second: Final = _team(guarded.proxy, cleanups)
first_key: Final = _team_key(guarded.proxy, first)
name, identity = cleanups.enter_context(_shared_model(guarded, _ptu_terms({first: 2.0, second: 1}, 3)))
info: Final = guarded.proxy.get("/model/info", {"litellm_model_id": identity})["data"]
assert isinstance(info, list) and len(info) == 1, info
assert object_value(object_value(info[0])["model_info"])["ptu_shares"] == {first: 2, second: 1}
assert [object_value(row["model_info"])["ptu_shares"] for row in _stored_models(name)] == [{first: 2, second: 1}]
assert set(_chat_statuses(guarded.proxy, first_key, name)) == {200}

View file

@ -0,0 +1,706 @@
import asyncio
import json
import os
import re
import threading
import uuid
from collections.abc import Callable, Iterator, Mapping
from concurrent.futures import Future, ThreadPoolExecutor
from dataclasses import dataclass
from pathlib import Path
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal
import httpx
import pytest
import yaml
from anthropic import Anthropic
from integration._support.anthropic_sse import (
LIFECYCLE,
Attempts,
delta_text,
dropping_reply,
event_types,
message_id,
message_json,
message_stream,
parse_sse,
stream_reply,
)
from integration._support.client import (
JSON_OBJECT,
Gateway,
eventually,
gateway_from_environment,
object_value,
string_value,
)
from integration._support.openai_wire import chat_reply, openai_error, responses_reply
from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
from integration._support.wire import Reply, Request, wire_server
from openai import AsyncOpenAI, OpenAI
from pydantic import JsonValue, TypeAdapter
from redis import Redis
if TYPE_CHECKING:
from litellm.llms.azure.ptu_capacity import PTUCapacity
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
_SHARED: Final = "ptu-shared"
_MIXED: Final = "ptu-mixed"
_PAIR: Final = "ptu-pair"
_CLAUDE: Final = "ptu-claude"
_SHARED_DEPLOYMENT: Final = "shared-gpt41"
_MIXED_SHARED_DEPLOYMENT: Final = "mixed-shared-gpt41"
_MIXED_OPEN_DEPLOYMENT: Final = "mixed-open-gpt41"
_PAIR_DEPLOYMENTS: Final = ("pair-a-gpt41", "pair-b-gpt41")
_CLAUDE_BACKEND: Final = "claude-ptu-backend"
_AZURE_KEY: Final = "synthetic-azure-key"
_ANTHROPIC_KEY: Final = "synthetic-anthropic-key"
_UPSTREAM_URL_PLACEHOLDER: Final = "upstream-url-placeholder"
_ATTRIBUTION_ON: Final = MappingProxyType({"LITELLM_ENABLE_PTU_COST_ATTRIBUTION": "True"})
_ATTRIBUTION_OFF: Final = MappingProxyType({"LITELLM_ENABLE_PTU_COST_ATTRIBUTION": "False"})
_RESERVED: Final = "is reserved for the teams holding a PTU share of it"
_TEXT: Final = "served"
_SMALL_MAX_TOKENS: Final = 16
_SETTLE_USAGE: Final[Mapping[str, JsonValue]] = MappingProxyType(
{"prompt_tokens": 5, "completion_tokens": 700, "total_tokens": 705}
)
_CACHED_USAGE: Final[Mapping[str, JsonValue]] = MappingProxyType(
{
"prompt_tokens": 1000,
"completion_tokens": 10,
"total_tokens": 1010,
"prompt_tokens_details": {"cached_tokens": 800},
}
)
_SHARE_ROLES: Final = (
"holder",
"settle",
"cached",
"refund",
"twin",
"oversize",
"retry",
"queue",
"burst",
"worker",
)
_ROLES: Final = (*_SHARE_ROLES, "outsider", "mixed", "pair", "messages")
_COUNTER: Final = TypeAdapter(int)
_CONFIG: Final = TypeAdapter(dict[str, JsonValue])
Surface = Literal["chat", "responses", "messages"]
Behavior = Literal["ok", "settle", "cached", "fail", "fail-once", "hold", "drop-once"]
def _gpt41() -> "PTUCapacity":
from litellm.llms.azure.ptu_capacity import AZURE_PTU_CAPACITY
return AZURE_PTU_CAPACITY["gpt-4.1"]
def _text_of(content: JsonValue) -> str:
if isinstance(content, str):
return content
if isinstance(content, list):
return "".join(_text_of(part) for part in content)
if isinstance(content, dict):
return _text_of(content.get("text", content.get("content", "")))
return ""
def _prompt(body: Mapping[str, JsonValue]) -> str:
messages: Final = body.get("messages")
if isinstance(messages, list) and messages:
return _text_of(messages[-1])
return _text_of(body.get("input", ""))
def _deployment_of(request: Request, body: Mapping[str, JsonValue]) -> str:
_, found, rest = request.target.partition("/openai/deployments/")
return rest.split("/", 1)[0] if found else str(body.get("model"))
def _usage_reply(identity: str, model: str, usage: Mapping[str, JsonValue]) -> Reply:
return Reply(
body=json.dumps(
{
"id": identity,
"object": "chat.completion",
"created": 1,
"model": model,
"choices": [{"index": 0, "message": {"role": "assistant", "content": _TEXT}, "finish_reason": "stop"}],
"usage": dict(usage),
}
).encode()
)
def _anthropic_reply(identity: str, behavior: str, attempt: int, stream: bool) -> Reply:
full: Final = message_stream(identity, _CLAUDE_BACKEND, _TEXT)
if behavior == "drop-once" and attempt == 1:
return dropping_reply(full, abort_after=1)
if stream:
return stream_reply(full)
return Reply(body=message_json(identity, _CLAUDE_BACKEND, _TEXT))
def _served_id(deployment: str, marker: str, attempt: int) -> str:
return f"chatcmpl-{deployment}-{marker}-a{attempt}"
def _openai_reply(request: Request, behavior: str, identity: str, deployment: str, attempt: int, stream: bool) -> Reply:
match behavior:
case "settle":
return _usage_reply(identity, deployment, _SETTLE_USAGE)
case "cached":
return _usage_reply(identity, deployment, _CACHED_USAGE)
case "fail":
return openai_error(500)
case "fail-once" if attempt == 1:
return openai_error(500)
case _:
pass
if "/responses" in request.target:
return responses_reply(identity, deployment, _TEXT, stream=stream)
return chat_reply(identity, deployment, _TEXT, stream=stream)
def _respond(attempts: Attempts, release: threading.Event) -> Callable[[Request], Reply]:
def respond(request: Request) -> Reply:
assert request.method == "POST", request
body: Final = JSON_OBJECT.validate_json(request.body)
behavior, _, marker = _prompt(body).partition(":")
deployment: Final = _deployment_of(request, body)
attempt: Final = attempts.record(f"{deployment}:{marker}")
identity: Final = _served_id(deployment, marker, attempt)
stream: Final = body.get("stream") is True
if behavior == "hold":
release.wait(timeout=30)
if request.target.startswith("/v1/messages"):
return _anthropic_reply(identity, behavior, attempt, stream)
return _openai_reply(request, behavior, identity, deployment, attempt, stream)
return respond
def _ptu_terms(shares: Mapping[str, int]) -> dict[str, JsonValue]:
return {
"base_model": "azure/gpt-4.1",
"ptu_count": sum(shares.values()),
"cost_per_ptu_per_hour": 2.0,
"ptu_effective_from": "2026-01-01T00:00:00Z",
"ptu_shares": dict(shares),
}
def _azure(model_name: str, deployment: str, model_info: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
return {
"model_name": model_name,
"litellm_params": {
"model": f"azure/{deployment}",
"api_base": _UPSTREAM_URL_PLACEHOLDER,
"api_key": _AZURE_KEY,
"api_version": "2025-03-01-preview",
},
"model_info": {"id": f"{deployment}-{uuid.uuid4().hex}", **model_info},
}
def _model_list(teams: Mapping[str, str]) -> list[JsonValue]:
one_each: Final = {teams[role]: 1 for role in _SHARE_ROLES}
return [
_azure(_SHARED, _SHARED_DEPLOYMENT, _ptu_terms(one_each)),
_azure(_MIXED, _MIXED_SHARED_DEPLOYMENT, _ptu_terms({teams["mixed"]: 1})),
_azure(_MIXED, _MIXED_OPEN_DEPLOYMENT, {"base_model": "azure/gpt-4.1"}),
*(_azure(_PAIR, deployment, _ptu_terms({teams["pair"]: 1})) for deployment in _PAIR_DEPLOYMENTS),
{
"model_name": _CLAUDE,
"litellm_params": {
"model": f"anthropic/{_CLAUDE_BACKEND}",
"api_base": _UPSTREAM_URL_PLACEHOLDER,
"api_key": _ANTHROPIC_KEY,
},
"model_info": {"id": f"claude-{uuid.uuid4().hex}", **_ptu_terms({teams["messages"]: 1})},
},
]
def _config(directory: Path, wire_url: str, model_list: list[JsonValue], name: str) -> Path:
config: Final = {
**_CONFIG.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())),
"model_list": model_list,
"router_settings": {"disable_cooldowns": True, "num_retries": 0, "ignore_invalid_deployments": True},
}
path: Final = directory / f"{name}.yaml"
path.write_text(yaml.safe_dump(config).replace(_UPSTREAM_URL_PLACEHOLDER, wire_url))
return path
def _team_ids() -> Mapping[str, str]:
run: Final = uuid.uuid4().hex[:12]
return MappingProxyType({role: f"ptu-{role}-{run}" for role in _ROLES})
def _create_team(proxy: Gateway, team_id: str) -> str:
proxy.post("/team/new", {"team_id": team_id, "team_alias": team_id})
return string_value(proxy.post("/key/generate", {"team_id": team_id})["key"])
def _delete_teams(gateway: Gateway, team_ids: tuple[str, ...]) -> None:
gateway.post("/team/delete", {"team_ids": list(team_ids)})
@dataclass(frozen=True, slots=True)
class _Rig:
owned: OwnedProxy
attempts: Attempts
release: threading.Event
teams: Mapping[str, str]
keys: Mapping[str, str]
@property
def proxy(self) -> Gateway:
return self.owned.gateway
def chat(self, model: str, role: str | None, behavior: Behavior, marker: str, **extra: JsonValue) -> httpx.Response:
return self.proxy.request(
"POST",
"/v1/chat/completions",
{
"model": model,
"max_tokens": _SMALL_MAX_TOKENS,
"messages": [{"role": "user", "content": f"{behavior}:{marker}"}],
**extra,
},
key=self.keys[role] if role is not None else None,
)
def send(self, surface: Surface, model: str, role: str | None, marker: str) -> httpx.Response:
key: Final = self.keys[role] if role is not None else None
prompt: Final = f"ok:{marker}"
match surface:
case "chat":
return self.chat(model, role, "ok", marker)
case "responses":
return self.proxy.request(
"POST",
"/v1/responses",
{"model": model, "input": prompt, "max_output_tokens": _SMALL_MAX_TOKENS},
key=key,
)
case "messages":
return self.proxy.request(
"POST",
"/v1/messages",
{
"model": model,
"max_tokens": _SMALL_MAX_TOKENS,
"messages": [{"role": "user", "content": prompt}],
},
key=key,
)
def hits(self, deployment: str, marker: str) -> int:
return self.attempts.count(f"{deployment}:{marker}")
@pytest.fixture(scope="module")
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]:
attempts: Final = Attempts()
release: Final = threading.Event()
teams: Final = _team_ids()
directory: Final = tmp_path_factory.mktemp("ptu-shares")
with gateway_from_environment() as gateway, wire_server(_respond(attempts, release)) as wire:
config: Final = _config(directory, wire.url, _model_list(teams), "ptu-shares")
try:
with owned_proxy_process(gateway, directory, _ATTRIBUTION_ON, config=config, workers=2) as owned:
keys: Final = MappingProxyType(
{role: _create_team(owned.gateway, team) for role, team in teams.items()}
)
yield _Rig(owned, attempts, release, teams, keys)
finally:
release.set()
_delete_teams(gateway, tuple(teams.values()))
def _ptu_counter(team_id: str, model_group: str) -> int | None:
with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]), decode_responses=True) as client:
keys: Final = tuple(client.scan_iter(match=f"*model_per_team_ptu:{team_id}:{model_group}}}:tokens"))
assert len(keys) <= 1, keys
raw: Final[object] = client.get(keys[0]) if keys else None
return None if raw is None else _COUNTER.validate_python(raw)
def _settled_counter(team_id: str, model_group: str, expected: int) -> int | None:
return eventually(
lambda: _ptu_counter(team_id, model_group),
lambda observed: observed == expected,
seconds=15,
return_last_on_timeout=True,
)
def _error_message(response: httpx.Response) -> str:
return string_value(object_value(JSON_OBJECT.validate_json(response.content)["error"])["message"])
@pytest.mark.parametrize("client", ("sync-sdk", "async-sdk-stream"))
def test_a_share_holder_is_served_chat_by_the_shared_deployment(rig: _Rig, client: str) -> None:
marker: Final = uuid.uuid4().hex
base_url: Final = f"{rig.proxy.client.base_url}/v1"
if client == "sync-sdk":
with OpenAI(api_key=rig.keys["holder"], base_url=base_url, max_retries=0, timeout=30) as sdk:
completion: Final = sdk.chat.completions.create(
model=_SHARED,
max_tokens=_SMALL_MAX_TOKENS,
messages=[{"role": "user", "content": f"ok:{marker}"}],
)
assert (completion.id, completion.choices[0].message.content) == (
_served_id(_SHARED_DEPLOYMENT, marker, 1),
_TEXT,
), completion
async def stream() -> tuple[str, str]:
async with AsyncOpenAI(api_key=rig.keys["holder"], base_url=base_url, max_retries=0, timeout=30) as sdk:
chunks: Final = await sdk.chat.completions.create(
model=_SHARED,
max_tokens=_SMALL_MAX_TOKENS,
messages=[{"role": "user", "content": f"ok:{marker}"}],
stream=True,
)
received: Final = tuple([chunk async for chunk in chunks])
return (
received[0].id,
"".join(choice.delta.content or "" for chunk in received for choice in chunk.choices),
)
if client == "async-sdk-stream":
assert asyncio.run(stream()) == (_served_id(_SHARED_DEPLOYMENT, marker, 1), _TEXT)
assert rig.hits(_SHARED_DEPLOYMENT, marker) == 1
@pytest.mark.parametrize("stream", (False, True), ids=("json", "stream"))
def test_a_share_holder_is_served_responses_by_the_shared_deployment(rig: _Rig, stream: bool) -> None:
marker: Final = uuid.uuid4().hex
response: Final = rig.proxy.request(
"POST",
"/v1/responses",
{"model": _SHARED, "input": f"ok:{marker}", "max_output_tokens": _SMALL_MAX_TOKENS, "stream": stream},
key=rig.keys["holder"],
)
assert response.status_code == 200, response.text
assert _served_id(_SHARED_DEPLOYMENT, marker, 1) in response.text, response.text
assert rig.hits(_SHARED_DEPLOYMENT, marker) == 1
def _anthropic_text(sdk: Anthropic, marker: str, *, stream: bool) -> str:
if stream:
with sdk.messages.stream(
model=_SHARED,
max_tokens=_SMALL_MAX_TOKENS,
messages=[{"role": "user", "content": f"ok:{marker}"}],
) as events:
return "".join(events.text_stream)
message: Final = sdk.messages.create(
model=_SHARED,
max_tokens=_SMALL_MAX_TOKENS,
messages=[{"role": "user", "content": f"ok:{marker}"}],
)
return "".join(block.text for block in message.content if block.type == "text")
@pytest.mark.parametrize("stream", (False, True), ids=("json", "stream"))
def test_a_share_holder_is_served_messages_by_the_shared_deployment(rig: _Rig, stream: bool) -> None:
marker: Final = uuid.uuid4().hex
with Anthropic(api_key=rig.keys["holder"], base_url=str(rig.proxy.client.base_url), max_retries=0) as sdk:
received: Final = _anthropic_text(sdk, marker, stream=stream)
assert received == _TEXT
assert rig.hits(_SHARED_DEPLOYMENT, marker) == 1
@pytest.mark.parametrize("caller", ("outsider", "master-key"))
@pytest.mark.parametrize("surface", ("chat", "responses", "messages"))
def test_a_caller_without_a_share_is_refused_the_shared_deployment(rig: _Rig, caller: str, surface: Surface) -> None:
marker: Final = uuid.uuid4().hex
response: Final = rig.send(surface, _SHARED, "outsider" if caller == "outsider" else None, marker)
assert response.status_code == 400, response.text
assert f"Deployment {_SHARED} {_RESERVED}" in response.text, response.text
assert rig.hits(_SHARED_DEPLOYMENT, marker) == 0
def test_a_mixed_group_serves_each_caller_from_its_own_deployment(rig: _Rig) -> None:
served: Final = {
caller: rig.chat(_MIXED, role, "ok", marker)
for caller, role, marker in (
("holder", "mixed", uuid.uuid4().hex),
("outsider", "outsider", uuid.uuid4().hex),
("master", None, uuid.uuid4().hex),
)
}
assert {caller: response.status_code for caller, response in served.items()} == dict.fromkeys(served, 200), {
caller: response.text for caller, response in served.items()
}
deployments: Final = {
caller: re.sub(
r"^chatcmpl-(.+)-[0-9a-f]{32}-a1$", r"\1", string_value(JSON_OBJECT.validate_json(response.content)["id"])
)
for caller, response in served.items()
}
assert deployments == {
"holder": _MIXED_SHARED_DEPLOYMENT,
"outsider": _MIXED_OPEN_DEPLOYMENT,
"master": _MIXED_OPEN_DEPLOYMENT,
}
@pytest.mark.parametrize(("model_group", "role", "ptus"), ((_SHARED, "oversize", 1), (_PAIR, "pair", 2)))
def test_a_request_bigger_than_the_whole_share_is_refused_up_front(
rig: _Rig, model_group: str, role: str, ptus: int
) -> None:
capacity: Final = _gpt41()
ceiling: Final = capacity.input_tpm_for(ptus)
marker: Final = uuid.uuid4().hex
response: Final = rig.chat(model_group, role, "ok", marker, max_tokens=ceiling)
assert response.status_code == 400, response.text
found: Final = re.search(
rf"This request needs (\d+) normalized tokens, more than your team's PTU share of {model_group} "
rf"allows in a minute \({ceiling}\), so it can never be admitted\. Lower max_tokens or shorten the prompt\.",
_error_message(response),
)
assert found is not None, response.text
assert int(found.group(1)) >= capacity.output_to_input_ratio * ceiling, response.text
assert rig.hits(_SHARED_DEPLOYMENT, marker) + sum(rig.hits(name, marker) for name in _PAIR_DEPLOYMENTS) == 0
assert _ptu_counter(rig.teams[role], model_group) in (None, 0)
def test_usage_settles_in_normalized_tokens_and_throttles_the_next_request(rig: _Rig) -> None:
capacity: Final = _gpt41()
ceiling: Final = capacity.input_tpm_for(1)
team: Final = rig.teams["settle"]
first: Final = uuid.uuid4().hex
settled: Final = rig.chat(_SHARED, "settle", "settle", first)
assert settled.status_code == 200, settled.text
assert settled.json()["usage"]["completion_tokens"] == _SETTLE_USAGE["completion_tokens"], settled.text
normalized: Final = round(
int(str(_SETTLE_USAGE["prompt_tokens"]))
+ capacity.output_to_input_ratio * int(str(_SETTLE_USAGE["completion_tokens"]))
)
assert _settled_counter(team, _SHARED, normalized) == normalized
max_tokens: Final = int((ceiling - normalized) / capacity.output_to_input_ratio) + 1
second: Final = uuid.uuid4().hex
throttled: Final = rig.chat(_SHARED, "settle", "ok", second, max_tokens=max_tokens)
assert throttled.status_code == 429, throttled.text
assert (
f"Rate limit exceeded for model_per_team_ptu: {team}:{_SHARED}. Limit type: tokens. Current limit: {ceiling},"
in _error_message(throttled)
), throttled.text
assert int(str(_SETTLE_USAGE["total_tokens"])) + max_tokens * capacity.output_to_input_ratio < ceiling
assert (rig.hits(_SHARED_DEPLOYMENT, first), rig.hits(_SHARED_DEPLOYMENT, second)) == (1, 0)
def test_cached_input_is_weighted_by_the_cached_ratio(rig: _Rig) -> None:
capacity: Final = _gpt41()
marker: Final = uuid.uuid4().hex
response: Final = rig.chat(_SHARED, "cached", "cached", marker)
assert response.status_code == 200, response.text
cached: Final = 800
expected: Final = round(
(int(str(_CACHED_USAGE["prompt_tokens"])) - cached)
+ capacity.cached_input_ratio * cached
+ capacity.output_to_input_ratio * int(str(_CACHED_USAGE["completion_tokens"]))
)
assert _settled_counter(rig.teams["cached"], _SHARED, expected) == expected
assert rig.hits(_SHARED_DEPLOYMENT, marker) == 1
def test_a_failed_upstream_call_refunds_the_reservation(rig: _Rig) -> None:
marker: Final = uuid.uuid4().hex
response: Final = rig.chat(_SHARED, "refund", "fail", marker)
assert response.status_code == 500, response.text
assert "scripted 500" in response.text, response.text
assert _settled_counter(rig.teams["refund"], _SHARED, 0) == 0
assert rig.hits(_SHARED_DEPLOYMENT, marker) == 1
def test_a_response_cache_hit_settles_against_the_share_like_every_token_limit(rig: _Rig) -> None:
capacity: Final = _gpt41()
marker: Final = uuid.uuid4().hex
first: Final = rig.chat(_SHARED, "twin", "settle", marker)
normalized: Final = round(
int(str(_SETTLE_USAGE["prompt_tokens"]))
+ capacity.output_to_input_ratio * int(str(_SETTLE_USAGE["completion_tokens"]))
)
assert _settled_counter(rig.teams["twin"], _SHARED, normalized) == normalized
repeat: Final = rig.chat(_SHARED, "twin", "settle", marker)
assert (first.status_code, repeat.status_code) == (200, 200), (first.text, repeat.text)
assert repeat.json()["id"] == first.json()["id"] == _served_id(_SHARED_DEPLOYMENT, marker, 1)
assert rig.hits(_SHARED_DEPLOYMENT, marker) == 1
assert _settled_counter(rig.teams["twin"], _SHARED, 2 * normalized) == 2 * normalized
def test_a_cached_response_is_not_served_to_a_caller_without_a_share(rig: _Rig) -> None:
marker: Final = uuid.uuid4().hex
cached: Final = rig.chat(_SHARED, "holder", "ok", marker)
assert cached.status_code == 200, cached.text
refused: Final = rig.chat(_SHARED, "outsider", "ok", marker)
assert refused.status_code == 400, refused.text
assert _RESERVED in refused.text, refused.text
assert rig.hits(_SHARED_DEPLOYMENT, marker) == 1
def test_a_retry_after_an_upstream_error_stays_on_the_shared_deployment(rig: _Rig) -> None:
marker: Final = uuid.uuid4().hex
response: Final = rig.chat(_SHARED, "retry", "fail-once", marker, num_retries=1)
assert response.status_code == 200, response.text
assert response.json()["id"] == _served_id(_SHARED_DEPLOYMENT, marker, 2), response.text
assert rig.hits(_SHARED_DEPLOYMENT, marker) == 2
def test_the_priority_scheduler_serves_share_holders_and_refuses_everyone_else(rig: _Rig) -> None:
served_marker: Final = uuid.uuid4().hex
refused_marker: Final = uuid.uuid4().hex
served: Final = rig.chat(_SHARED, "queue", "ok", served_marker, priority=1)
refused: Final = rig.chat(_SHARED, "outsider", "ok", refused_marker, priority=1)
assert served.status_code == 200, served.text
assert served.json()["id"] == _served_id(_SHARED_DEPLOYMENT, served_marker, 1), served.text
assert "x-litellm-request-prioritization-used" in served.headers, served.headers
assert refused.status_code == 400, refused.text
assert _RESERVED in refused.text, refused.text
assert (rig.hits(_SHARED_DEPLOYMENT, served_marker), rig.hits(_SHARED_DEPLOYMENT, refused_marker)) == (1, 0)
def test_a_dropped_messages_stream_retries_on_the_shared_deployment(rig: _Rig) -> None:
marker: Final = uuid.uuid4().hex
response: Final = rig.proxy.request(
"POST",
"/v1/messages",
{
"model": _CLAUDE,
"max_tokens": _SMALL_MAX_TOKENS,
"stream": True,
"num_retries": 1,
"messages": [{"role": "user", "content": f"drop-once:{marker}"}],
},
key=rig.keys["messages"],
)
assert response.status_code == 200, response.text
events: Final = parse_sse(response.text)
assert event_types(events) == LIFECYCLE, events
assert message_id(events) == _served_id(_CLAUDE_BACKEND, marker, 2), events
assert delta_text(events) == _TEXT, events
assert rig.hits(_CLAUDE_BACKEND, marker) == 2
def test_concurrent_requests_admit_only_what_fits_under_the_share(rig: _Rig) -> None:
capacity: Final = _gpt41()
ceiling: Final = capacity.input_tpm_for(1)
max_tokens: Final = int(ceiling / (2.5 * capacity.output_to_input_ratio))
markers: Final = tuple(uuid.uuid4().hex for _ in range(5))
with ThreadPoolExecutor(max_workers=len(markers)) as pool:
futures: Final[tuple[Future[httpx.Response], ...]] = tuple(
pool.submit(rig.chat, _SHARED, "burst", "hold", marker, max_tokens=max_tokens) for marker in markers
)
eventually(
lambda: (
sum(future.done() for future in futures)
+ sum(rig.hits(_SHARED_DEPLOYMENT, marker) for marker in markers)
),
lambda settled: settled >= len(markers),
seconds=30,
)
rig.release.set()
responses: Final = tuple(future.result(timeout=60) for future in futures)
assert sorted(response.status_code for response in responses) == [200, 200, 429, 429, 429], tuple(
response.text for response in responses
)
admitted: Final = tuple(
marker for marker, response in zip(markers, responses, strict=True) if response.status_code == 200
)
assert tuple(rig.hits(_SHARED_DEPLOYMENT, marker) for marker in markers) == tuple(
1 if marker in admitted else 0 for marker in markers
)
@pytest.mark.timeout(4 * graceful_stop_seconds() + 240)
def test_a_replacement_worker_rebuilds_the_share_rules(rig: _Rig) -> None:
victim: Final = min(rig.owned.worker_pids())
replacement: Final = rig.owned.kill_worker(victim)
assert replacement != victim
answers: Final = tuple(
(rig.chat(_SHARED, "worker", "ok", uuid.uuid4().hex), rig.chat(_SHARED, "outsider", "ok", uuid.uuid4().hex))
for _ in range(6)
)
assert tuple((served.status_code, refused.status_code) for served, refused in answers) == ((200, 400),) * 6, tuple(
(served.text, refused.text) for served, refused in answers
)
assert all(
string_value(JSON_OBJECT.validate_json(served.content)["id"]).startswith(f"chatcmpl-{_SHARED_DEPLOYMENT}-")
for served, _ in answers
)
def test_with_attribution_off_the_access_rule_holds_without_a_ceiling(tmp_path: Path) -> None:
attempts: Final = Attempts()
teams: Final = _team_ids()
capacity: Final = _gpt41()
with gateway_from_environment() as gateway, wire_server(_respond(attempts, threading.Event())) as wire:
config: Final = _config(tmp_path, wire.url, _model_list(teams), "ptu-shares-off")
try:
with owned_proxy_process(gateway, tmp_path, _ATTRIBUTION_OFF, config=config, workers=2) as owned:
holder: Final = _create_team(owned.gateway, teams["oversize"])
outsider: Final = _create_team(owned.gateway, teams["outsider"])
served_marker: Final = uuid.uuid4().hex
served: Final = owned.gateway.request(
"POST",
"/v1/chat/completions",
{
"model": _SHARED,
"max_tokens": capacity.input_tpm_for(1),
"messages": [{"role": "user", "content": f"ok:{served_marker}"}],
},
key=holder,
)
refused: Final = owned.gateway.request(
"POST",
"/v1/chat/completions",
{"model": _SHARED, "messages": [{"role": "user", "content": f"ok:{uuid.uuid4().hex}"}]},
key=outsider,
)
finally:
_delete_teams(gateway, (teams["oversize"], teams["outsider"]))
assert served.status_code == 200, served.text
assert served.json()["id"] == _served_id(_SHARED_DEPLOYMENT, served_marker, 1), served.text
assert refused.status_code == 400, refused.text
assert _RESERVED in refused.text, refused.text
assert _ptu_counter(teams["oversize"], _SHARED) is None
def test_a_split_that_does_not_add_up_is_dropped_at_load_with_attribution_off(tmp_path: Path) -> None:
attempts: Final = Attempts()
holder: Final = _team_ids()["oversize"]
marker: Final = uuid.uuid4().hex
with gateway_from_environment() as gateway, wire_server(_respond(attempts, threading.Event())) as wire:
split: Final = _azure(_SHARED, _SHARED_DEPLOYMENT, {**_ptu_terms({holder: 1}), "ptu_count": 2})
config: Final = _config(tmp_path, wire.url, [split], "ptu-split-mismatch")
try:
with owned_proxy_process(gateway, tmp_path, _ATTRIBUTION_OFF, config=config, workers=2) as owned:
key: Final = _create_team(owned.gateway, holder)
refused: Final = owned.gateway.request(
"POST",
"/v1/chat/completions",
{"model": _SHARED, "messages": [{"role": "user", "content": f"ok:{marker}"}]},
key=key,
)
log: Final = owned.log.read_text()
finally:
_delete_teams(gateway, (holder,))
assert refused.status_code == 400, refused.text
assert _RESERVED not in refused.text, refused.text
assert attempts.count(f"{_SHARED_DEPLOYMENT}:{marker}") == 0
assert (
f"PTU configuration on model '{_SHARED}' is invalid: ptu_shares must add up to ptu_count (1 of 2 allocated)"
in log
), log[-4000:]

View file

@ -0,0 +1,143 @@
import statistics
import threading
import uuid
from collections.abc import Iterator, Mapping
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from pathlib import Path
from types import MappingProxyType
from typing import Final
import httpx
import pytest
import yaml
from integration._support.client import Gateway, gateway_from_environment, string_value
from integration._support.openai_wire import chat_reply
from integration._support.process import graceful_stop_seconds, owned_proxy_process
from integration._support.timing import assert_within_factor_of_control, timed
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue, TypeAdapter
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 240)
_WORST_GROUP: Final = "ptu-regex-worst"
_CONTROL_GROUP: Final = "ptu-regex-unshared-worst"
_UPSTREAM_URL_PLACEHOLDER: Final = "upstream-url-placeholder"
_WORST_BASE_MODEL: Final = "azure/" + "-0000-00-00" * 5_958 + "x"
_ATTRIBUTION_ON: Final = MappingProxyType({"LITELLM_ENABLE_PTU_COST_ATTRIBUTION": "True"})
_REQUESTS: Final = 8
_CONFIG: Final = TypeAdapter(dict[str, JsonValue])
_FACTOR: Final = 3.0
def _reply(_request: Request) -> Reply:
return chat_reply(f"chatcmpl-{uuid.uuid4().hex}", "gpt-4.1", "served", stream=False)
def _share_terms(holder: str) -> dict[str, JsonValue]:
return {
"ptu_count": 1,
"cost_per_ptu_per_hour": 2.0,
"ptu_effective_from": "2026-01-01T00:00:00Z",
"ptu_shares": {holder: 1},
}
def _deployment(model_name: str, model_info: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
return {
"model_name": model_name,
"litellm_params": {
"model": f"azure/{model_name}",
"api_base": _UPSTREAM_URL_PLACEHOLDER,
"api_key": "synthetic-azure-key",
"api_version": "2025-03-01-preview",
},
"model_info": {"id": f"{model_name}-{uuid.uuid4().hex}", "base_model": _WORST_BASE_MODEL, **model_info},
}
def _config(directory: Path, wire_url: str, holder: str) -> Path:
config: Final = {
**_CONFIG.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())),
"model_list": [
_deployment(_CONTROL_GROUP, {}),
_deployment(_WORST_GROUP, _share_terms(holder)),
],
"router_settings": {"disable_cooldowns": True, "num_retries": 0, "ignore_invalid_deployments": True},
}
path: Final = directory / "ptu-regex.yaml"
path.write_text(yaml.safe_dump(config, width=2_000_000).replace(_UPSTREAM_URL_PLACEHOLDER, wire_url))
return path
@dataclass(frozen=True, slots=True)
class _Timing:
proxy: Gateway
key: str
@pytest.fixture(scope="module")
def timing(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Timing]:
directory: Final = tmp_path_factory.mktemp("ptu-regex")
holder: Final = f"ptu-regex-holder-{uuid.uuid4().hex}"
with gateway_from_environment() as gateway, wire_server(_reply) as wire:
try:
with owned_proxy_process(
gateway, directory, _ATTRIBUTION_ON, config=_config(directory, wire.url, holder), workers=2
) as owned:
owned.gateway.post("/team/new", {"team_id": holder, "team_alias": holder})
key: Final = string_value(owned.gateway.post("/key/generate", {"team_id": holder})["key"])
yield _Timing(owned.gateway, key)
finally:
gateway.post("/team/delete", {"team_ids": [holder]})
def _chat_seconds(proxy: Gateway, key: str, model: str) -> float:
result: Final = timed(
lambda: proxy.request(
"POST",
"/v1/chat/completions",
{"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": uuid.uuid4().hex}]},
key=key,
)
)
assert result.value.status_code == 200, result.value.text
return result.seconds
def _liveliness_seconds(base_url: str, stop: threading.Event) -> tuple[float, ...]:
def probe() -> float:
with httpx.Client(
base_url=base_url, timeout=30, trust_env=False, limits=httpx.Limits(max_keepalive_connections=0)
) as client:
result: Final = timed(lambda: client.get("/health/liveliness"))
assert result.value.status_code == 200, result.value.text
return result.seconds
return tuple(probe() for _ in iter(stop.is_set, True))
def _under_load(timing: _Timing, model: str) -> tuple[float, float]:
stop: Final = threading.Event()
with ThreadPoolExecutor(max_workers=1) as pool:
liveliness: Final = pool.submit(_liveliness_seconds, str(timing.proxy.client.base_url), stop)
try:
chats: Final = tuple(_chat_seconds(timing.proxy, timing.key, model) for _ in range(_REQUESTS))
finally:
stop.set()
probes: Final = liveliness.result()
assert probes, "no liveliness probe completed while the requests ran"
return statistics.median(chats), statistics.median(probes)
def test_a_64kb_base_model_on_a_shared_deployment_keeps_requests_and_liveliness_near_the_unshared_one(
timing: _Timing,
) -> None:
warmed: Final = tuple(_chat_seconds(timing.proxy, timing.key, model) for model in (_CONTROL_GROUP, _WORST_GROUP))
assert len(warmed) == 2
control_chat, control_liveliness = _under_load(timing, _CONTROL_GROUP)
worst_chat, worst_liveliness = _under_load(timing, _WORST_GROUP)
assert_within_factor_of_control(worst_chat, control=control_chat, factor=_FACTOR)
assert_within_factor_of_control(worst_liveliness, control=control_liveliness, factor=_FACTOR)

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

View file

@ -13,6 +13,7 @@
"total_compression_savings_spend": 0.0,
"total_failed_requests": 0,
"total_flat_cost": 0.0,
"total_ptu_hours": 0.0,
"total_gateway_injected_caching_savings_spend": 0.0,
"total_pages": 1,
"total_prompt_caching_savings_spend": 0.0,
@ -45,6 +46,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 2,
@ -73,6 +75,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 3,
@ -101,6 +104,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 4,
@ -129,6 +133,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1000,
@ -157,6 +162,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 5,
@ -189,6 +195,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 2,
@ -217,6 +224,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 3,
@ -245,6 +253,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 4,
@ -273,6 +282,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1000,
@ -301,6 +311,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 5,
@ -323,6 +334,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1014,
@ -355,6 +367,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 2,
@ -383,6 +396,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 3,
@ -411,6 +425,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 4,
@ -439,6 +454,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1000,
@ -467,6 +483,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 5,
@ -491,6 +508,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1014,
@ -524,6 +542,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1000,
@ -546,6 +565,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1000,
@ -576,6 +596,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 2,
@ -604,6 +625,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 3,
@ -632,6 +654,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 4,
@ -654,6 +677,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 9,
@ -677,6 +701,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 0,
@ -707,6 +732,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 5,
@ -729,6 +755,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 5,
@ -761,6 +788,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1000,
@ -783,6 +811,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1000,
@ -813,6 +842,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 2,
@ -841,6 +871,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 3,
@ -869,6 +900,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 4,
@ -891,6 +923,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 9,
@ -914,6 +947,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 0,
@ -944,6 +978,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 5,
@ -966,6 +1001,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 5,
@ -998,6 +1034,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 2,
@ -1026,6 +1063,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 3,
@ -1054,6 +1092,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 4,
@ -1082,6 +1121,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1000,
@ -1110,6 +1150,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 5,
@ -1132,6 +1173,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1014,
@ -1155,6 +1197,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1014,

View file

@ -13,6 +13,7 @@
"total_compression_savings_spend": 0.0,
"total_failed_requests": 0,
"total_flat_cost": 0.0,
"total_ptu_hours": 0.0,
"total_gateway_injected_caching_savings_spend": 0.0,
"total_pages": 1,
"total_prompt_caching_savings_spend": 0.0,
@ -45,6 +46,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 2,
@ -73,6 +75,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 3,
@ -101,6 +104,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 4,
@ -129,6 +133,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1000,
@ -157,6 +162,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 5,
@ -189,6 +195,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 0,
@ -217,6 +224,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 2,
@ -245,6 +253,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 3,
@ -273,6 +282,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 4,
@ -301,6 +311,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1000,
@ -329,6 +340,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 5,
@ -351,6 +363,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1014,
@ -383,6 +396,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 2,
@ -411,6 +425,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 3,
@ -439,6 +454,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 4,
@ -467,6 +483,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1000,
@ -495,6 +512,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 5,
@ -519,6 +537,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1014,
@ -552,6 +571,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1000,
@ -574,6 +594,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1000,
@ -604,6 +625,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 2,
@ -632,6 +654,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 3,
@ -660,6 +683,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 4,
@ -682,6 +706,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 9,
@ -705,6 +730,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 0,
@ -735,6 +761,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 5,
@ -757,6 +784,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 5,
@ -789,6 +817,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1000,
@ -811,6 +840,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1000,
@ -841,6 +871,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 2,
@ -869,6 +900,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 3,
@ -897,6 +929,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 4,
@ -919,6 +952,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 9,
@ -942,6 +976,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 0,
@ -972,6 +1007,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 5,
@ -994,6 +1030,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 5,
@ -1026,6 +1063,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 2,
@ -1054,6 +1092,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 3,
@ -1082,6 +1121,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 4,
@ -1110,6 +1150,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1000,
@ -1138,6 +1179,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 5,
@ -1160,6 +1202,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1014,
@ -1183,6 +1226,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1014,

View file

@ -13,6 +13,7 @@
"total_compression_savings_spend": 0.0,
"total_failed_requests": 0,
"total_flat_cost": 0.0,
"total_ptu_hours": 0.0,
"total_gateway_injected_caching_savings_spend": 0.0,
"total_pages": 1,
"total_prompt_caching_savings_spend": 0.0,
@ -45,6 +46,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 2,
@ -73,6 +75,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 3,
@ -101,6 +104,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 4,
@ -129,6 +133,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1000,
@ -157,6 +162,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 5,
@ -189,6 +195,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 2,
@ -217,6 +224,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 3,
@ -245,6 +253,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 4,
@ -273,6 +282,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1000,
@ -301,6 +311,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 5,
@ -323,6 +334,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1014,
@ -357,6 +369,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1000,
@ -379,6 +392,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1000,
@ -409,6 +423,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 2,
@ -437,6 +452,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 3,
@ -465,6 +481,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 4,
@ -487,6 +504,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 9,
@ -510,6 +528,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 0,
@ -540,6 +559,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 5,
@ -562,6 +582,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 5,
@ -594,6 +615,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1000,
@ -616,6 +638,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1000,
@ -646,6 +669,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 2,
@ -674,6 +698,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 3,
@ -702,6 +727,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 4,
@ -724,6 +750,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 9,
@ -747,6 +774,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 0,
@ -777,6 +805,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 5,
@ -799,6 +828,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 5,
@ -831,6 +861,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 2,
@ -859,6 +890,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 3,
@ -887,6 +919,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 4,
@ -915,6 +948,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1000,
@ -943,6 +977,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 5,
@ -965,6 +1000,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1014,
@ -988,6 +1024,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1014,

View file

@ -13,6 +13,7 @@
"total_compression_savings_spend": 0.0,
"total_failed_requests": 0,
"total_flat_cost": 0.0,
"total_ptu_hours": 0.0,
"total_gateway_injected_caching_savings_spend": 0.0,
"total_pages": 1,
"total_prompt_caching_savings_spend": 0.0,
@ -45,6 +46,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 2,
@ -73,6 +75,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 3,
@ -101,6 +104,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 4,
@ -129,6 +133,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1000,
@ -157,6 +162,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 5,
@ -189,6 +195,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 0,
@ -217,6 +224,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 2,
@ -245,6 +253,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 3,
@ -273,6 +282,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 4,
@ -301,6 +311,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1000,
@ -329,6 +340,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 5,
@ -351,6 +363,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1014,
@ -383,6 +396,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 2,
@ -411,6 +425,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 3,
@ -439,6 +454,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 4,
@ -467,6 +483,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1000,
@ -495,6 +512,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 5,
@ -520,6 +538,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1014,
@ -553,6 +572,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1000,
@ -575,6 +595,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1000,
@ -605,6 +626,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 2,
@ -633,6 +655,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 3,
@ -661,6 +684,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 4,
@ -683,6 +707,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 9,
@ -706,6 +731,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 0,
@ -736,6 +762,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 5,
@ -758,6 +785,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 5,
@ -790,6 +818,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1000,
@ -812,6 +841,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1000,
@ -842,6 +872,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 2,
@ -870,6 +901,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 3,
@ -898,6 +930,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 4,
@ -920,6 +953,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 9,
@ -943,6 +977,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 0,
@ -973,6 +1008,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 5,
@ -995,6 +1031,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 5,
@ -1027,6 +1064,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 2,
@ -1055,6 +1093,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 3,
@ -1083,6 +1122,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 4,
@ -1111,6 +1151,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1000,
@ -1139,6 +1180,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 5,
@ -1161,6 +1203,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1014,
@ -1184,6 +1227,7 @@
"compression_savings_spend": 0.0,
"failed_requests": 0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"gateway_injected_caching_savings_spend": 0.0,
"prompt_caching_savings_spend": 0.0,
"prompt_tokens": 1014,

View file

@ -390,6 +390,7 @@ def _daily_tag_metrics(
return {
"spend": spend,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"cache_read_input_tokens": 0,

View file

@ -0,0 +1,362 @@
import asyncio
import json
import math
import re
import uuid
from collections.abc import Iterator, Mapping
from contextlib import ExitStack
from dataclasses import dataclass
from datetime import date, datetime, timedelta, timezone
from pathlib import Path
from types import MappingProxyType
from typing import TYPE_CHECKING, Final
import pytest
import yaml
from integration._support.client import (
JSON_OBJECT,
Gateway,
eventually,
gateway_from_environment,
object_value,
string_value,
)
from integration._support.database import read_rows, write_rows
from integration._support.prisma_clients import connected_proxy_prisma
from integration._support.process import graceful_stop_seconds, owned_proxy_process
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue, TypeAdapter
from litellm.constants import PTU_SENTINEL_API_KEY
if TYPE_CHECKING:
from litellm.llms.azure.ptu_capacity import PTUCapacity
pytestmark: Final = pytest.mark.timeout(4 * graceful_stop_seconds() + 240)
_GROUP: Final = "ptu-ledger"
_SHARED_DEPLOYMENT: Final = "ledger-shared-gpt41"
_OPEN_DEPLOYMENT: Final = "ledger-open-gpt41"
_UPSTREAM_URL_PLACEHOLDER: Final = "upstream-url-placeholder"
_RATE: Final = 2.0
_ATTRIBUTION_ON: Final = MappingProxyType(
{"LITELLM_ENABLE_PTU_COST_ATTRIBUTION": "True", "PROXY_CONFIG_RELOAD_INTERVAL_SECONDS": "5"}
)
_ATTRIBUTION_OFF: Final = MappingProxyType({"LITELLM_ENABLE_PTU_COST_ATTRIBUTION": "False"})
_PROMPT_TOKENS: Final = 1000
_CACHED_TOKENS: Final = 800
_COMPLETION_TOKENS: Final = 10
_HOLDER_REQUESTS: Final = 2
_CONFIG: Final = TypeAdapter(dict[str, JsonValue])
_DEPLOYMENT_PATH: Final = re.compile(r"/openai/deployments/([^/]+)/")
def _gpt41() -> "PTUCapacity":
from litellm.llms.azure.ptu_capacity import AZURE_PTU_CAPACITY
return AZURE_PTU_CAPACITY["gpt-4.1"]
def _reply(request: Request) -> Reply:
matched: Final = _DEPLOYMENT_PATH.search(request.target)
deployment: Final = matched.group(1) if matched else "unknown"
return Reply(
body=json.dumps(
{
"id": f"chatcmpl-{deployment}-{uuid.uuid4().hex}",
"object": "chat.completion",
"created": 1,
"model": "gpt-4.1",
"choices": [
{"index": 0, "message": {"role": "assistant", "content": "served"}, "finish_reason": "stop"}
],
"usage": {
"prompt_tokens": _PROMPT_TOKENS,
"completion_tokens": _COMPLETION_TOKENS,
"total_tokens": _PROMPT_TOKENS + _COMPLETION_TOKENS,
"prompt_tokens_details": {"cached_tokens": _CACHED_TOKENS},
},
}
).encode()
)
def _ptu_terms(shares: Mapping[str, int], effective_from: str) -> dict[str, JsonValue]:
return {
"base_model": "azure/gpt-4.1",
"ptu_count": sum(shares.values()),
"cost_per_ptu_per_hour": _RATE,
"ptu_effective_from": effective_from,
"ptu_shares": dict(shares),
}
def _azure(
model_name: str, deployment: str, api_base: str, model_info: Mapping[str, JsonValue]
) -> dict[str, JsonValue]:
return {
"model_name": model_name,
"litellm_params": {
"model": f"azure/{deployment}",
"api_base": api_base,
"api_key": "synthetic-azure-key",
"api_version": "2025-03-01-preview",
},
"model_info": dict(model_info),
}
def _config(directory: Path, wire_url: str, holder: str, name: str) -> Path:
config: Final = {
**_CONFIG.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())),
"model_list": [
_azure(
_GROUP,
_SHARED_DEPLOYMENT,
_UPSTREAM_URL_PLACEHOLDER,
{"id": f"{_SHARED_DEPLOYMENT}-{uuid.uuid4().hex}", **_ptu_terms({holder: 1}, "2026-01-01T00:00:00Z")},
),
_azure(
_GROUP,
_OPEN_DEPLOYMENT,
_UPSTREAM_URL_PLACEHOLDER,
{"id": f"{_OPEN_DEPLOYMENT}-{uuid.uuid4().hex}", "base_model": "azure/gpt-4.1"},
),
],
"router_settings": {"disable_cooldowns": True, "num_retries": 0, "ignore_invalid_deployments": True},
}
path: Final = directory / f"{name}.yaml"
path.write_text(yaml.safe_dump(config).replace(_UPSTREAM_URL_PLACEHOLDER, wire_url))
return path
def _today() -> date:
return datetime.now(timezone.utc).date()
def _window() -> Mapping[str, str]:
return MappingProxyType(
{
"start_date": (_today() - timedelta(days=1)).isoformat(),
"end_date": (_today() + timedelta(days=1)).isoformat(),
}
)
def _team_rows(team_id: str) -> list[dict[str, JsonValue]]:
return read_rows('SELECT team_id FROM "LiteLLM_TeamTable" WHERE team_id = %s', (team_id,))
def _delete_team_once_unheld(proxy: Gateway, team_id: str) -> None:
if _team_rows(team_id):
eventually(
lambda: proxy.request("POST", "/team/delete", {"team_ids": [team_id]}).status_code,
lambda status: status == 200,
seconds=60,
)
assert _team_rows(team_id) == []
def _delete_user(proxy: Gateway, user_id: str) -> None:
response: Final = proxy.request("POST", "/user/delete", {"user_ids": [user_id]})
assert response.status_code == 200, response.text
def _team(proxy: Gateway, cleanups: ExitStack) -> str:
team_id: Final = f"ptu-ledger-team-{uuid.uuid4().hex}"
proxy.post("/team/new", {"team_id": team_id, "team_alias": team_id})
cleanups.callback(_delete_team_once_unheld, proxy, team_id)
return team_id
def _chat(proxy: Gateway, key: str) -> str:
response: Final = proxy.request(
"POST",
"/v1/chat/completions",
{"model": _GROUP, "max_tokens": 16, "messages": [{"role": "user", "content": uuid.uuid4().hex}]},
key=key,
)
assert response.status_code == 200, response.text
return string_value(JSON_OBJECT.validate_json(response.content)["id"])
def _model_groups(day: JsonValue) -> dict[str, JsonValue]:
return object_value(object_value(object_value(day)["breakdown"])["model_groups"])
def _model_group_metrics(activity: Mapping[str, JsonValue]) -> tuple[dict[str, JsonValue], ...]:
results: Final = activity["results"]
assert isinstance(results, list), activity
return tuple(
object_value(object_value(_model_groups(day)[_GROUP])["metrics"])
for day in results
if _GROUP in _model_groups(day)
)
def _requests(activity: Mapping[str, JsonValue]) -> int:
return sum(int(str(metrics["api_requests"])) for metrics in _model_group_metrics(activity))
def _ptu_hours_per_request() -> float:
capacity: Final = _gpt41()
normalized: Final = (
(_PROMPT_TOKENS - _CACHED_TOKENS)
+ capacity.cached_input_ratio * _CACHED_TOKENS
+ capacity.output_to_input_ratio * _COMPLETION_TOKENS
)
return normalized / (capacity.input_tpm_per_ptu * 60)
@dataclass(frozen=True, slots=True)
class _Ledger:
proxy: Gateway
config: Path
holder: str
outsider: str
member: str
holder_ids: tuple[str, ...]
outsider_id: str
def _activity(proxy: Gateway, path: str, params: Mapping[str, str], requests: int) -> dict[str, JsonValue]:
return eventually(
lambda: proxy.get(path, {**_window(), **params}),
lambda activity: _requests(activity) == requests,
seconds=60,
)
@pytest.fixture(scope="module")
def ledger(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Ledger]:
directory: Final = tmp_path_factory.mktemp("ptu-ledger")
holder: Final = f"ptu-ledger-holder-{uuid.uuid4().hex}"
outsider: Final = f"ptu-ledger-outsider-{uuid.uuid4().hex}"
member: Final = f"ptu-ledger-member-{uuid.uuid4().hex}"
with ExitStack() as stack:
gateway: Final = stack.enter_context(gateway_from_environment())
wire: Final = stack.enter_context(wire_server(_reply))
config: Final = _config(directory, wire.url, holder, "ptu-ledger")
stack.callback(_delete_team_once_unheld, gateway, outsider)
stack.callback(_delete_team_once_unheld, gateway, holder)
owned: Final = stack.enter_context(
owned_proxy_process(gateway, directory, _ATTRIBUTION_ON, config=config, workers=2)
)
proxy: Final = owned.gateway
proxy.post("/team/new", {"team_id": holder, "team_alias": holder})
proxy.post("/team/new", {"team_id": outsider, "team_alias": outsider})
proxy.post("/user/new", {"user_id": member, "auto_create_key": False})
stack.callback(_delete_user, proxy, member)
holder_key: Final = string_value(proxy.post("/key/generate", {"team_id": holder, "user_id": member})["key"])
outsider_key: Final = string_value(proxy.post("/key/generate", {"team_id": outsider})["key"])
holder_ids: Final = tuple(_chat(proxy, holder_key) for _ in range(_HOLDER_REQUESTS))
outsider_id: Final = _chat(proxy, outsider_key)
_activity(proxy, "/team/daily/activity", {"team_ids": holder}, _HOLDER_REQUESTS)
_activity(proxy, "/team/daily/activity", {"team_ids": outsider}, 1)
yield _Ledger(proxy, config, holder, outsider, member, holder_ids, outsider_id)
@pytest.mark.parametrize(
"path", ["/team/daily/activity", "/team/daily/activity/aggregated"], ids=["daily", "daily-aggregated"]
)
def test_a_share_holder_reads_the_ptu_hours_its_tokens_used(ledger: _Ledger, path: str) -> None:
assert all(identity.startswith(f"chatcmpl-{_SHARED_DEPLOYMENT}-") for identity in ledger.holder_ids)
activity: Final = _activity(ledger.proxy, path, {"team_ids": ledger.holder}, _HOLDER_REQUESTS)
(metrics,) = _model_group_metrics(activity)
expected: Final = _HOLDER_REQUESTS * _ptu_hours_per_request()
assert (metrics["prompt_tokens"], metrics["completion_tokens"], metrics["cache_read_input_tokens"]) == (
_HOLDER_REQUESTS * _PROMPT_TOKENS,
_HOLDER_REQUESTS * _COMPLETION_TOKENS,
_HOLDER_REQUESTS * _CACHED_TOKENS,
)
assert math.isclose(float(str(metrics["ptu_hours"])), expected), activity
assert math.isclose(float(str(object_value(activity["metadata"])["total_ptu_hours"])), expected), activity
def test_a_team_served_by_the_open_deployment_reads_no_ptu_hours(ledger: _Ledger) -> None:
assert ledger.outsider_id.startswith(f"chatcmpl-{_OPEN_DEPLOYMENT}-")
activity: Final = _activity(ledger.proxy, "/team/daily/activity", {"team_ids": ledger.outsider}, 1)
(metrics,) = _model_group_metrics(activity)
assert metrics["prompt_tokens"] == _PROMPT_TOKENS
assert metrics["ptu_hours"] == 0, activity
assert object_value(activity["metadata"])["total_ptu_hours"] == 0, activity
def test_user_activity_reports_no_ptu_hours(ledger: _Ledger) -> None:
activity: Final = _activity(ledger.proxy, "/user/daily/activity", {"user_id": ledger.member}, _HOLDER_REQUESTS)
(metrics,) = _model_group_metrics(activity)
assert metrics["prompt_tokens"] == _HOLDER_REQUESTS * _PROMPT_TOKENS
assert metrics["ptu_hours"] == 0, activity
def test_with_attribution_off_team_activity_reports_no_ptu_hours(ledger: _Ledger, tmp_path: Path) -> None:
with (
gateway_from_environment() as gateway,
owned_proxy_process(gateway, tmp_path, _ATTRIBUTION_OFF, config=ledger.config, workers=2) as owned,
):
activity: Final = _activity(
owned.gateway, "/team/daily/activity", {"team_ids": ledger.holder}, _HOLDER_REQUESTS
)
(metrics,) = _model_group_metrics(activity)
assert metrics["prompt_tokens"] == _HOLDER_REQUESTS * _PROMPT_TOKENS
assert metrics["ptu_hours"] == 0, activity
assert object_value(activity["metadata"])["total_ptu_hours"] == 0, activity
async def _rollup(day: date) -> None:
from litellm.proxy.spend_tracking.ptu_flat_cost_rollup import run_ptu_flat_cost_rollup
async with connected_proxy_prisma() as prisma:
await run_ptu_flat_cost_rollup(prisma, target_date=day)
def _flat_costs(model_id: str, day: date) -> dict[str, JsonValue]:
rows: Final = read_rows(
'SELECT team_id, ptu_flat_cost FROM "LiteLLM_DailyTeamSpend" WHERE api_key = %s AND model = %s AND date = %s',
(PTU_SENTINEL_API_KEY, model_id, day.isoformat()),
)
return {string_value(row["team_id"]): row["ptu_flat_cost"] for row in rows}
def _delete_flat_costs(model_id: str) -> None:
write_rows(
'DELETE FROM "LiteLLM_DailyTeamSpend" WHERE api_key = %s AND model = %s', (PTU_SENTINEL_API_KEY, model_id)
)
def test_the_rollup_bills_each_team_its_share_and_follows_a_move(ledger: _Ledger) -> None:
yesterday: Final = _today() - timedelta(days=1)
effective_from: Final = f"{yesterday.isoformat()}T00:00:00Z"
with ExitStack() as cleanups:
first: Final = _team(ledger.proxy, cleanups)
second: Final = _team(ledger.proxy, cleanups)
third: Final = _team(ledger.proxy, cleanups)
created: Final = ledger.proxy.post(
"/model/new",
_azure(
f"ptu-ledger-rollup-{uuid.uuid4().hex}",
"ledger-rollup-gpt41",
"http://127.0.0.1:9",
_ptu_terms({first: 2, second: 1}, effective_from),
),
)
model_id: Final = string_value(object_value(created["model_info"])["id"])
cleanups.callback(_delete_flat_costs, model_id)
cleanups.callback(ledger.proxy.post, "/model/delete", {"id": model_id})
asyncio.run(_rollup(yesterday))
asyncio.run(_rollup(yesterday))
assert _flat_costs(model_id, yesterday) == {first: 2 * _RATE * 24, second: 1 * _RATE * 24}
moved: Final = ledger.proxy.request(
"PATCH", f"/model/{model_id}/update", {"model_info": {"ptu_shares": {first: 1, third: 2}}}
)
assert moved.status_code == 200, moved.text
asyncio.run(_rollup(yesterday))
assert _flat_costs(model_id, yesterday) == {first: 1 * _RATE * 24, third: 2 * _RATE * 24}

View file

@ -38,7 +38,7 @@ def test_a_complete_reservation_is_accepted():
terms = ptu_terms(_VALID)
assert terms is not None
assert terms.team_id == "team-alpha"
assert terms.shares == {"team-alpha": 100}
assert terms.ptu_count == 100
assert terms.effective_from == datetime(2026, 1, 1, tzinfo=timezone.utc)
assert terms.effective_to is None
@ -193,8 +193,8 @@ def test_a_deployment_with_no_ptu_fields_is_not_a_ptu_deployment():
@pytest.mark.parametrize(
"override, expected",
[
({"team_id": None}, "team_id is required when PTU fields are set (one model maps to one team)"),
({"team_id": ""}, "team_id is required when PTU fields are set (one model maps to one team)"),
({"team_id": None}, "team_id or ptu_shares is required when PTU fields are set"),
({"team_id": ""}, "team_id or ptu_shares is required when PTU fields are set"),
({"cost_per_ptu_per_hour": None}, "ptu_count and cost_per_ptu_per_hour must be set together"),
({"ptu_count": None}, "ptu_count and cost_per_ptu_per_hour must be set together"),
({"ptu_effective_to": "2025-01-01T00:00:00Z"}, "ptu_effective_to must be after ptu_effective_from"),
@ -356,3 +356,139 @@ def test_no_spillover_marker_returns_none():
is None
)
assert azure_spillover(response_headers=None, additional_headers=None) is None
# --- a reservation split into per-team shares ------------------------------------------
_SHARED = {**{k: v for k, v in _VALID.items() if k != "team_id"}, "ptu_shares": {"team-a": 60, "team-b": 40}}
def test_a_split_reservation_attributes_each_share_to_its_team():
terms = ptu_terms(_SHARED)
assert terms is not None
assert terms.shares == {"team-a": 60, "team-b": 40}
assert terms.ptu_count == 100
assert ptu_config_error(_SHARED) is None
def test_a_single_team_reservation_holds_the_whole_count_under_that_team():
terms = ptu_terms(_VALID)
assert terms is not None
assert sum(terms.shares.values()) == terms.ptu_count
@pytest.mark.parametrize(
"override, expected",
[
(
{"team_id": "team-alpha"},
"team_id and ptu_shares cannot both be set; ptu_shares lists every team the capacity is split across",
),
({"ptu_shares": {}}, "ptu_shares must map at least one team_id to a positive whole number of PTUs"),
(
{"ptu_shares": {"team-a": 0, "team-b": 100}},
"ptu_shares must map at least one team_id to a positive whole number of PTUs",
),
(
{"ptu_shares": {"team-a": 50.5, "team-b": 49.5}},
"ptu_shares must map at least one team_id to a positive whole number of PTUs",
),
(
{"ptu_shares": {"team-a": True}},
"ptu_shares must map at least one team_id to a positive whole number of PTUs",
),
(
{"ptu_shares": {"team-a": "60", "team-b": 40}},
"ptu_shares must map at least one team_id to a positive whole number of PTUs",
),
(
{"ptu_shares": {"team-a": float("inf"), "team-b": 40}},
"ptu_shares must map at least one team_id to a positive whole number of PTUs",
),
({"ptu_shares": {"": 100}}, "ptu_shares must map at least one team_id to a positive whole number of PTUs"),
({"ptu_shares": {None: 100}}, "ptu_shares must map at least one team_id to a positive whole number of PTUs"),
({"ptu_shares": {1: 100}}, "ptu_shares must map at least one team_id to a positive whole number of PTUs"),
({"ptu_shares": ["team-a"]}, "ptu_shares must map at least one team_id to a positive whole number of PTUs"),
({"ptu_shares": {"team-a": 60, "team-b": 30}}, "ptu_shares must add up to ptu_count (90 of 100 allocated)"),
({"ptu_shares": {"team-a": 60, "team-b": 50}}, "ptu_shares must add up to ptu_count (110 of 100 allocated)"),
],
ids=[
"team_id and shares",
"empty shares",
"zero share",
"fractional share",
"boolean share",
"string share",
"infinite share",
"blank team",
"null team",
"numeric team",
"not a mapping",
"shares short of the count",
"shares over the count",
],
)
def test_an_incoherent_split_names_its_reason_and_reserves_nothing(override, expected):
assert ptu_config_error({**_SHARED, **override}) == expected
assert ptu_terms({**_SHARED, **override}) is None
def test_a_whole_count_written_as_a_float_is_checked_against_the_shares_all_the_same():
assert ptu_config_error({**_SHARED, "ptu_count": 100.0, "ptu_shares": {"team-a": 60, "team-b": 30}}) == (
"ptu_shares must add up to ptu_count (90 of 100 allocated)"
)
terms = ptu_terms({**_SHARED, "ptu_count": 100.0})
assert terms is not None
assert terms.ptu_count == 100
def test_whole_shares_written_as_floats_split_the_reservation_like_integers():
written_as_floats = {**_SHARED, "ptu_count": 100.0, "ptu_shares": {"team-a": 60.0, "team-b": 40.0}}
terms = ptu_terms(written_as_floats)
assert ptu_config_error(written_as_floats) is None
assert terms is not None
assert terms.shares == {"team-a": 60, "team-b": 40}
assert {type(share) for share in terms.shares.values()} == {int}
def test_a_rate_that_is_not_a_number_reserves_nothing():
assert ptu_terms({**_VALID, "cost_per_ptu_per_hour": "two cents"}) is None
def test_a_fractional_count_reserves_nothing():
assert ptu_terms({**_VALID, "ptu_count": 100.5}) is None
assert ptu_terms({**_SHARED, "ptu_count": 100.5}) is None
@pytest.mark.parametrize(
"shares_only",
[
{"ptu_shares": {"team-a": 60, "team-b": 40}},
{"ptu_shares": {"team-a": 60, "team-b": 40}, "ptu_effective_from": "2026-01-01T00:00:00Z"},
],
ids=["shares alone", "shares with a window"],
)
def test_shares_without_the_count_and_rate_are_refused_instead_of_registering_unpriced(shares_only):
"""Shares split a count, so a split with no count has nothing to add up to and would register a
deployment the share filter and ceiling honour while flat cost never rolls up."""
assert ptu_config_error(shares_only) == "ptu_count and cost_per_ptu_per_hour are required when ptu_shares is set"
assert ptu_terms(shares_only) is None
def test_the_split_is_named_after_the_deployment_when_the_caller_supplies_one():
error = ptu_config_error({**_SHARED, "ptu_shares": {"team-a": 1}}, model_name="gpt-4.1-ptu")
assert error is not None
assert "'gpt-4.1-ptu'" in error
assert "1 of 100 allocated" in error
def test_a_split_reservation_is_zeroed_like_a_single_team_one():
zeroed = _with_flag(_SHARED)
assert zeroed["input_cost_per_token"] == 0.0
assert zeroed["output_cost_per_token"] == 0.0

View file

@ -0,0 +1,106 @@
"""Tests for the Azure PTU sizing table and the conversions built on it."""
from typing import Final
import pytest
from litellm.llms.azure.ptu_capacity import (
AZURE_PTU_CAPACITY,
PTUCapacity,
azure_ptu_capacity,
deployment_ptu_capacity,
is_azure_deployment,
normalized_tokens,
ptu_hours,
)
_GPT41: Final = AZURE_PTU_CAPACITY["gpt-4.1"]
_ROW: Final = PTUCapacity(input_tpm_per_ptu=1_000, output_to_input_ratio=4.0)
_CACHED_ROW: Final = PTUCapacity(input_tpm_per_ptu=1_000, output_to_input_ratio=5.0, cached_input_ratio=0.1)
@pytest.mark.parametrize(
"model",
["gpt-4.1", "azure/gpt-4.1", "azure/gpt-4.1-2025-04-14", "GPT-4.1", "azure/eastus/gpt-4.1-2025-04-14"],
)
def test_a_model_resolves_to_its_row_through_a_provider_prefix_and_a_dated_version(model):
assert azure_ptu_capacity(model) is _GPT41
def test_a_deployment_name_that_is_not_a_model_has_no_row():
assert azure_ptu_capacity("azure/my-ptu-deployment") is None
assert azure_ptu_capacity("") is None
def test_a_dated_version_only_drops_a_full_date_suffix():
"""``gpt-4o-mini`` must not lose its ``-mini`` the way ``-2024-07-18`` is dropped."""
assert azure_ptu_capacity("azure/gpt-4o-mini-2024-07-18") is AZURE_PTU_CAPACITY["gpt-4o-mini"]
assert azure_ptu_capacity("gpt-4o-mini") is not AZURE_PTU_CAPACITY["gpt-4o"]
def test_every_row_serves_its_input_tpm_for_an_hour():
for capacity in AZURE_PTU_CAPACITY.values():
assert capacity.normalized_tokens_per_ptu_hour == capacity.input_tpm_per_ptu * 60
assert capacity.input_tpm_per_ptu > 0
assert capacity.output_to_input_ratio >= 1.0
assert 0.0 <= capacity.cached_input_ratio < 1.0
def test_a_deployment_prefers_its_declared_base_model_over_its_deployment_name():
deployment: Final = {
"model_info": {"base_model": "azure/gpt-4.1"},
"litellm_params": {"model": "azure/gpt-4o"},
}
assert deployment_ptu_capacity(deployment) is _GPT41
def test_a_deployment_falls_back_to_its_litellm_model_when_no_base_model_is_declared():
assert deployment_ptu_capacity({"litellm_params": {"model": "azure/gpt-4o"}}) is AZURE_PTU_CAPACITY["gpt-4o"]
assert (
deployment_ptu_capacity({"model_info": {"base_model": ""}, "litellm_params": {"model": "azure/gpt-4o"}})
is (AZURE_PTU_CAPACITY["gpt-4o"])
)
def test_a_deployment_with_no_recognisable_model_has_no_row():
assert deployment_ptu_capacity({"litellm_params": {"model": "azure/team-a-ptu"}}) is None
assert deployment_ptu_capacity({"model_info": None, "litellm_params": None}) is None
assert deployment_ptu_capacity({}) is None
def test_output_is_weighted_by_the_models_ratio_and_uncached_input_counts_in_full():
assert normalized_tokens(_ROW, prompt_tokens=100, completion_tokens=10) == pytest.approx(140.0)
def test_cached_input_is_free_unless_the_row_prices_it():
assert normalized_tokens(_ROW, prompt_tokens=100, completion_tokens=0, cache_read_tokens=60) == pytest.approx(40.0)
assert normalized_tokens(
_CACHED_ROW, prompt_tokens=100, completion_tokens=0, cache_read_tokens=60
) == pytest.approx(46.0)
def test_cached_input_never_exceeds_the_prompt_and_negatives_count_as_zero():
assert normalized_tokens(_ROW, prompt_tokens=100, completion_tokens=0, cache_read_tokens=250) == 0.0
assert normalized_tokens(_ROW, prompt_tokens=-5, completion_tokens=-5, cache_read_tokens=-5) == 0.0
def test_one_ptu_hour_is_one_ptus_input_tpm_served_for_sixty_minutes():
assert ptu_hours(_ROW, _ROW.input_tpm_per_ptu * 60) == pytest.approx(1.0)
assert ptu_hours(_ROW, _ROW.input_tpm_per_ptu * 30) == pytest.approx(0.5)
assert ptu_hours(_ROW, 0.0) == 0.0
@pytest.mark.parametrize(
("deployment", "expected"),
[
({"litellm_params": {"model": "azure/gpt-4.1"}}, True),
({"litellm_params": {"model": "azure_ai/gpt-4.1"}}, True),
({"litellm_params": {"model": "my-ptu-deployment", "custom_llm_provider": "azure"}}, True),
({"litellm_params": {"model": "azure/gpt-4.1", "custom_llm_provider": "openai"}}, False),
({"litellm_params": {"model": "gpt-4.1"}}, False),
({"model_info": {"base_model": "azure/gpt-4.1"}}, False),
],
ids=["azure prefix", "azure_ai prefix", "azure provider", "provider over prefix", "no prefix", "no litellm_params"],
)
def test_a_deployment_is_azure_by_its_provider_first_and_its_model_prefix_otherwise(deployment, expected):
assert is_azure_deployment(deployment) is expected

View file

@ -10,7 +10,9 @@ from collections.abc import Iterator, Sequence
from contextlib import contextmanager
from datetime import datetime, timedelta, timezone
from typing import Any, Dict, Final, List, Optional
from unittest.mock import patch
import httpx
import pytest
from fastapi import HTTPException
from pydantic import TypeAdapter
@ -20,6 +22,9 @@ from litellm import Router
from litellm.caching.caching import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
from litellm.integrations.custom_secret_manager import CustomSecretManager
from litellm.litellm_core_utils import ptu_pricing
from litellm.llms.azure.ptu_capacity import AZURE_PTU_CAPACITY
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
@ -39,14 +44,18 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import (
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
PROXY_MaxParallelRequestsHandler_v3 as _PROXY_MaxParallelRequestsHandler,
)
from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR
from litellm.proxy.utils import InternalUsageCache, ProxyLogging, hash_token
from litellm.router_utils.ptu_shares import PTUTeamCeiling
from litellm.types.caching import RedisPipelineIncrementOperation
from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.types.mcp import MCPPreCallRequestObject, MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPServer
from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem
from litellm.types.utils import (
EmbeddingResponse,
ModelResponse,
PromptTokensDetailsWrapper,
TextCompletionResponse,
Usage,
)
@ -4458,7 +4467,7 @@ async def test_tpm_over_limit_rejection_releases_parallel_slot_v3(monkeypatch):
)
counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests"
async def over_limit_reservation(descriptors, estimated_tokens, parent_otel_span=None):
async def over_limit_reservation(descriptors, estimated_tokens, parent_otel_span=None, **_kwargs):
return {
"overall_code": "OVER_LIMIT",
"statuses": [
@ -4687,8 +4696,8 @@ async def test_async_data_generator_releases_counter_when_wrapped_v3():
from the outer generator: the counter returns to 0 (not -1), proving the
nested hook does not also refund and there is no double decrement.
"""
from litellm.integrations.custom_logger import CustomLogger
import litellm.proxy.proxy_server as proxy_server
from litellm.integrations.custom_logger import CustomLogger
class _PassthroughIteratorOverride(CustomLogger):
async def async_post_call_streaming_iterator_hook(
@ -7681,6 +7690,770 @@ async def test_success_tpm_accounting_keeps_the_admission_target_after_an_alias_
assert not any(":target-b" in key for key in charged)
# --- a team's PTU share on a shared Azure provisioned deployment ---------------------------
def _ptu_ceiling_for(team_id: str, model_group: str, tpm_limit: int, ratio: float, cached_ratio: float = 0.0):
calls: list[tuple[str, str]] = []
def resolve(requested_team: str, requested_group: str) -> PTUTeamCeiling | None:
calls.append((requested_team, requested_group))
if (requested_team, requested_group) != (team_id, model_group):
return None
return PTUTeamCeiling(
model_group=model_group, tpm_limit=tpm_limit, output_to_input_ratio=ratio, cached_input_ratio=cached_ratio
)
return resolve, calls
def _ptu_request(model: str = "test-model") -> dict:
return {"model": model, "messages": [{"role": "user", "content": "hi"}], "max_tokens": 300}
@pytest.mark.asyncio
async def test_a_teams_ptu_share_is_a_hard_tpm_ceiling_on_the_shared_model():
cache = DualCache()
resolve, _ = _ptu_ceiling_for("t", "test-model", tpm_limit=2000, ratio=4.0)
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(cache), ptu_team_ceiling_resolver=resolve
)
key = UserAPIKeyAuth(api_key=hash_token("sk-ptu"), team_id="t")
await handler.async_pre_call_hook(user_api_key_dict=key, cache=cache, data=_ptu_request(), call_type="acompletion")
with pytest.raises(HTTPException) as exc:
await handler.async_pre_call_hook(
user_api_key_dict=key, cache=cache, data=_ptu_request(), call_type="acompletion"
)
assert exc.value.status_code == 429
assert "model_per_team_ptu" in str(exc.value.detail)
assert exc.value.headers["rate_limit_type"] == "tokens"
@pytest.mark.asyncio
async def test_a_ptu_ceiling_on_one_model_leaves_the_teams_other_models_alone():
cache = DualCache()
resolve, calls = _ptu_ceiling_for("t", "test-model", tpm_limit=500, ratio=4.0)
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(cache), ptu_team_ceiling_resolver=resolve
)
key = UserAPIKeyAuth(api_key=hash_token("sk-ptu"), team_id="t")
for _ in range(3):
await handler.async_pre_call_hook(
user_api_key_dict=key, cache=cache, data=_ptu_request("other-model"), call_type="acompletion"
)
assert calls == [("t", "other-model")] * 3
@pytest.mark.asyncio
async def test_a_team_without_a_share_and_a_key_without_a_team_get_no_ptu_ceiling():
cache = DualCache()
resolve, calls = _ptu_ceiling_for("t", "test-model", tpm_limit=500, ratio=4.0)
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(cache), ptu_team_ceiling_resolver=resolve
)
other_team = UserAPIKeyAuth(api_key=hash_token("sk-other"), team_id="u")
no_team = UserAPIKeyAuth(api_key=hash_token("sk-no-team"))
for _ in range(3):
for caller in (other_team, no_team):
await handler.async_pre_call_hook(
user_api_key_dict=caller, cache=cache, data=_ptu_request(), call_type="acompletion"
)
assert set(calls) == {("u", "test-model")}
def _shared_ptu_router(model_group: str) -> Router:
return Router(
model_list=[
{
"model_name": model_group,
"litellm_params": {"model": "azure/gpt-4.1", "api_key": "sk-ptu", "api_base": "https://ptu.example"},
"model_info": {
"id": "shared-ptu",
"base_model": "azure/gpt-4.1",
"ptu_count": 1,
"cost_per_ptu_per_hour": 1.0,
"ptu_effective_from": "2026-01-01T00:00:00Z",
"ptu_shares": {"t": 1},
},
}
],
model_group_alias={f"{model_group}-alias": model_group},
)
def _two_thirds_of_a_ptu_minute() -> dict:
"""An output budget worth two thirds of a gpt-4.1 PTU minute once weighted at the model's
output-to-input ratio, the way the ceiling counts it."""
capacity = AZURE_PTU_CAPACITY["gpt-4.1"]
return {**_ptu_request(), "max_tokens": int(capacity.input_tpm_per_ptu * 2 / 3 / capacity.output_to_input_ratio)}
@pytest.mark.asyncio
async def test_the_reservation_weighs_output_the_way_the_ceiling_does():
"""300 output tokens are 1200 normalized tokens at 4:1, more than a 1000-token ceiling the raw
301-token estimate would clear, so the request is refused as one no minute can admit; the same
request at 1:1 is admitted."""
key = UserAPIKeyAuth(api_key=hash_token("sk-ptu"), team_id="t")
unweighted_cache = DualCache()
unweighted = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(unweighted_cache),
ptu_team_ceiling_resolver=_ptu_ceiling_for("t", "test-model", tpm_limit=1000, ratio=1.0)[0],
)
await unweighted.async_pre_call_hook(
user_api_key_dict=key, cache=unweighted_cache, data=_ptu_request(), call_type="acompletion"
)
weighted_cache = DualCache()
weighted = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(weighted_cache),
ptu_team_ceiling_resolver=_ptu_ceiling_for("t", "test-model", tpm_limit=1000, ratio=4.0)[0],
)
with pytest.raises(HTTPException) as rejected:
await weighted.async_pre_call_hook(
user_api_key_dict=key, cache=weighted_cache, data=_ptu_request(), call_type="acompletion"
)
assert rejected.value.status_code == 400
assert "needs 1201 normalized tokens" in str(rejected.value.detail)
@pytest.mark.asyncio
async def test_a_one_ptu_share_admits_four_uncapped_requests_a_minute_and_rejects_the_fifth():
"""One PTU of gpt-4.1 is 3,000 normalized tokens a minute. A request without max_tokens gets the
proxy's output floor measured against the share in output tokens (750), so its cap costs at most
a quarter of the share once weighted 4:1 and four such requests fit before the fifth is refused."""
cache = DualCache()
resolve, _ = _ptu_ceiling_for("t", "test-model", tpm_limit=3000, ratio=4.0)
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(cache), ptu_team_ceiling_resolver=resolve
)
key = UserAPIKeyAuth(api_key=hash_token("sk-ptu"), team_id="t")
uncapped = {"model": "test-model", "messages": [{"role": "user", "content": "hi"}]}
admitted = [dict(uncapped) for _ in range(4)]
for data in admitted:
await handler.async_pre_call_hook(user_api_key_dict=key, cache=cache, data=data, call_type="acompletion")
with pytest.raises(HTTPException) as fifth:
await handler.async_pre_call_hook(
user_api_key_dict=key, cache=cache, data=dict(uncapped), call_type="acompletion"
)
assert fifth.value.status_code == 429
assert "model_per_team_ptu" in str(fifth.value.detail)
assert all(data["max_tokens"] * 4 <= 3000 // 4 for data in admitted)
@pytest.mark.asyncio
async def test_a_request_whose_own_cap_outgrows_the_share_is_refused_as_invalid_and_charges_nothing():
"""A 429 would tell the client to retry a call no empty minute can ever admit."""
cache = DualCache()
resolve, _ = _ptu_ceiling_for("t", "test-model", tpm_limit=3000, ratio=4.0)
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(cache), ptu_team_ceiling_resolver=resolve
)
key = UserAPIKeyAuth(api_key=hash_token("sk-ptu"), team_id="t")
oversized = {**_ptu_request(), "max_tokens": 750}
with pytest.raises(HTTPException) as refused:
await handler.async_pre_call_hook(user_api_key_dict=key, cache=cache, data=oversized, call_type="acompletion")
for _ in range(2):
await handler.async_pre_call_hook(
user_api_key_dict=key, cache=cache, data={**_ptu_request(), "max_tokens": 360}, call_type="acompletion"
)
assert refused.value.status_code == 400
assert "max_tokens" in str(refused.value.detail)
assert refused.value.headers is None or "retry-after" not in refused.value.headers
@pytest.mark.asyncio
async def test_an_uncapped_prompt_longer_than_the_output_floor_reserves_the_cap_the_proxy_writes():
"""Without a cap the output budget defaults to the prompt's own length, so a long prompt would be
reserved twice, weighted 4:1, and refused on an empty window. The ceiling counts the cap the proxy
writes into the request instead, the same output the deployment can produce."""
cache = DualCache()
resolve, _ = _ptu_ceiling_for("t", "test-model", tpm_limit=3000, ratio=4.0)
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(cache), ptu_team_ceiling_resolver=resolve
)
key = UserAPIKeyAuth(api_key=hash_token("sk-ptu"), team_id="t")
data = {"model": "test-model", "messages": [{"role": "user", "content": "word " * 1200}]}
await handler.async_pre_call_hook(user_api_key_dict=key, cache=cache, data=data, call_type="acompletion")
stash = get_request_stash()
assert stash is not None
assert data["max_tokens"] * 4 <= 3000 // 4
assert stash.ptu_reserved_tokens == stash.reserved_tokens + 3 * data["max_tokens"]
assert stash.ptu_reserved_tokens <= 3000
@pytest.mark.asyncio
async def test_the_ptu_counter_holds_the_normalized_reservation_beside_the_raw_one():
cache = DualCache()
resolve, _ = _ptu_ceiling_for("t", "test-model", tpm_limit=2000, ratio=4.0)
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(cache), ptu_team_ceiling_resolver=resolve
)
key = UserAPIKeyAuth(api_key=hash_token("sk-ptu"), team_id="t")
await handler.async_pre_call_hook(user_api_key_dict=key, cache=cache, data=_ptu_request(), call_type="acompletion")
stash = get_request_stash()
assert stash is not None
assert stash.ptu_reserved_tokens == stash.reserved_tokens + 3 * 300
ptu_key = handler.create_rate_limit_keys("model_per_team_ptu", "t:test-model", "tokens")
assert int(await cache.async_get_cache(key=ptu_key) or 0) == stash.ptu_reserved_tokens
@pytest.mark.asyncio
async def test_the_proxy_router_turns_a_teams_share_into_its_ceiling_when_attribution_is_on(monkeypatch):
"""With no resolver injected the ceiling comes from the proxy router's own deployments: one
PTU of gpt-4.1 a minute, so two requests each reserving two thirds of it are one too many."""
monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true")
cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache))
key = UserAPIKeyAuth(api_key=hash_token("sk-ptu"), team_id="t")
with patch("litellm.proxy.proxy_server.llm_router", _shared_ptu_router("test-model")):
await handler.async_pre_call_hook(
user_api_key_dict=key, cache=cache, data=_two_thirds_of_a_ptu_minute(), call_type="acompletion"
)
with pytest.raises(HTTPException) as exc:
await handler.async_pre_call_hook(
user_api_key_dict=key, cache=cache, data=_two_thirds_of_a_ptu_minute(), call_type="acompletion"
)
assert exc.value.status_code == 429
assert "model_per_team_ptu" in str(exc.value.detail)
class _SecretManagerHoldingThePTUFlag(CustomSecretManager):
def __init__(self) -> None:
super().__init__(secret_manager_name="ptu_flag_holder")
self.flag_reads = 0
async def async_read_secret(
self,
secret_name: str,
optional_params: dict | None = None,
timeout: float | httpx.Timeout | None = None,
) -> str | None:
return self.sync_read_secret(secret_name, optional_params, timeout)
def sync_read_secret(
self,
secret_name: str,
optional_params: dict | None = None,
timeout: float | httpx.Timeout | None = None,
) -> str | None:
if secret_name != PTU_COST_ATTRIBUTION_ENV_VAR:
return None
self.flag_reads += 1
return "true"
@pytest.fixture
def ptu_flag_in_secret_manager(monkeypatch: pytest.MonkeyPatch) -> Iterator[_SecretManagerHoldingThePTUFlag]:
manager = _SecretManagerHoldingThePTUFlag()
monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False)
monkeypatch.setattr(litellm, "secret_manager_client", manager)
monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.CUSTOM)
monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode="read_only"))
ptu_pricing._enabled_in_secret_manager.cache_clear()
yield manager
ptu_pricing._enabled_in_secret_manager.cache_clear()
@pytest.mark.asyncio
@pytest.mark.parametrize(("team_id", "reads_the_flag"), [("u", False), ("t", True)], ids=["no-share", "share-holder"])
async def test_only_a_share_holders_request_asks_the_secret_manager_for_the_ptu_flag(
ptu_flag_in_secret_manager: _SecretManagerHoldingThePTUFlag, team_id: str, reads_the_flag: bool
):
"""A proxy keeping its secrets in a secret manager pays for a flag read only on the requests a
PTU share could limit, so every other team's traffic never reaches the manager."""
router: Final = _shared_ptu_router("test-model")
ptu_pricing._enabled_in_secret_manager.cache_clear()
reads_at_boot: Final = ptu_flag_in_secret_manager.flag_reads
cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache))
key = UserAPIKeyAuth(api_key=hash_token("sk-ptu"), team_id=team_id)
with patch("litellm.proxy.proxy_server.llm_router", router):
await handler.async_pre_call_hook(
user_api_key_dict=key, cache=cache, data=_ptu_request(), call_type="acompletion"
)
assert (ptu_flag_in_secret_manager.flag_reads > reads_at_boot) is reads_the_flag
@pytest.mark.asyncio
@pytest.mark.parametrize("deployment_name", ["shared-ptu", "azure/gpt-4.1"])
async def test_naming_the_shared_deployment_directly_draws_on_the_same_ceiling_as_its_group(
monkeypatch, deployment_name
):
"""The router also serves a deployment named by its id or its provider model, so a team that
spent its share by group name cannot keep going under the deployment's other names, even
though the router lists the deployment's alias copy ahead of it."""
monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true")
cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache))
key = UserAPIKeyAuth(api_key=hash_token("sk-ptu"), team_id="t")
with patch("litellm.proxy.proxy_server.llm_router", _shared_ptu_router("test-model")):
await handler.async_pre_call_hook(
user_api_key_dict=key, cache=cache, data=_two_thirds_of_a_ptu_minute(), call_type="acompletion"
)
with pytest.raises(HTTPException) as exc:
await handler.async_pre_call_hook(
user_api_key_dict=key,
cache=cache,
data={**_two_thirds_of_a_ptu_minute(), "model": deployment_name},
call_type="acompletion",
)
assert exc.value.status_code == 429
assert "model_per_team_ptu" in str(exc.value.detail)
ptu_keys = [cache_key for cache_key in cache.in_memory_cache.cache_dict if "model_per_team_ptu" in cache_key]
assert handler.create_rate_limit_keys("model_per_team_ptu", "t:test-model", "tokens") in ptu_keys
assert all(cache_key.startswith("{model_per_team_ptu:t:test-model}") for cache_key in ptu_keys)
@pytest.mark.asyncio
async def test_naming_the_shared_deployment_by_id_draws_on_its_ceiling_beside_a_catch_all_wildcard(monkeypatch):
"""The router serves a deployment id before any wildcard, so a proxy-wide `*` route that also
matches the id string does not let a team that spent its share keep going under the id."""
monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true")
cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache))
key = UserAPIKeyAuth(api_key=hash_token("sk-ptu"), team_id="t")
catch_all = {
"model_name": "*",
"litellm_params": {"model": "azure/*", "api_key": "sk-open", "api_base": "https://open.example"},
"model_info": {"id": "open-catch-all"},
}
router = Router(model_list=[*_shared_ptu_router("test-model").model_list, catch_all])
with patch("litellm.proxy.proxy_server.llm_router", router):
await handler.async_pre_call_hook(
user_api_key_dict=key, cache=cache, data=_two_thirds_of_a_ptu_minute(), call_type="acompletion"
)
with pytest.raises(HTTPException) as exc:
await handler.async_pre_call_hook(
user_api_key_dict=key,
cache=cache,
data={**_two_thirds_of_a_ptu_minute(), "model": "shared-ptu"},
call_type="acompletion",
)
assert exc.value.status_code == 429
assert "model_per_team_ptu" in str(exc.value.detail)
@pytest.mark.asyncio
async def test_every_name_a_wildcard_route_serves_draws_on_its_shared_deployments_ceiling(monkeypatch):
"""A shared deployment listed under a wildcard serves every name the pattern matches, so a
team that spent its share under one of those names cannot keep going under another."""
monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true")
cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache))
key = UserAPIKeyAuth(api_key=hash_token("sk-ptu"), team_id="t")
with patch("litellm.proxy.proxy_server.llm_router", _shared_ptu_router("ptu-*")):
await handler.async_pre_call_hook(
user_api_key_dict=key,
cache=cache,
data={**_two_thirds_of_a_ptu_minute(), "model": "ptu-chat"},
call_type="acompletion",
)
with pytest.raises(HTTPException) as exc:
await handler.async_pre_call_hook(
user_api_key_dict=key,
cache=cache,
data={**_two_thirds_of_a_ptu_minute(), "model": "ptu-summarize"},
call_type="acompletion",
)
assert exc.value.status_code == 429
assert "model_per_team_ptu" in str(exc.value.detail)
@pytest.mark.asyncio
async def test_the_proxy_router_sets_no_ceiling_while_attribution_is_off(monkeypatch):
monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False)
cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache))
key = UserAPIKeyAuth(api_key=hash_token("sk-ptu"), team_id="t")
with patch("litellm.proxy.proxy_server.llm_router", _shared_ptu_router("test-model")):
for _ in range(3):
await handler.async_pre_call_hook(
user_api_key_dict=key, cache=cache, data=_two_thirds_of_a_ptu_minute(), call_type="acompletion"
)
assert not any("model_per_team_ptu" in cache_key for cache_key in cache.in_memory_cache.cache_dict)
def _ptu_success_kwargs() -> dict:
return {
"standard_logging_object": {
"metadata": {"user_api_key_hash": hash_token("sk-ptu"), "user_api_key_team_id": "t"}
},
"litellm_params": {
"metadata": {"model_group": "test-model", "user_api_key_metadata": {}, "user_api_key_team_metadata": {}}
},
"model": "test-model",
}
def _ptu_response(usage: Usage) -> ModelResponse:
return ModelResponse(
id="ptu-share",
object="chat.completion",
created=1_760_000_000,
model="test-model",
usage=usage,
choices=[],
)
def _ptu_increment(handler: _PROXY_MaxParallelRequestsHandler, ops: list) -> float | None:
ptu_key = handler.create_rate_limit_keys("model_per_team_ptu", "t:test-model", "tokens")
return next((op["increment_value"] for op in ops if op["key"] == ptu_key), None)
def test_success_accounting_settles_the_ptu_counter_in_azure_normalized_tokens():
"""100 uncached input tokens plus 50 output tokens at a 4:1 output ratio are 300
normalized tokens, which is what the share was sized in."""
resolve, _ = _ptu_ceiling_for("t", "test-model", tpm_limit=500, ratio=4.0)
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(DualCache()), ptu_team_ceiling_resolver=resolve
)
response = _ptu_response(Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150))
ops = handler._build_success_event_pipeline_operations(
kwargs=_ptu_success_kwargs(), response_obj=response, rate_limit_type="output"
)
assert _ptu_increment(handler, ops) == 300
def test_pass_through_usage_settles_the_ptu_counter_in_azure_normalized_tokens():
"""A pass-through response carries its usage beside the body rather than in it, and its 100
input and 50 output tokens still settle at 300 normalized tokens, not a raw count."""
resolve, _ = _ptu_ceiling_for("t", "test-model", tpm_limit=500, ratio=4.0)
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(DualCache()), ptu_team_ceiling_resolver=resolve
)
ops = handler._build_success_event_pipeline_operations(
kwargs={
**_ptu_success_kwargs(),
"combined_usage_object": Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150),
},
response_obj={"response": "upstream body was never parsed"},
rate_limit_type="output",
)
assert _ptu_increment(handler, ops) == 300
def test_cached_input_is_not_charged_to_the_ptu_counter():
resolve, _ = _ptu_ceiling_for("t", "test-model", tpm_limit=500, ratio=4.0)
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(DualCache()), ptu_team_ceiling_resolver=resolve
)
response = _ptu_response(
Usage(
prompt_tokens=100,
completion_tokens=50,
total_tokens=150,
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=40),
)
)
ops = handler._build_success_event_pipeline_operations(
kwargs=_ptu_success_kwargs(), response_obj=response, rate_limit_type="output"
)
assert _ptu_increment(handler, ops) == 60 + 4 * 50
def test_cached_input_is_charged_at_the_models_cached_ratio():
"""40 of the 100 input tokens were cache reads; at a tenth each they are 4 normalized tokens
beside the 60 uncached ones and the 200 for 50 outputs at 4:1."""
resolve, _ = _ptu_ceiling_for("t", "test-model", tpm_limit=500, ratio=4.0, cached_ratio=0.1)
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(DualCache()), ptu_team_ceiling_resolver=resolve
)
response = _ptu_response(
Usage(
prompt_tokens=100,
completion_tokens=50,
total_tokens=150,
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=40),
)
)
ops = handler._build_success_event_pipeline_operations(
kwargs=_ptu_success_kwargs(), response_obj=response, rate_limit_type="output"
)
assert _ptu_increment(handler, ops) == 60 + 4 + 4 * 50
@pytest.mark.asyncio
async def test_a_reservation_is_settled_even_after_the_teams_share_is_gone():
"""The share can be removed between admission and completion; the reserved tokens still
come off the counter instead of standing in the window."""
ceiling: dict[str, PTUTeamCeiling | None] = {
"current": PTUTeamCeiling(
model_group="test-model", tpm_limit=2000, output_to_input_ratio=4.0, cached_input_ratio=0.0
)
}
cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(cache),
ptu_team_ceiling_resolver=lambda _team, _group: ceiling["current"],
)
key = UserAPIKeyAuth(api_key=hash_token("sk-ptu"), team_id="t")
await handler.async_pre_call_hook(user_api_key_dict=key, cache=cache, data=_ptu_request(), call_type="acompletion")
stash = get_request_stash()
assert stash is not None
assert ("model_per_team_ptu", "t:test-model") in stash.reserved_scopes
assert stash.ptu_reserved_tokens > stash.reserved_tokens > 150
ceiling["current"] = None
ops = handler._build_success_event_pipeline_operations(
kwargs=_ptu_success_kwargs(),
response_obj=_ptu_response(Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150)),
rate_limit_type="total",
)
assert _ptu_increment(handler, ops) == 300 - stash.ptu_reserved_tokens
@pytest.mark.asyncio
async def test_a_reservation_is_settled_when_the_success_metadata_lost_the_team():
"""Admission reserved on the team's PTU counter, so settlement credits that scope even when the
success event's logging metadata carries no team id, instead of leaving the reservation standing."""
cache = DualCache()
resolve, _ = _ptu_ceiling_for("t", "test-model", tpm_limit=2000, ratio=4.0)
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(cache), ptu_team_ceiling_resolver=resolve
)
key = UserAPIKeyAuth(api_key=hash_token("sk-ptu"), team_id="t")
await handler.async_pre_call_hook(user_api_key_dict=key, cache=cache, data=_ptu_request(), call_type="acompletion")
stash = get_request_stash()
assert stash is not None
assert stash.ptu_reserved_tokens > 300
ops = handler._build_success_event_pipeline_operations(
kwargs={
**_ptu_success_kwargs(),
"standard_logging_object": {"metadata": {"user_api_key_hash": hash_token("sk-ptu")}},
},
response_obj=_ptu_response(Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150)),
rate_limit_type="total",
)
assert _ptu_increment(handler, ops) == 300 - stash.ptu_reserved_tokens
@pytest.mark.asyncio
async def test_success_settles_the_scope_the_reservation_was_taken_on():
"""Settlement credits the scope admission reserved, not one rebuilt from the name the
request used, so a request naming the deployment by id cannot leave its reservation standing."""
ceiling: dict[str, PTUTeamCeiling | None] = {
"current": PTUTeamCeiling(
model_group="test-model", tpm_limit=2000, output_to_input_ratio=4.0, cached_input_ratio=0.0
)
}
cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(cache),
ptu_team_ceiling_resolver=lambda _team, _group: ceiling["current"],
)
key = UserAPIKeyAuth(api_key=hash_token("sk-ptu"), team_id="t")
await handler.async_pre_call_hook(
user_api_key_dict=key, cache=cache, data=_ptu_request("shared-ptu"), call_type="acompletion"
)
stash = get_request_stash()
assert stash is not None
assert ("model_per_team_ptu", "t:test-model") in stash.reserved_scopes
ceiling["current"] = None
stash.ptu_ceiling = None
kwargs = _ptu_success_kwargs()
kwargs["litellm_params"]["metadata"]["model_group"] = "shared-ptu"
ops = handler._build_success_event_pipeline_operations(
kwargs=kwargs,
response_obj=_ptu_response(Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150)),
rate_limit_type="total",
)
assert _ptu_increment(handler, ops) == 150 - stash.ptu_reserved_tokens
async def _reserve_a_ptu_minute(
cache: DualCache, call_id: str, clock: list[datetime] | None = None
) -> tuple[_PROXY_MaxParallelRequestsHandler, str]:
resolve, _ = _ptu_ceiling_for("t", "test-model", tpm_limit=2000, ratio=4.0)
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(cache),
ptu_team_ceiling_resolver=resolve,
time_provider=(lambda: clock[-1]) if clock is not None else None,
)
key = UserAPIKeyAuth(api_key=hash_token("sk-ptu"), team_id="t")
await handler.async_pre_call_hook(
user_api_key_dict=key,
cache=cache,
data={**_ptu_request(), "litellm_call_id": call_id},
call_type="acompletion",
)
stash = get_request_stash()
assert stash is not None and stash.ptu_reserved_tokens > 0
return handler, handler.create_rate_limit_keys("model_per_team_ptu", "t:test-model", "tokens")
@pytest.mark.asyncio
async def test_a_failed_stream_settles_the_ptu_counter_in_normalized_tokens():
"""The partial usage a failed stream recovered is 20 input and 7 output tokens: 48 normalized
at 4:1, which is what stays in the window instead of the raw 27 or the whole reservation."""
cache = DualCache()
handler, ptu_key = await _reserve_a_ptu_minute(cache, "ptu-partial")
await handler.async_log_failure_event(
kwargs={
"litellm_call_id": "ptu-partial",
"standard_logging_object": {
"metadata": {"user_api_key_hash": hash_token("sk-ptu"), "user_api_key_team_id": "t"}
},
"combined_usage_object": Usage(prompt_tokens=20, completion_tokens=7, total_tokens=27),
},
response_obj=None,
start_time=None,
end_time=None,
)
assert int(await cache.async_get_cache(key=ptu_key) or 0) == 20 + 4 * 7
@pytest.mark.asyncio
async def test_a_proxy_side_rejection_refunds_the_whole_normalized_ptu_reservation():
cache = DualCache()
handler, ptu_key = await _reserve_a_ptu_minute(cache, "ptu-rejected")
await handler.async_post_call_failure_hook(
request_data={**_ptu_request(), "litellm_call_id": "ptu-rejected"},
original_exception=Exception("guardrail rejected the request"),
user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("sk-ptu"), team_id="t"),
)
assert int(await cache.async_get_cache(key=ptu_key) or 0) == 0
async def _settle_a_ptu_request(handler: _PROXY_MaxParallelRequestsHandler, call_id: str, outcome: str) -> None:
team_metadata = {"metadata": {"user_api_key_hash": hash_token("sk-ptu"), "user_api_key_team_id": "t"}}
if outcome == "success":
await handler.async_log_success_event(
kwargs={**_ptu_success_kwargs(), "litellm_call_id": call_id},
response_obj=_ptu_response(Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150)),
start_time=None,
end_time=None,
)
elif outcome == "upstream_failure":
await handler.async_log_failure_event(
kwargs={"litellm_call_id": call_id, "standard_logging_object": team_metadata},
response_obj=None,
start_time=None,
end_time=None,
)
else:
await handler.async_post_call_failure_hook(
request_data={**_ptu_request(), "litellm_call_id": call_id},
original_exception=Exception("guardrail rejected the request"),
user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("sk-ptu"), team_id="t"),
)
@pytest.mark.asyncio
@pytest.mark.parametrize("outcome", ["success", "upstream_failure", "proxy_rejection"])
async def test_a_ptu_reservation_settled_after_its_minute_rolled_leaves_the_next_minute_alone(outcome: str):
"""A request admitted in one minute that settles after the next minute opened must not move
the new minute's PTU counter: its reservation already expired with the window it was taken in."""
cache = DualCache()
clock = [datetime(2026, 10, 14, 12, 0, 0)]
handler, ptu_key = await _reserve_a_ptu_minute(cache, "ptu-late", clock)
clock.append(clock[-1] + timedelta(seconds=handler.window_size + 1))
await handler.atomic_check_and_increment_by_n(
descriptors=[
{
"key": "model_per_team_ptu",
"value": "t:test-model",
"rate_limit": {"tokens_per_unit": 2000, "window_size": handler.window_size},
}
],
increments=[{"tokens": 20}],
)
await _settle_a_ptu_request(handler, "ptu-late", outcome)
assert int(await cache.async_get_cache(key=ptu_key) or 0) == 20
@pytest.mark.asyncio
@pytest.mark.parametrize(
("outcome", "settled_tokens"), [("success", 100 + 4 * 50), ("upstream_failure", 0), ("proxy_rejection", 0)]
)
async def test_a_ptu_reservation_settled_inside_its_minute_lands_at_the_normalized_usage(
outcome: str, settled_tokens: int
):
cache = DualCache()
handler, ptu_key = await _reserve_a_ptu_minute(cache, "ptu-on-time")
await _settle_a_ptu_request(handler, "ptu-on-time", outcome)
assert int(await cache.async_get_cache(key=ptu_key) or 0) == settled_tokens
def test_usage_that_only_reports_a_total_charges_that_total_to_the_ptu_counter():
resolve, _ = _ptu_ceiling_for("t", "test-model", tpm_limit=500, ratio=4.0)
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(DualCache()), ptu_team_ceiling_resolver=resolve
)
response = _ptu_response(Usage(prompt_tokens=0, completion_tokens=0, total_tokens=150))
ops = handler._build_success_event_pipeline_operations(
kwargs=_ptu_success_kwargs(), response_obj=response, rate_limit_type="output"
)
assert _ptu_increment(handler, ops) == 150
def test_success_accounting_charges_no_ptu_counter_without_a_ceiling():
resolve, _ = _ptu_ceiling_for("u", "test-model", tpm_limit=500, ratio=4.0)
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(DualCache()), ptu_team_ceiling_resolver=resolve
)
response = _ptu_response(Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150))
ops = handler._build_success_event_pipeline_operations(
kwargs=_ptu_success_kwargs(), response_obj=response, rate_limit_type="output"
)
assert _ptu_increment(handler, ops) is None
@pytest.mark.parametrize("self_call", [False, True])
async def test_managed_invocations_enforce_actor_and_target_rate_policies(
monkeypatch: pytest.MonkeyPatch, self_call: bool

View file

@ -3006,6 +3006,23 @@ def test_zero_usage_keeps_reservations_unless_measured_fallback_exists(
assert tuple(operation["increment_value"] for operation in operations) == expected_increments
def test_a_responses_reply_reporting_no_tokens_keeps_the_reservations(rate_limiter):
handler, _cache = rate_limiter
stash = get_or_create_request_stash()
stash.itpm_reserved_tokens = 100
stash.itpm_reserved_scopes = frozenset({(PROJECT_ITPM_DESCRIPTOR_KEY, "project:model")})
stash.otpm_reserved_tokens = 60
stash.otpm_reserved_scopes = frozenset({(PROJECT_OTPM_DESCRIPTOR_KEY, "project:model")})
response_obj = ResponsesAPIResponse(
id="resp_zero",
created_at=int(datetime.now().timestamp()),
output=[],
usage=ResponseAPIUsage(input_tokens=0, output_tokens=0, total_tokens=0),
)
assert tuple(handler._build_io_token_reservation_ops({}, response_obj)) == ()
@pytest.mark.parametrize(
("usage", "expected_increments"),
[

View file

@ -3083,6 +3083,52 @@ async def test_delete_group_recomputes_roles_for_members(mocker):
assert list(recompute_mock.call_args[0][1]) == ["user1", "user2"]
@pytest.mark.asyncio
async def test_delete_group_is_refused_while_the_team_holds_a_ptu_share(mocker, monkeypatch):
"""A SCIM group is a team, so an identity provider deleting one that still holds a PTU share is
refused like /team/delete, with the team and its members' team lists left as they were."""
from litellm import Router
shared_router: Final = Router(
model_list=[
{
"model_name": "gpt-4.1-ptu",
"litellm_params": {"model": "azure/gpt-4.1", "api_key": "sk-azure", "api_base": "https://x"},
"model_info": {
"id": "shared-ptu",
"ptu_count": 50,
"cost_per_ptu_per_hour": 1.0,
"ptu_effective_from": "2026-01-01T00:00:00Z",
"ptu_shares": {"test-team-123": 30, "team-2": 20},
},
}
]
)
existing_team = mocker.MagicMock()
existing_team.members_with_roles = [Member(user_id="user1", role="user")]
member = mocker.MagicMock()
member.teams = ["test-team-123"]
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing_team)
mock_prisma_client.db.litellm_teamtable.delete = AsyncMock()
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=member)
mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=())
mock_prisma_client.db.litellm_usertable.update = AsyncMock()
mocker.patch(
"litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception",
AsyncMock(return_value=mock_prisma_client),
)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", shared_router)
with pytest.raises(ProxyException) as refused:
await delete_group(group_id="test-team-123")
assert refused.value.code == "400"
assert "test-team-123" in str(refused.value.message) and "shared-ptu" in str(refused.value.message)
mock_prisma_client.db.litellm_teamtable.delete.assert_not_called()
mock_prisma_client.db.litellm_usertable.update.assert_not_called()
@pytest.mark.asyncio
async def test_handle_existing_user_by_email_applies_role_when_admin_group_set(mocker):
"""When admin_group is configured, re-upserting an existing email persists the

View file

@ -3,7 +3,7 @@ import io
from collections.abc import AsyncIterator, Iterator, Mapping, Sequence
from dataclasses import dataclass, fields
from itertools import chain
from types import SimpleNamespace
from types import MappingProxyType, SimpleNamespace
from typing import Final
from unittest.mock import AsyncMock
@ -12,14 +12,17 @@ from fastapi import FastAPI, Request
from fastapi.testclient import TestClient
from litellm import constants
from litellm.llms.azure.ptu_capacity import AZURE_PTU_CAPACITY
from litellm.proxy._types import LiteLLM_TeamTable, LiteLLM_UserTable, LitellmUserRoles, Member, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.management_endpoints.daily_activity_routes import (
_csv_cell,
get_daily_activity_llm_router,
get_daily_activity_prisma_client,
get_daily_activity_repository,
router,
)
from litellm.router import Router
from litellm.types.proxy.management_endpoints.common_daily_activity import KeySpendMetrics, SpendMetrics
from litellm.types.repositories.daily_activity import (
AggregatedRows,
@ -193,6 +196,7 @@ def _grouping_row(
api_key: str | None,
group_level: int,
distinct_api_keys: int | None,
model_group: str | None = None,
) -> GroupingSetsRow:
metric_values: Final = _metrics(rows)
return GroupingSetsRow(
@ -200,7 +204,7 @@ def _grouping_row(
api_key=api_key,
**metric_values,
model=None,
model_group=None,
model_group=model_group,
custom_llm_provider=None,
mcp_namespaced_tool_name=None,
endpoint=None,
@ -214,17 +218,29 @@ def _grouping_rows_for_day(
) -> tuple[GroupingSetsRow, ...]:
date_rows: Final = tuple(row for row in rows if row.date == date)
return (
_grouping_row(date_rows, date=date, api_key=None, group_level=63, distinct_api_keys=distinct_api_keys),
) + tuple(
_grouping_row(
tuple(row for row in date_rows if row.api_key == api_key),
date=date,
api_key=api_key,
group_level=31,
distinct_api_keys=None,
(_grouping_row(date_rows, date=date, api_key=None, group_level=63, distinct_api_keys=distinct_api_keys),)
+ tuple(
_grouping_row(
tuple(row for row in date_rows if row.api_key == api_key),
date=date,
api_key=api_key,
group_level=31,
distinct_api_keys=None,
)
for api_key in top_keys
if any(row.api_key == api_key for row in date_rows)
)
+ tuple(
_grouping_row(
tuple(row for row in date_rows if row.model_group == model_group),
date=date,
api_key=None,
group_level=55,
distinct_api_keys=None,
model_group=model_group,
)
for model_group in sorted({row.model_group for row in date_rows})
)
for api_key in top_keys
if any(row.api_key == api_key for row in date_rows)
)
@ -611,6 +627,103 @@ def test_aggregated_routes_return_scoped_results(
assert len(body["results"]) == 2, response.text
def _ptu_sized_router(model_group: str, shares: Mapping[str, int] = MappingProxyType({"team-a": 30, "team-b": 20})) -> Router:
return Router(
model_list=[
{
"model_name": model_group,
"litellm_params": {"model": "azure/gpt-4.1", "api_key": "sk-ptu", "api_base": "https://ptu.example"},
"model_info": {
"id": "shared-ptu",
"base_model": "azure/gpt-4.1",
"ptu_count": 50,
"cost_per_ptu_per_hour": 1.0,
"ptu_effective_from": "2024-01-01T00:00:00Z",
"ptu_shares": dict(shares),
},
}
]
)
def _inject_llm_router(client: TestClient, llm_router: Router) -> None:
app: Final = client.app
assert isinstance(app, FastAPI)
app.dependency_overrides[get_daily_activity_llm_router] = lambda: llm_router
@pytest.mark.parametrize("attribution_enabled", [True, False])
def test_team_aggregated_route_reports_ptu_hours_for_the_sized_model_group_only_while_attribution_is_on(
daily_activity_client: tuple[TestClient, _FakeRepository],
monkeypatch: pytest.MonkeyPatch,
attribution_enabled: bool,
) -> None:
"""team-a sent one rare-group request of 10 input and 5 output tokens on 2025-01-01, so with
gpt-4.1's sizing that day, its model group, and the total carry those tokens normalized
(output at the model's output ratio) over one PTU-hour of input; the unsized popular-group,
the day without rare-group traffic, and the flag-off run all report zero."""
client, _ = daily_activity_client
monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True" if attribution_enabled else "False")
_inject_llm_router(client, _ptu_sized_router("rare-group"))
response: Final = client.get("/team/daily/activity/aggregated", params=_entity_params("team_ids", "team-a"))
assert response.status_code == 200, response.text
body: Final = response.json()
capacity: Final = AZURE_PTU_CAPACITY["gpt-4.1"]
expected: Final = (
(10 + 5 * capacity.output_to_input_ratio) / capacity.normalized_tokens_per_ptu_hour
if attribution_enabled
else 0.0
)
by_date: Final = {day["date"]: day for day in body["results"]}
rare_day: Final = by_date["2025-01-01"]
assert body["metadata"]["total_ptu_hours"] == pytest.approx(expected), response.text
assert rare_day["metrics"]["ptu_hours"] == pytest.approx(expected), response.text
assert rare_day["breakdown"]["model_groups"]["rare-group"]["metrics"]["ptu_hours"] == pytest.approx(expected)
assert rare_day["breakdown"]["model_groups"]["popular-group"]["metrics"]["ptu_hours"] == 0, response.text
assert by_date["2025-01-02"]["metrics"]["ptu_hours"] == 0, response.text
assert body["metadata"]["total_tokens"] == 90, response.text
def test_team_aggregated_route_reports_no_ptu_hours_to_a_team_holding_no_share(
daily_activity_client: tuple[TestClient, _FakeRepository],
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""rare-group is split between team-b and team-c, so team-a's page alone converts none of its
rare-group tokens even though the group is sized."""
client, _ = daily_activity_client
monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True")
_inject_llm_router(client, _ptu_sized_router("rare-group", shares=MappingProxyType({"team-b": 30, "team-c": 20})))
response: Final = client.get("/team/daily/activity/aggregated", params=_entity_params("team_ids", "team-a"))
assert response.status_code == 200, response.text
body: Final = response.json()
rare_day: Final = {day["date"]: day for day in body["results"]}["2025-01-01"]
assert body["metadata"]["total_ptu_hours"] == 0, response.text
assert rare_day["breakdown"]["model_groups"]["rare-group"]["metrics"]["ptu_hours"] == 0, response.text
assert body["metadata"]["total_tokens"] == 90, response.text
def test_ptu_hours_stay_zero_outside_the_team_routes(
daily_activity_client: tuple[TestClient, _FakeRepository], monkeypatch: pytest.MonkeyPatch
) -> None:
"""PTU reservations are split per team, so only the team routes convert tokens to PTU-hours;
a user's own activity on the same sized model group keeps the field at zero."""
client, _ = daily_activity_client
monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True")
_inject_llm_router(client, _ptu_sized_router("rare-group"))
response: Final = client.get("/user/daily/activity/aggregated", params=_entity_params("user_id", "user-a"))
assert response.status_code == 200, response.text
body: Final = response.json()
rare_day: Final = next(day for day in body["results"] if day["date"] == "2025-01-01")
assert body["metadata"]["total_ptu_hours"] == 0, response.text
assert rare_day["breakdown"]["model_groups"]["rare-group"]["metrics"]["ptu_hours"] == 0, response.text
@pytest.mark.parametrize(("prefix", "query_name", "entity_id"), _ENTITY_CASES)
def test_admin_aggregates_all_entities_when_filter_is_omitted(
daily_activity_client: tuple[TestClient, _FakeRepository], prefix: str, query_name: str, entity_id: str
@ -642,6 +755,7 @@ def test_search_folds_each_entity_key_across_days(
"metrics": {
"spend": 3.0,
"flat_cost": 0.0,
"ptu_hours": 0.0,
"prompt_tokens": 20,
"completion_tokens": 10,
"cache_read_input_tokens": 0,

View file

@ -1573,3 +1573,45 @@ async def test_delete_organization_evicts_the_cache_of_the_keys_it_deletes(monke
assert all(cache.get_cache(key=cache_key) is None for cache_key in doomed_cache_keys)
assert all(cache.get_cache(key=cache_key) == {"retained": True} for cache_key in kept_cache_keys)
assert jwt_table.rows == (kept_row,)
@pytest.mark.asyncio
async def test_delete_organization_is_refused_while_one_of_its_teams_holds_a_ptu_share(monkeypatch):
"""Deleting an organization deletes its teams, so a team still holding a PTU share stops it the
way /team/delete does, before any team, member, key, or the organization itself is removed."""
from litellm import Router
from litellm.proxy._types import DeleteOrganizationRequest, LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.management_endpoints.organization_endpoints import delete_organization
shared_router: Final = Router(
model_list=[
{
"model_name": "gpt-4.1-ptu",
"litellm_params": {"model": "azure/gpt-4.1", "api_key": "sk-azure", "api_base": "https://x"},
"model_info": {
"id": "shared-ptu",
"ptu_count": 50,
"cost_per_ptu_per_hour": 1.0,
"ptu_effective_from": "2026-01-01T00:00:00Z",
"ptu_shares": {"team-in-org": 30, "team-elsewhere": 20},
},
}
]
)
prisma_client: Final = AsyncMock()
prisma_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[SimpleNamespace(team_id="team-in-org")])
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True, raising=False)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", shared_router)
with pytest.raises(HTTPException) as refused:
await delete_organization(
data=DeleteOrganizationRequest(organization_ids=["org-with-a-share"]),
user_api_key_dict=UserAPIKeyAuth(api_key="sk-admin", user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert refused.value.status_code == 400
assert "team-in-org" in str(refused.value.detail) and "shared-ptu" in str(refused.value.detail)
prisma_client.db.litellm_teamtable.delete_many.assert_not_called()
prisma_client.db.litellm_verificationtoken.delete_many.assert_not_called()
prisma_client.db.litellm_organizationtable.delete.assert_not_called()

View file

@ -0,0 +1,381 @@
"""Tests for attaching PTU-hours to a daily activity response."""
from datetime import date
from typing import Final
import pytest
from litellm import Router
from litellm.llms.azure.ptu_capacity import AZURE_PTU_CAPACITY, PTUCapacity
from litellm.proxy.management_endpoints.ptu_consumption import attach_ptu_hours, with_ptu_consumption
from litellm.types.proxy.management_endpoints.common_daily_activity import (
BreakdownMetrics,
DailySpendData,
DailySpendMetadata,
KeyMetricWithMetadata,
MetricWithMetadata,
SpendAnalyticsPaginatedResponse,
SpendMetrics,
)
_ROW: Final = PTUCapacity(input_tpm_per_ptu=1_000, output_to_input_ratio=4.0)
_CACHED_ROW: Final = PTUCapacity(input_tpm_per_ptu=1_000, output_to_input_ratio=4.0, cached_input_ratio=0.1)
_CAPACITY: Final = {"gpt-4.1-ptu": _ROW, "gpt-6-ptu": _CACHED_ROW}
_ONE_PTU_HOUR_OF_INPUT: Final = AZURE_PTU_CAPACITY["gpt-4.1"].normalized_tokens_per_ptu_hour
def _capacity_every_day(model_group: str, _day: date) -> PTUCapacity | None:
return _CAPACITY.get(model_group)
def _metrics(prompt: int, completion: int, cached: int = 0) -> SpendMetrics:
return SpendMetrics(
prompt_tokens=prompt,
completion_tokens=completion,
cache_read_input_tokens=cached,
total_tokens=prompt + completion,
api_requests=1,
successful_requests=1,
)
def _bucket(metrics: SpendMetrics, keys: dict[str, SpendMetrics] | None = None) -> MetricWithMetadata:
return MetricWithMetadata(
metrics=metrics,
metadata={},
api_key_breakdown={key: KeyMetricWithMetadata(metrics=m, metadata={}) for key, m in (keys or {}).items()},
)
def _day(date: str, model_groups: dict[str, MetricWithMetadata]) -> DailySpendData:
total: Final = _metrics(
sum(b.metrics.prompt_tokens for b in model_groups.values()),
sum(b.metrics.completion_tokens for b in model_groups.values()),
)
return DailySpendData(date=date, metrics=total, breakdown=BreakdownMetrics(model_groups=model_groups))
def _response(*days: DailySpendData) -> SpendAnalyticsPaginatedResponse:
return SpendAnalyticsPaginatedResponse(results=list(days), metadata=DailySpendMetadata())
def test_a_ptu_model_groups_tokens_become_ptu_hours_on_the_group_the_day_and_the_total():
"""60,000 normalized tokens on a 1,000 input-TPM row is one PTU-hour."""
day: Final = _day("2026-09-23", {"gpt-4.1-ptu": _bucket(_metrics(prompt=40_000, completion=5_000))})
attached: Final = attach_ptu_hours(_response(day), _capacity_every_day)
group: Final = attached.results[0].breakdown.model_groups["gpt-4.1-ptu"]
assert group.metrics.ptu_hours == pytest.approx(1.0)
assert attached.results[0].metrics.ptu_hours == pytest.approx(1.0)
assert attached.metadata.total_ptu_hours == pytest.approx(1.0)
def test_a_model_group_without_a_sizing_row_keeps_zero_ptu_hours_and_the_rest_of_the_day_intact():
day: Final = _day(
"2026-09-23",
{
"gpt-4.1-ptu": _bucket(_metrics(prompt=30_000, completion=0)),
"gpt-4o-mini": _bucket(_metrics(prompt=1_000_000, completion=1_000_000)),
},
)
attached: Final = attach_ptu_hours(_response(day), _capacity_every_day)
groups: Final = attached.results[0].breakdown.model_groups
assert groups["gpt-4o-mini"].metrics.ptu_hours == 0.0
assert groups["gpt-4o-mini"].metrics.prompt_tokens == 1_000_000
assert groups["gpt-4.1-ptu"].metrics.ptu_hours == pytest.approx(0.5)
assert attached.results[0].metrics.ptu_hours == pytest.approx(0.5)
assert attached.results[0].metrics.total_tokens == day.metrics.total_tokens
def test_cached_input_is_charged_at_the_rows_cached_ratio():
day: Final = _day("2026-09-23", {"gpt-6-ptu": _bucket(_metrics(prompt=60_000, completion=0, cached=60_000))})
attached: Final = attach_ptu_hours(_response(day), _capacity_every_day)
assert attached.results[0].breakdown.model_groups["gpt-6-ptu"].metrics.ptu_hours == pytest.approx(0.1)
def test_each_api_key_under_a_ptu_group_gets_its_own_ptu_hours():
day: Final = _day(
"2026-09-23",
{
"gpt-4.1-ptu": _bucket(
_metrics(prompt=60_000, completion=0),
keys={"key-a": _metrics(prompt=45_000, completion=0), "key-b": _metrics(prompt=15_000, completion=0)},
)
},
)
attached: Final = attach_ptu_hours(_response(day), _capacity_every_day)
breakdown: Final = attached.results[0].breakdown.model_groups["gpt-4.1-ptu"].api_key_breakdown
assert breakdown["key-a"].metrics.ptu_hours == pytest.approx(0.75)
assert breakdown["key-b"].metrics.ptu_hours == pytest.approx(0.25)
def test_the_total_sums_every_day_on_the_page():
first: Final = _day("2026-09-22", {"gpt-4.1-ptu": _bucket(_metrics(prompt=60_000, completion=0))})
second: Final = _day("2026-09-23", {"gpt-4.1-ptu": _bucket(_metrics(prompt=0, completion=15_000))})
attached: Final = attach_ptu_hours(_response(first, second), _capacity_every_day)
assert [day.metrics.ptu_hours for day in attached.results] == [pytest.approx(1.0), pytest.approx(1.0)]
assert attached.metadata.total_ptu_hours == pytest.approx(2.0)
def test_a_page_with_no_ptu_group_is_returned_unchanged():
day: Final = _day("2026-09-23", {"gpt-4o-mini": _bucket(_metrics(prompt=100, completion=100))})
response: Final = _response(day)
attached: Final = attach_ptu_hours(response, _capacity_every_day)
assert attached.results[0] is day
assert attached.metadata.total_ptu_hours == 0.0
assert response.metadata.total_ptu_hours == 0.0
def _shared_ptu_router() -> Router:
return Router(
model_list=[
{
"model_name": "gpt-4.1-ptu",
"litellm_params": {"model": "azure/gpt-4.1", "api_key": "sk-ptu", "api_base": "https://ptu.example"},
"model_info": {
"id": "shared-ptu",
"base_model": "azure/gpt-4.1",
"ptu_count": 50,
"cost_per_ptu_per_hour": 1.0,
"ptu_effective_from": "2026-01-01T00:00:00Z",
"ptu_shares": {"team-a": 30, "team-b": 20},
},
}
],
model_group_alias={"ptu-alias": {"model": "gpt-4.1-ptu", "hidden": True}},
)
def test_rows_keyed_by_an_alias_a_deployment_id_or_a_provider_model_are_sized_like_the_group(monkeypatch):
"""The ceiling charges a request however it names the shared deployment, so the usage row
that request lands in, keyed by the name it used, reports the same PTU-hours as the group."""
monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True")
names: Final = ("gpt-4.1-ptu", "ptu-alias", "shared-ptu", "azure/gpt-4.1")
one_hour_each: Final = {name: _bucket(_metrics(prompt=_ONE_PTU_HOUR_OF_INPUT, completion=0)) for name in names}
attached: Final = with_ptu_consumption(_response(_day("2026-09-23", one_hour_each)), _shared_ptu_router(), "team-a")
groups: Final = attached.results[0].breakdown.model_groups
assert [groups[name].metrics.ptu_hours for name in names] == [pytest.approx(1.0)] * len(names)
assert attached.results[0].metrics.ptu_hours == pytest.approx(float(len(names)))
assert attached.metadata.total_ptu_hours == pytest.approx(float(len(names)))
@pytest.mark.parametrize("team_id", ["team-a", None], ids=["one-team", "all-teams"])
def test_only_the_days_a_reservation_covers_read_ptu_hours(monkeypatch, team_id):
"""Tokens a team sent before its reservation started or after it ended were pay-as-you-go,
so only the days inside the window convert to PTU-hours."""
monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True")
router: Final = Router(
model_list=[
{
"model_name": "gpt-4.1-ptu",
"litellm_params": {"model": "azure/gpt-4.1", "api_key": "sk-ptu", "api_base": "https://ptu.example"},
"model_info": {
"id": "shared-ptu",
"base_model": "azure/gpt-4.1",
"ptu_count": 50,
"cost_per_ptu_per_hour": 1.0,
"ptu_effective_from": "2026-09-23T12:00:00Z",
"ptu_effective_to": "2026-09-24T00:00:00Z",
"ptu_shares": {"team-a": 30, "team-b": 20},
},
}
]
)
days: Final = tuple(
_day(day, {"gpt-4.1-ptu": _bucket(_metrics(prompt=_ONE_PTU_HOUR_OF_INPUT, completion=0))})
for day in ("2026-09-22", "2026-09-23", "2026-09-24")
)
attached: Final = with_ptu_consumption(_response(*days), router, team_id)
assert [day.metrics.ptu_hours for day in attached.results] == [0.0, pytest.approx(1.0), 0.0]
assert attached.metadata.total_ptu_hours == pytest.approx(1.0)
def test_rows_keyed_by_a_name_a_wildcard_route_serves_are_sized_by_its_shared_deployment(monkeypatch):
"""A request the router served through a wildcard lands in a row keyed by the name it used,
which reports the PTU-hours of the shared deployment behind the pattern."""
monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True")
router: Final = Router(
model_list=[
{
"model_name": "ptu-*",
"litellm_params": {"model": "azure/gpt-4.1", "api_key": "sk-ptu", "api_base": "https://ptu.example"},
"model_info": {
"id": "shared-ptu",
"base_model": "azure/gpt-4.1",
"ptu_count": 50,
"cost_per_ptu_per_hour": 1.0,
"ptu_effective_from": "2026-01-01T00:00:00Z",
"ptu_shares": {"team-a": 30, "team-b": 20},
},
}
]
)
one_hour: Final = _bucket(_metrics(prompt=_ONE_PTU_HOUR_OF_INPUT, completion=0))
attached: Final = with_ptu_consumption(_response(_day("2026-09-23", {"ptu-chat": one_hour})), router, "team-a")
assert attached.results[0].breakdown.model_groups["ptu-chat"].metrics.ptu_hours == pytest.approx(1.0)
@pytest.mark.parametrize(
("catch_all_routes", "expected_ptu_hours"), [(False, 1.0), (True, 0.0)], ids=["team-wildcard", "behind-catch-all"]
)
def test_a_teams_wildcard_reservation_reads_ptu_hours_only_while_the_router_serves_from_it(
monkeypatch, catch_all_routes, expected_ptu_hours
):
"""The router tries the proxy-wide wildcards before a team's own, so a team's `ptu-*`
reservation behind an open `*` route never serves `ptu-chat`, and its row reads no PTU-hours."""
monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True")
reservation: Final = {
"model_name": "ptu-team-r",
"litellm_params": {"model": "azure/gpt-4.1", "api_key": "sk-r", "api_base": "https://r.example"},
"model_info": {
"id": "reserved-r",
"base_model": "azure/gpt-4.1",
"team_id": "team-r",
"team_public_model_name": "ptu-*",
"ptu_count": 50,
"cost_per_ptu_per_hour": 1.0,
"ptu_effective_from": "2026-01-01T00:00:00Z",
},
}
catch_all: Final = {
"model_name": "*",
"litellm_params": {"model": "azure/*", "api_key": "sk-open", "api_base": "https://open.example"},
"model_info": {"id": "open-catch-all"},
}
router: Final = Router(model_list=[reservation, *([catch_all] if catch_all_routes else [])])
one_hour: Final = _bucket(_metrics(prompt=_ONE_PTU_HOUR_OF_INPUT, completion=0))
attached: Final = with_ptu_consumption(_response(_day("2026-09-23", {"ptu-chat": one_hour})), router, "team-r")
assert attached.results[0].breakdown.model_groups["ptu-chat"].metrics.ptu_hours == pytest.approx(
expected_ptu_hours
)
def _mixed_ptu_router() -> Router:
"""One group split between team-a on a gpt-4.1 PTU deployment and team-b on a gpt-5.5 one,
beside an open pay-as-you-go deployment of gpt-4.1 in its own group."""
return Router(
model_list=[
{
"model_name": "ptu",
"litellm_params": {"model": "azure/gpt-4.1", "api_key": "sk-ptu", "api_base": "https://ptu.example"},
"model_info": {
"id": "ptu-41",
"base_model": "azure/gpt-4.1",
"ptu_count": 50,
"cost_per_ptu_per_hour": 1.0,
"ptu_effective_from": "2026-01-01T00:00:00Z",
"ptu_shares": {"team-a": 50},
},
},
{
"model_name": "ptu",
"litellm_params": {"model": "azure/gpt-5.5", "api_key": "sk-ptu", "api_base": "https://ptu.example"},
"model_info": {
"id": "ptu-55",
"base_model": "azure/gpt-5.5",
"ptu_count": 50,
"cost_per_ptu_per_hour": 1.0,
"ptu_effective_from": "2026-01-01T00:00:00Z",
"ptu_shares": {"team-b": 50},
},
},
{
"model_name": "gpt-4.1",
"litellm_params": {"model": "azure/gpt-4.1", "api_key": "sk-payg", "api_base": "https://payg.example"},
"model_info": {"id": "payg"},
},
]
)
def _reserved_ptu_router() -> Router:
"""A gpt-4.1 PTU deployment reserved for team-x alone the single-team way, beside an open
pay-as-you-go deployment of gpt-4.1 that serves every other team."""
return Router(
model_list=[
{
"model_name": "reserved",
"litellm_params": {"model": "azure/gpt-4.1", "api_key": "sk-ptu", "api_base": "https://ptu.example"},
"model_info": {
"id": "reserved-41",
"team_id": "team-x",
"ptu_count": 50,
"cost_per_ptu_per_hour": 1.0,
"ptu_effective_from": "2026-01-01T00:00:00Z",
},
},
{
"model_name": "gpt-4.1",
"litellm_params": {"model": "azure/gpt-4.1", "api_key": "sk-payg", "api_base": "https://payg.example"},
"model_info": {"id": "payg"},
},
]
)
@pytest.mark.parametrize(("team_id", "expected_on_the_provider_model"), [("team-x", 1.0), ("team-y", 0.0), (None, 1.0)])
def test_another_teams_single_team_reservation_sizes_nothing_on_a_teams_page(
monkeypatch, team_id: str | None, expected_on_the_provider_model: float
):
"""A provider-model row reaches the reserved deployment only for the team it is reserved for: any
other team was served by the open deployment, so its page counts no PTU-hours on that row, while
a page spanning teams keeps the reserved deployment's sizing."""
monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True")
one_hour: Final = {"azure/gpt-4.1": _bucket(_metrics(prompt=_ONE_PTU_HOUR_OF_INPUT, completion=0))}
attached: Final = with_ptu_consumption(_response(_day("2026-09-23", one_hour)), _reserved_ptu_router(), team_id)
row: Final = attached.results[0].breakdown.model_groups["azure/gpt-4.1"]
assert row.metrics.ptu_hours == pytest.approx(expected_on_the_provider_model)
assert attached.metadata.total_ptu_hours == pytest.approx(expected_on_the_provider_model)
_ONE_GPT_41_PTU_HOUR_ON_GPT_55: Final = _ONE_PTU_HOUR_OF_INPUT / AZURE_PTU_CAPACITY["gpt-5.5"].normalized_tokens_per_ptu_hour
@pytest.mark.parametrize(
("team_id", "expected_on_the_group", "expected_on_the_provider_model"),
[
("team-a", 1.0, 1.0),
("team-b", _ONE_GPT_41_PTU_HOUR_ON_GPT_55, 0.0),
("team-c", 0.0, 0.0),
(None, 1.0, 1.0),
],
)
def test_a_teams_rows_are_sized_by_the_deployment_it_is_served_from(
monkeypatch, team_id: str | None, expected_on_the_group: float, expected_on_the_provider_model: float
):
"""A team's page converts its tokens through the deployment the ceiling served it from: team-b's
share is on the gpt-5.5 deployment so its group row counts at that rate, its provider-model row
reached only the open deployment so it counts nothing, team-c holds no share so it counts nothing,
and a page spanning teams keeps the group's first sized deployment."""
monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True")
one_hour_each: Final = {
name: _bucket(_metrics(prompt=_ONE_PTU_HOUR_OF_INPUT, completion=0)) for name in ("ptu", "azure/gpt-4.1")
}
attached: Final = with_ptu_consumption(_response(_day("2026-09-23", one_hour_each)), _mixed_ptu_router(), team_id)
groups: Final = attached.results[0].breakdown.model_groups
assert groups["ptu"].metrics.ptu_hours == pytest.approx(expected_on_the_group)
assert groups["azure/gpt-4.1"].metrics.ptu_hours == pytest.approx(expected_on_the_provider_model)
assert attached.metadata.total_ptu_hours == pytest.approx(expected_on_the_group + expected_on_the_provider_model)

View file

@ -2,37 +2,43 @@
import datetime
import json
from collections.abc import Mapping, Sequence
from contextlib import ExitStack
from typing import Final, cast
from unittest.mock import AsyncMock, MagicMock, patch
from unittest.mock import patch as patch_ctx
import pytest
from fastapi import HTTPException
from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token
from litellm.litellm_core_utils.ptu_pricing import ptu_terms
from litellm.llms.gemini.cost_calculator import cost_per_web_search_request
from litellm.proxy._types import (
LiteLLM_ProxyModelTable,
LiteLLM_TeamTable,
LitellmUserRoles,
ProxyException,
ReconcileOutcome,
UserAPIKeyAuth,
)
from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token
from litellm.proxy.auth.auth_checks import is_model_cost_zero
from litellm.llms.gemini.cost_calculator import cost_per_web_search_request
from litellm.proxy.management_endpoints.model_management_endpoints import (
_PTU_ZEROED_PRICING_FIELDS,
_SEARCH_CONTEXT_SIZES,
_is_nonzero_price,
_merged_ptu_model_info,
_update_team_model_in_db,
_ptu_priced_deployment,
_ptu_zeroed_pricing,
_raise_if_ptu_cost_attribution_disabled,
_raise_if_ptu_share_teams_missing,
_update_team_model_in_db,
_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.utils import PromptTokensDetailsWrapper
from litellm.repositories.chunked_in import IN_LIST_CHUNK_SIZE
from litellm.router import Router
from litellm.types.router import (
SPECIAL_MODEL_INFO_PARAMS,
@ -42,7 +48,7 @@ from litellm.types.router import (
updateDeployment,
updateLiteLLMParams,
)
from litellm.types.utils import Usage
from litellm.types.utils import PromptTokensDetailsWrapper, Usage
async def _passthrough_row(update_data):
@ -115,6 +121,21 @@ def test_model_info_allows_partial_delta_for_patch():
assert info.cost_per_ptu_per_hour is None
@pytest.mark.parametrize("share", [True, "2", 2.5, float("inf")])
def test_model_info_rejects_a_share_that_is_not_a_whole_number(share):
with pytest.raises(ValueError, match="ptu_shares"):
ModelInfo(id="x", ptu_shares={"team-a": share})
def test_model_info_keeps_whole_number_shares_and_refuses_a_fractional_count():
assert ModelInfo(id="x", ptu_shares={"team-a": 2}).ptu_shares == {"team-a": 2}
written_as_float = ModelInfo(id="x", ptu_shares={"team-a": 2.0}).ptu_shares
assert written_as_float == {"team-a": 2}
assert written_as_float is not None and type(written_as_float["team-a"]) is int
with pytest.raises(ValueError, match="ptu_count"):
ModelInfo(id="x", ptu_count=100.5)
def test_validate_helper_no_ptu_is_noop():
_validate_ptu_model_info({"team_id": "t"})
@ -639,7 +660,7 @@ class TestAddNewModelPtuGate:
monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False)
@staticmethod
def _patched_proxy(model_id: str):
def _patched_proxy(model_id: str, prisma_client: MagicMock | None = None):
"""Patch everything /model/new touches except the PTU gate, and hand back the DB writers."""
db_row = LiteLLM_ProxyModelTable(
model_id=model_id,
@ -666,7 +687,7 @@ class TestAddNewModelPtuGate:
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}.prisma_client", prisma_client if prisma_client is not None else 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()),
@ -1284,3 +1305,144 @@ class TestPtuDeploymentsAreNotBilledPerToken:
assert "input_cost_per_token" in str(exc.value)
add_team_model_to_db.assert_not_called()
_SHARED_START = "2026-08-01T00:00:00Z"
def test_validate_helper_accepts_shares_in_place_of_a_team_id():
shared = {
"ptu_count": 5,
"cost_per_ptu_per_hour": 2.0,
"ptu_effective_from": _SHARED_START,
"ptu_shares": {"team-a": 3, "team-b": 2},
}
_validate_ptu_model_info(shared)
assert ptu_terms(shared) is not None
def test_validate_helper_refuses_shares_that_do_not_add_up_to_the_count():
with pytest.raises(HTTPException) as exc:
_validate_ptu_model_info(
{
"ptu_count": 5,
"cost_per_ptu_per_hour": 2.0,
"ptu_effective_from": _SHARED_START,
"ptu_shares": {"team-a": 3, "team-b": 1},
}
)
assert exc.value.status_code == 400
assert "4 of 5 allocated" in exc.value.detail
def test_validate_helper_refuses_shares_declared_without_a_count_and_rate():
with pytest.raises(HTTPException) as exc:
_validate_ptu_model_info({"ptu_shares": {"team-a": 3, "team-b": 2}, "ptu_effective_from": _SHARED_START})
assert exc.value.status_code == 400
assert exc.value.detail == "ptu_count and cost_per_ptu_per_hour are required when ptu_shares is set"
class _TeamLookup:
def __init__(self, existing: frozenset[str]) -> None:
self.existing: Final = existing
self.looked_up: tuple[str, ...] = ()
self.batch_sizes: tuple[int, ...] = ()
async def find_unique(self, *, where: Mapping[str, object]) -> LiteLLM_TeamTable | None:
raise AssertionError(f"one lookup per team is what the review asked to avoid: {where}")
async def find_many(self, *, where: Mapping[str, object]) -> Sequence[LiteLLM_TeamTable]:
team_filter: Final = where["team_id"]
assert isinstance(team_filter, Mapping)
requested: Final = tuple(str(team_id) for team_id in cast(Sequence[object], team_filter["in"]))
self.looked_up = (*self.looked_up, *requested)
self.batch_sizes = (*self.batch_sizes, len(requested))
return tuple(LiteLLM_TeamTable(team_id=team_id) for team_id in requested if team_id in self.existing)
def _shared_model_info(shares: Mapping[str, int]) -> Mapping[str, object]:
return {
"ptu_count": sum(shares.values()),
"cost_per_ptu_per_hour": 2.0,
"ptu_effective_from": _SHARED_START,
"ptu_shares": dict(shares),
}
@pytest.mark.asyncio
async def test_share_team_check_refuses_a_team_that_does_not_exist():
with pytest.raises(HTTPException) as exc:
await _raise_if_ptu_share_teams_missing(
_shared_model_info({"team-a": 3, "ghost-team": 2}), lambda: _TeamLookup(frozenset({"team-a"}))
)
assert exc.value.status_code == 400
assert exc.value.detail == {"error": "Team id=ghost-team does not exist in db"}
@pytest.mark.asyncio
async def test_share_team_check_accepts_shares_naming_existing_teams():
lookup: Final = _TeamLookup(frozenset({"team-a", "team-b"}))
await _raise_if_ptu_share_teams_missing(_shared_model_info({"team-a": 3, "team-b": 2}), lambda: lookup)
assert lookup.looked_up == ("team-a", "team-b")
@pytest.mark.asyncio
async def test_share_team_check_splits_a_share_list_longer_than_one_in_list_chunk():
team_ids: Final = tuple(f"team-{index}" for index in range(IN_LIST_CHUNK_SIZE + 1))
lookup: Final = _TeamLookup(frozenset(team_ids))
await _raise_if_ptu_share_teams_missing(_shared_model_info(dict.fromkeys(team_ids, 1)), lambda: lookup)
assert lookup.looked_up == team_ids
assert max(lookup.batch_sizes) <= IN_LIST_CHUNK_SIZE
@pytest.mark.asyncio
async def test_share_team_check_leaves_a_team_id_holder_to_the_team_model_check():
lookup: Final = _TeamLookup(frozenset())
await _raise_if_ptu_share_teams_missing(
{"team_id": "team-a", "ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "ptu_effective_from": _SHARED_START},
lambda: lookup,
)
assert lookup.looked_up == ()
@pytest.mark.asyncio
async def test_model_new_refuses_shares_naming_a_team_that_does_not_exist(monkeypatch):
monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true")
prisma_client = MagicMock()
prisma_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[])
(add_model_to_db, add_team_model_to_db), patches = TestAddNewModelPtuGate._patched_proxy(
"ptu-shared-model", prisma_client=prisma_client
)
admin = UserAPIKeyAuth(user_id="test-admin", user_role=LitellmUserRoles.PROXY_ADMIN)
base = TestAddNewModelPtuGate._ptu_deployment("ptu-shared-model")
shared = base.model_copy(
update={"model_info": base.model_info.model_copy(update={"team_id": None, "ptu_shares": {"ghost-team": 15}})}
)
with ExitStack() as stack:
for active_patch in patches:
stack.enter_context(active_patch)
with pytest.raises(ProxyException, match="Team id=ghost-team does not exist in db") as exc:
await add_new_model(model_params=shared, user_api_key_dict=admin)
assert exc.value.code == "400"
prisma_client.db.litellm_teamtable.find_many.assert_awaited_once_with(where={"team_id": {"in": ["ghost-team"]}})
add_model_to_db.assert_not_called()
add_team_model_to_db.assert_not_called()
def test_validate_helper_refuses_a_team_id_beside_shares():
with pytest.raises(HTTPException) as exc:
_validate_ptu_model_info(
{
"team_id": "team-a",
"ptu_count": 5,
"cost_per_ptu_per_hour": 2.0,
"ptu_effective_from": _SHARED_START,
"ptu_shares": {"team-a": 5},
}
)
assert exc.value.status_code == 400
assert "cannot both be set" in exc.value.detail

View file

@ -17,6 +17,7 @@ from pydantic import ValidationError
from litellm._uuid import uuid
from litellm.integrations.custom_logger import CustomLogger
from litellm.llms.azure.ptu_capacity import AZURE_PTU_CAPACITY
from litellm.proxy._types import (
LiteLLM_BudgetTable,
LiteLLM_BudgetTableFull,
@ -9086,6 +9087,52 @@ async def test_delete_team_persists_deleted_teams(
assert records[0]["litellm_changed_by"] == "admin-user"
@pytest.mark.asyncio
async def test_delete_team_is_refused_while_the_team_holds_a_ptu_share(
monkeypatch,
disable_audit_logging_for_mocked_team,
):
from litellm import Router
from litellm.proxy._types import DeleteTeamRequest
shared_router = Router(
model_list=[
{
"model_name": "gpt-4.1-ptu",
"litellm_params": {"model": "azure/gpt-4.1", "api_key": "sk-azure", "api_base": "https://x"},
"model_info": {
"id": "shared-ptu",
"ptu_count": 50,
"cost_per_ptu_per_hour": 1.0,
"ptu_effective_from": "2026-01-01T00:00:00Z",
"ptu_shares": {"team-1": 30, "team-2": 20},
},
}
]
)
mock_prisma_client = AsyncMock()
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(
side_effect=lambda where: LiteLLM_TeamTable(team_id=where["team_id"], members_with_roles=[], metadata={})
)
mock_prisma_client.db.litellm_deletedteamtable.create_many = AsyncMock()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", shared_router)
admin = UserAPIKeyAuth(user_id="admin-user", api_key="sk-admin", user_role=LitellmUserRoles.PROXY_ADMIN.value)
with pytest.raises(HTTPException) as refused:
await delete_team(
data=DeleteTeamRequest(team_ids=["team-3", "team-2"]),
http_request=MagicMock(),
user_api_key_dict=admin,
litellm_changed_by="admin-user",
)
assert refused.value.status_code == 400
assert "team-2" in str(refused.value.detail) and "shared-ptu" in str(refused.value.detail)
mock_prisma_client.db.litellm_deletedteamtable.create_many.assert_not_called()
mock_prisma_client.delete_data.assert_not_called()
@pytest.mark.asyncio
async def test_delete_team_sweeps_references_outside_members_with_roles(
monkeypatch,
@ -17968,6 +18015,177 @@ def test_list_team_v2_answers_503_no_db_connection_when_the_callers_user_read_hi
assert response.json() == _DB_OUTAGE_503_BODY
# --- PTU-equivalent consumption on /team/daily/activity ------------------------------------
_ONE_PTU_HOUR_OF_INPUT: Final = AZURE_PTU_CAPACITY["gpt-4.1"].normalized_tokens_per_ptu_hour
def _db_with_one_ptu_hour_of_team_spend(team_id: str) -> MagicMock:
"""The team's daily spend table holds one day of three requests on the PTU group, carrying
one PTU-hour of gpt-4.1 uncached input; the read path behind the endpoint runs as written."""
row: Final = SimpleNamespace(
date="2026-09-23",
team_id=team_id,
api_key="ptu-key",
model="azure/gpt-4.1",
model_group="gpt-4.1-ptu",
custom_llm_provider="azure",
mcp_namespaced_tool_name=None,
endpoint="/chat/completions",
spend=0.0,
prompt_tokens=_ONE_PTU_HOUR_OF_INPUT,
completion_tokens=0,
cache_read_input_tokens=0,
cache_creation_input_tokens=0,
compression_saved_tokens=0,
compression_savings_spend=0.0,
prompt_caching_savings_spend=0.0,
gateway_injected_caching_savings_spend=0.0,
autorouter_savings_spend=0.0,
total_response_time_ms=0,
timed_requests=0,
api_requests=3,
successful_requests=3,
failed_requests=0,
)
db_client: Final = MagicMock()
db_client.db.litellm_dailyteamspend.count = AsyncMock(return_value=1)
db_client.db.litellm_dailyteamspend.find_many = AsyncMock(return_value=[row])
db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
db_client.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[])
db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[])
return db_client
def _shared_ptu_router() -> Router:
return Router(
model_list=[
{
"model_name": "gpt-4.1-ptu",
"litellm_params": {"model": "azure/gpt-4.1", "api_key": "sk-ptu", "api_base": "https://ptu.example"},
"model_info": {
"id": "shared-ptu",
"base_model": "azure/gpt-4.1",
"ptu_count": 50,
"cost_per_ptu_per_hour": 1.0,
"ptu_effective_from": "2026-01-01T00:00:00Z",
"ptu_shares": {"team-a": 30, "team-b": 20},
},
}
]
)
@pytest.mark.asyncio
@pytest.mark.parametrize("attribution_enabled", [True, False])
async def test_team_daily_activity_reports_ptu_hours_only_while_attribution_is_on(
mock_admin_auth, monkeypatch, attribution_enabled
):
"""An hour of one PTU's input rate on gpt-4.1, in uncached input tokens, is one PTU-hour;
the figure appears beside tokens only once the PTU flag is on, and the token totals are
untouched either way."""
from litellm.proxy.management_endpoints.team_endpoints import get_team_daily_activity
monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True" if attribution_enabled else "False")
with (
patch("litellm.proxy.proxy_server.prisma_client", _db_with_one_ptu_hour_of_team_spend("team-a")),
patch("litellm.proxy.proxy_server.llm_router", _shared_ptu_router()),
):
result = await get_team_daily_activity(
team_ids="team-a",
start_date="2026-09-23",
end_date="2026-09-24",
model=None,
api_key=None,
page=1,
page_size=10,
exclude_team_ids=None,
user_api_key_dict=mock_admin_auth,
)
expected_ptu_hours = 1.0 if attribution_enabled else 0.0
assert result.metadata.total_ptu_hours == expected_ptu_hours
assert result.results[0].metrics.ptu_hours == expected_ptu_hours
assert result.results[0].breakdown.model_groups["gpt-4.1-ptu"].metrics.ptu_hours == expected_ptu_hours
assert result.metadata.total_tokens == _ONE_PTU_HOUR_OF_INPUT
assert result.results[0].metrics.total_tokens == _ONE_PTU_HOUR_OF_INPUT
@pytest.mark.asyncio
async def test_team_daily_activity_sizes_a_team_scoped_deployment_by_its_public_name(mock_admin_auth, monkeypatch):
"""A deployment registered for one team is renamed to a unique internal name and keeps the
name callers use in ``team_public_model_name``, which is the name the activity rows carry,
so its sizing row is still found."""
from litellm.proxy.management_endpoints.team_endpoints import get_team_daily_activity
monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True")
team_scoped_router: Final = Router(
model_list=[
{
"model_name": "gpt-4.1-ptu-3f9c1b",
"litellm_params": {"model": "azure/gpt-4.1", "api_key": "sk-ptu", "api_base": "https://ptu.example"},
"model_info": {
"id": "team-a-ptu",
"base_model": "azure/gpt-4.1",
"team_id": "team-a",
"team_public_model_name": "gpt-4.1-ptu",
"ptu_count": 50,
"cost_per_ptu_per_hour": 1.0,
"ptu_effective_from": "2026-01-01T00:00:00Z",
},
}
]
)
with (
patch("litellm.proxy.proxy_server.prisma_client", _db_with_one_ptu_hour_of_team_spend("team-a")),
patch("litellm.proxy.proxy_server.llm_router", team_scoped_router),
):
result = await get_team_daily_activity(
team_ids="team-a",
start_date="2026-09-23",
end_date="2026-09-24",
model=None,
api_key=None,
page=1,
page_size=10,
exclude_team_ids=None,
user_api_key_dict=mock_admin_auth,
)
assert result.metadata.total_ptu_hours == 1.0
assert result.results[0].breakdown.model_groups["gpt-4.1-ptu"].metrics.ptu_hours == 1.0
@pytest.mark.asyncio
async def test_team_daily_activity_reports_no_ptu_hours_to_a_team_holding_no_share(mock_admin_auth, monkeypatch):
"""The deployment is split between team-a and team-b, so team-c's page alone converts none of
its tokens on the group even though the group is sized."""
from litellm.proxy.management_endpoints.team_endpoints import get_team_daily_activity
monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True")
with (
patch("litellm.proxy.proxy_server.prisma_client", _db_with_one_ptu_hour_of_team_spend("team-c")),
patch("litellm.proxy.proxy_server.llm_router", _shared_ptu_router()),
):
result = await get_team_daily_activity(
team_ids="team-c",
start_date="2026-09-23",
end_date="2026-09-24",
model=None,
api_key=None,
page=1,
page_size=10,
exclude_team_ids=None,
user_api_key_dict=mock_admin_auth,
)
assert result.metadata.total_ptu_hours == 0.0
assert result.results[0].breakdown.model_groups["gpt-4.1-ptu"].metrics.ptu_hours == 0.0
assert result.results[0].metrics.total_tokens == _ONE_PTU_HOUR_OF_INPUT
@pytest.mark.parametrize(
("start_date", "end_date"),
(

View file

@ -1,11 +1,55 @@
"""Tests for the opt-in flag that gates PTU flat-cost attribution."""
from collections.abc import Iterator
import httpx
import pytest
import litellm
from litellm.integrations.custom_secret_manager import CustomSecretManager
from litellm.litellm_core_utils import ptu_pricing
from litellm.secret_managers.main import get_secret_bool
from litellm.proxy.spend_tracking.ptu_feature_flag import (
PTU_COST_ATTRIBUTION_ENV_VAR,
is_ptu_cost_attribution_enabled,
)
from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem
class _SecretManagerHolding(CustomSecretManager):
def __init__(self, secrets: dict[str, str]) -> None:
super().__init__(secret_manager_name="flag_holder")
self.secrets = secrets
self.reads = 0
async def async_read_secret(
self,
secret_name: str,
optional_params: dict | None = None,
timeout: float | httpx.Timeout | None = None,
) -> str | None:
return self.sync_read_secret(secret_name, optional_params, timeout)
def sync_read_secret(
self,
secret_name: str,
optional_params: dict | None = None,
timeout: float | httpx.Timeout | None = None,
) -> str | None:
self.reads += 1
return self.secrets.get(secret_name)
@pytest.fixture
def secret_manager(monkeypatch: pytest.MonkeyPatch) -> Iterator[_SecretManagerHolding]:
manager = _SecretManagerHolding({PTU_COST_ATTRIBUTION_ENV_VAR: "true"})
monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False)
monkeypatch.setattr(litellm, "secret_manager_client", manager)
monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.CUSTOM)
monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode="read_only"))
ptu_pricing._enabled_in_secret_manager.cache_clear()
yield manager
ptu_pricing._enabled_in_secret_manager.cache_clear()
def test_disabled_when_env_var_is_unset(monkeypatch):
@ -31,3 +75,32 @@ def test_reads_the_env_var_on_every_call(monkeypatch):
monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true")
assert is_ptu_cost_attribution_enabled() is True
def test_the_secret_manager_outranks_the_env_like_every_other_setting(
monkeypatch: pytest.MonkeyPatch, secret_manager: _SecretManagerHolding
) -> None:
monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "false")
assert is_ptu_cost_attribution_enabled() is get_secret_bool(PTU_COST_ATTRIBUTION_ENV_VAR) is True
def test_a_secret_manager_set_to_false_turns_attribution_off_despite_the_env(
monkeypatch: pytest.MonkeyPatch, secret_manager: _SecretManagerHolding
) -> None:
secret_manager.secrets[PTU_COST_ATTRIBUTION_ENV_VAR] = "false"
monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true")
assert is_ptu_cost_attribution_enabled() is False
def test_a_flag_kept_in_a_secret_manager_still_turns_attribution_on(secret_manager: _SecretManagerHolding) -> None:
assert is_ptu_cost_attribution_enabled() is True
def test_the_secret_manager_is_asked_once_for_many_requests(secret_manager: _SecretManagerHolding) -> None:
assert all(is_ptu_cost_attribution_enabled() for _ in range(50))
assert secret_manager.reads == 1
def test_a_secret_manager_without_the_flag_leaves_it_off(secret_manager: _SecretManagerHolding) -> None:
secret_manager.secrets.clear()
assert is_ptu_cost_attribution_enabled() is False

View file

@ -16,7 +16,7 @@ from litellm.proxy.spend_tracking.ptu_flat_cost_rollup import (
PTUModel,
_active_hours_on_day,
_compute_daily_flat_cost,
_parse_ptu_model,
_parse_ptu_models,
run_ptu_flat_cost_backfill,
run_ptu_flat_cost_rollup,
run_scheduled_ptu_rollup,
@ -43,6 +43,14 @@ def _ptu_enabled(monkeypatch):
_VALID_PTU = {"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "team_id": "t"}
def _parse_ptu_model(row):
"""The single holding a team_id deployment parses to, else None; shared deployments
parse to one holding per team and have their own tests below."""
parsed = _parse_ptu_models(row)
assert len(parsed) <= 1
return parsed[0] if parsed else None
def _model_row(model_id="m1", model_name="gpt-4o-mini-ptu", model_info=None, with_start=True):
row = MagicMock()
row.model_id = model_id
@ -163,6 +171,49 @@ def _sentinel_row(row_id, team_id, model):
return row
class _FakeBatch:
"""prisma's ``batch_()``: writes queue inside the block and commit as one transaction on a clean exit.
A queued delete reaches the table as ``batch_delete_many``, kept apart from the table's own
``delete_many`` the way prisma keeps batch actions apart from direct ones, so the prune's
assertions still see only the prune."""
def __init__(self, table):
self._table = table
self._queued = []
self.litellm_dailyteamspend = self
def upsert(self, where, data):
self._queued.append((where, data))
def delete_many(self, where):
self._queued.append((where, None))
async def __aenter__(self):
return self
async def __aexit__(self, exc_type, exc, tb):
if exc_type is not None:
return False
committed = (
(dict(self._table.rows), len(self._table.upsert_keys))
if isinstance(self._table, _FakeSentinelTable)
else None
)
try:
for where, data in self._queued:
if data is None:
self._table.batch_delete_many(where)
else:
await self._table.upsert(where=where, data=data)
except Exception:
if committed is not None:
self._table.rows = committed[0]
del self._table.upsert_keys[committed[1] :]
raise
return False
def _prisma_with_models(rows, existing_sentinel_rows=()):
prisma = MagicMock()
model_table = MagicMock()
@ -171,7 +222,9 @@ def _prisma_with_models(rows, existing_sentinel_rows=()):
daily.find_many = AsyncMock(return_value=list(existing_sentinel_rows))
daily.upsert = AsyncMock()
daily.delete_many = AsyncMock()
prisma.db = types.SimpleNamespace(litellm_proxymodeltable=model_table, litellm_dailyteamspend=daily)
prisma.db = types.SimpleNamespace(
litellm_proxymodeltable=model_table, litellm_dailyteamspend=daily, batch_=lambda: _FakeBatch(daily)
)
return prisma, daily
@ -687,9 +740,10 @@ async def test_rollup_skips_the_prune_when_a_replacement_write_failed():
class _FakeSentinelTable:
"""In-memory LiteLLM_DailyTeamSpend that honours the sentinel key and prune predicate."""
def __init__(self, upsert_gate=None):
def __init__(self, upsert_gate=None, failing_team=None):
self.rows = {}
self._upsert_gate = upsert_gate
self.failing_team = failing_team
self.upsert_keys = []
self.delete_many_calls = []
self.find_many_calls = []
@ -698,6 +752,8 @@ class _FakeSentinelTable:
if self._upsert_gate is not None:
await self._upsert_gate.wait()
key = where["team_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint"]
if key["team_id"] == self.failing_team:
raise RuntimeError(f"write rejected for {self.failing_team}")
row_key = (key["team_id"], key["date"], key["api_key"], key["model"])
self.upsert_keys.append(row_key)
self.rows[row_key] = {
@ -724,6 +780,16 @@ class _FakeSentinelTable:
del self.rows[k]
return len(doomed)
def batch_delete_many(self, where):
"""The batch delete of a deployment day's rows. Every clause is honoured and an unexpected
one fails, so a dropped scope cannot pass silently."""
assert set(where) == {"date", "api_key", "model"}
doomed = [
k for k in self.rows if k[1] == where["date"] and k[2] == where["api_key"] and k[3] == where["model"]
]
for k in doomed:
del self.rows[k]
async def find_many(self, where=None):
"""Read back sentinel rows the way prisma would, honouring api_key and a date range."""
self.find_many_calls.append(where)
@ -758,7 +824,11 @@ def _prisma_for(model_rows, daily_table):
prisma = MagicMock()
model_table = MagicMock()
model_table.find_many = AsyncMock(return_value=model_rows)
prisma.db = types.SimpleNamespace(litellm_proxymodeltable=model_table, litellm_dailyteamspend=daily_table)
prisma.db = types.SimpleNamespace(
litellm_proxymodeltable=model_table,
litellm_dailyteamspend=daily_table,
batch_=lambda: _FakeBatch(daily_table),
)
return prisma
@ -1594,6 +1664,70 @@ async def test_backfill_does_not_double_price_a_day_after_a_rename():
assert table.upsert_keys == [("t", _day(0).isoformat(), PTU_SENTINEL_API_KEY, "dep-1")]
@pytest.mark.asyncio
async def test_backfill_does_not_bill_a_priced_day_again_after_its_shares_moved_to_other_teams():
"""Days already billed to the team that held the deployment stay its only charge once the
shares move, and only the unpriced day is billed to the teams holding the shares now."""
table = _FakeSentinelTable()
for offset in (-2, -1):
table.seed("team-b", _day(offset), "shared", 1200.0)
prisma = _prisma_for(
[
_model_row(
model_id="shared",
model_info={
"ptu_count": 50,
"cost_per_ptu_per_hour": 1.0,
"ptu_shares": {"team-a": 30, "team-c": 20},
"ptu_effective_from": _midnight(_day(-2)).isoformat(),
},
)
],
table,
)
result = await run_ptu_flat_cost_backfill(prisma, today=TODAY)
assert sorted(table.upsert_keys) == [
("team-a", _day(0).isoformat(), PTU_SENTINEL_API_KEY, "shared"),
("team-c", _day(0).isoformat(), PTU_SENTINEL_API_KEY, "shared"),
]
assert result.rows_written == 2
@pytest.mark.asyncio
async def test_a_split_deployment_day_whose_write_fails_for_one_team_is_billed_to_every_team_later():
"""A day holding only some teams' rows would read as priced to the backfill, so the teams
whose rows were lost would never be billed for it."""
table = _FakeSentinelTable(failing_team="team-c")
prisma = _prisma_for(
[
_model_row(
model_id="shared",
model_info={
"ptu_count": 50,
"cost_per_ptu_per_hour": 1.0,
"ptu_shares": {"team-a": 30, "team-c": 20},
"ptu_effective_from": _midnight(_day(0)).isoformat(),
},
)
],
table,
)
rollup = await run_ptu_flat_cost_rollup(prisma, target_date=_day(0))
assert table.rows == {}
assert (rollup.rows_written, rollup.rows_failed) == (0, 2)
table.failing_team = None
backfill = await run_ptu_flat_cost_backfill(prisma, today=TODAY)
assert sorted(team for team, *_ in table.rows) == ["team-a", "team-c"]
assert sum(row["ptu_flat_cost"] for row in table.rows.values()) == pytest.approx(50 * 24)
assert backfill.rows_written == 2
@pytest.mark.asyncio
async def test_backfill_still_prices_a_genuinely_missing_day_for_a_renamed_deployment():
"""Rename safety must not swallow real gaps: a day with no row for the deployment at
@ -1746,6 +1880,38 @@ async def test_an_unguarded_run_writes_but_does_not_prune():
assert ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-live") in table.rows
@pytest.mark.asyncio
async def test_an_unguarded_rerun_after_a_share_move_bills_the_day_to_the_new_holders_alone():
"""Without the lock the prune is skipped, so the old holder's row has to go with the day's own
write, or rerunning a priced day after its shares moved bills it to the old and new holders."""
table = _FakeSentinelTable()
table.seed("team-a", DAY, "shared", 720.0)
table.seed("team-b", DAY, "shared", 480.0)
table.seed("team-b", DAY, "other", 240.0)
prisma = _prisma_for(
[
_model_row(
model_id="shared",
model_info={
"ptu_count": 50,
"cost_per_ptu_per_hour": 1.0,
"ptu_shares": {"team-a": 30, "team-c": 20},
"ptu_effective_from": _midnight(DAY).isoformat(),
},
)
],
table,
)
await run_scheduled_ptu_rollup(prisma, target_date=DAY)
assert {key[0]: row["ptu_flat_cost"] for key, row in table.rows.items() if key[3] == "shared"} == {
"team-a": pytest.approx(720.0),
"team-c": pytest.approx(480.0),
}
assert ("team-b", DAY.isoformat(), PTU_SENTINEL_API_KEY, "other") in table.rows
@pytest.mark.asyncio
async def test_a_run_holding_the_lock_still_prunes():
"""Losing the sweep entirely would leave stale charges forever, so the guarded path,
@ -2105,3 +2271,54 @@ async def test_the_catch_up_pass_reaches_a_config_declared_deployment():
assert len(charged) == 3, charged
assert charged[-1] == yesterday
assert all(row["ptu_flat_cost"] == pytest.approx(48.0) for row in table.rows.values())
# --- a deployment split into per-team PTU shares -----------------------------------------
_SHARED_PTU = {"ptu_count": 50, "cost_per_ptu_per_hour": 1.0, "ptu_shares": {"team-b": 20, "team-a": 30}}
def test_a_shared_deployment_parses_to_one_holding_per_team_carrying_its_share():
parsed = _parse_ptu_models(_model_row(model_info=dict(_SHARED_PTU)))
assert [(m.team_id, m.ptu_count) for m in parsed] == [("team-a", 30), ("team-b", 20)]
assert {m.model_id for m in parsed} == {"m1"}
assert {m.model_name for m in parsed} == {"gpt-4o-mini-ptu"}
assert {m.cost_per_ptu_per_hour for m in parsed} == {1.0}
def test_a_shared_deployment_whose_shares_do_not_add_up_is_not_priced():
assert _parse_ptu_models(_model_row(model_info={**_SHARED_PTU, "ptu_shares": {"team-a": 30}})) == ()
@pytest.mark.asyncio
async def test_rollup_splits_a_shared_deployments_flat_cost_by_share():
"""50 PTUs at $1/hour for a day is $1,200; team-a's 30 PTUs are $720 of it and team-b's
20 are $480, each keyed on the same deployment id and public name."""
prisma, table = _prisma_with_models([_model_row(model_info=dict(_SHARED_PTU))])
result = await run_ptu_flat_cost_rollup(prisma, target_date=DAY)
assert result.rows_written == 2
assert result.models_processed == 1
created = {
call.kwargs["data"]["create"]["team_id"]: call.kwargs["data"]["create"] for call in table.upsert.await_args_list
}
assert created["team-a"]["ptu_flat_cost"] == pytest.approx(720.0)
assert created["team-b"]["ptu_flat_cost"] == pytest.approx(480.0)
assert sum(row["ptu_flat_cost"] for row in created.values()) == pytest.approx(50 * 1.0 * 24)
assert {row["model"] for row in created.values()} == {"m1"}
assert {row["model_group"] for row in created.values()} == {"gpt-4o-mini-ptu"}
assert {row["api_key"] for row in created.values()} == {PTU_SENTINEL_API_KEY}
@pytest.mark.asyncio
async def test_a_shared_deployment_with_a_closed_window_is_reported_lapsed_once():
closed = {**_SHARED_PTU, "ptu_effective_from": "2020-01-01T00:00:00Z", "ptu_effective_to": "2020-02-01T00:00:00Z"}
prisma, _ = _prisma_with_models([_model_row(model_id="dep-shared", model_info=closed)])
alert = AsyncMock()
result = await run_scheduled_ptu_rollup(prisma, target_date=DAY, alert=alert)
assert result.lapsed == ("gpt-4o-mini-ptu",)
alert.assert_awaited_once()

View file

@ -0,0 +1,301 @@
"""Tests for per-team PTU shares: who a shared deployment is served to and what a share is worth."""
from datetime import date
from typing import Final
import pytest
from litellm.llms.azure.ptu_capacity import AZURE_PTU_CAPACITY
from litellm.router_utils.ptu_shares import (
filter_ptu_shared_deployments,
model_group_deployments,
model_group_ptu_capacity,
ptu_capacity_warning,
ptu_share_teams,
PTUTeamCeiling,
share_holder_request_kwargs,
team_ptu_ceiling,
team_servable_deployments,
)
_GPT41: Final = AZURE_PTU_CAPACITY["gpt-4.1"]
_GPT4O: Final = AZURE_PTU_CAPACITY["gpt-4o"]
_GPT6SOL: Final = AZURE_PTU_CAPACITY["gpt-6-sol"]
_SHARES: Final = {"team-a": 30, "team-b": 20}
def _shared(
model: str = "azure/gpt-4.1", shares: object = _SHARES, deployment_id: str = "shared", ptu_count: int | None = None
) -> dict:
return {
"model_name": "gpt-4.1-ptu",
"litellm_params": {"model": model},
"model_info": {
"id": deployment_id,
"ptu_count": ptu_count or (sum(shares.values()) if isinstance(shares, dict) else 50),
"cost_per_ptu_per_hour": 1.0,
"ptu_effective_from": "2026-01-01T00:00:00Z",
"ptu_shares": shares,
},
}
def _single_team(model: str = "azure/gpt-4.1") -> dict:
return {
"model_name": "gpt-4.1-ptu",
"litellm_params": {"model": model},
"model_info": {
"id": "single",
"team_id": "team-a",
"ptu_count": 50,
"cost_per_ptu_per_hour": 1.0,
"ptu_effective_from": "2026-01-01T00:00:00Z",
},
}
_OPEN: Final = {"model_name": "gpt-4.1-ptu", "litellm_params": {"model": "azure/gpt-4.1"}, "model_info": {"id": "open"}}
def _no_wildcard_route(_: str) -> tuple[()]:
return ()
def _unaliased_ceiling(deployments: list[dict], team_id: str, requested_model: str) -> PTUTeamCeiling | None:
return team_ptu_ceiling(deployments, deployments, team_id, requested_model, _no_wildcard_route)
def test_a_team_holding_a_share_is_served_from_its_share_alone():
"""An open deployment beside the shared one would put the team's pay-as-you-go traffic
under its PTU ceiling and into its PTU-hours, so only the teams holding no share get it."""
result: Final = filter_ptu_shared_deployments([_shared(), _OPEN], "team-a")
assert [d["model_info"]["id"] for d in result.deployments] == ["shared"]
def test_a_team_without_a_share_only_sees_the_unshared_deployments():
result: Final = filter_ptu_shared_deployments([_shared(), _OPEN], "team-c")
assert [d["model_info"]["id"] for d in result.deployments] == ["open"]
assert result.withheld is True
def test_a_caller_with_no_team_holds_no_share():
for team_id in (None, ""):
result = filter_ptu_shared_deployments([_shared()], team_id)
assert result.deployments == ()
assert result.withheld is True
def test_a_single_team_deployment_and_a_malformed_share_map_are_not_filtered_here():
"""A ``team_id`` deployment is scoped by the router's team filter, and an unusable
``ptu_shares`` is refused at registration, so neither is withheld by the share filter."""
result: Final = filter_ptu_shared_deployments([_single_team(), _shared(shares={"team-a": 0})], "team-z")
assert [d["model_info"]["id"] for d in result.deployments] == ["single", "shared"]
assert result.withheld is False
def test_only_a_share_holders_request_keeps_its_team_for_the_retry_lookup():
"""A team owning a deployment through ``team_id`` keeps the router's unscoped retry lookup."""
share_teams: Final = ptu_share_teams([_shared(shares={"team-b": 50}), _single_team()])
holder: Final = {"metadata": {"user_api_key_team_id": "team-b"}}
messages_holder: Final = {"litellm_metadata": {"user_api_key_team_id": "team-b"}}
assert share_holder_request_kwargs(holder, share_teams) is holder
assert share_holder_request_kwargs(messages_holder, share_teams) is messages_holder
assert share_holder_request_kwargs({"metadata": {"user_api_key_team_id": "team-a"}}, share_teams) is None
assert share_holder_request_kwargs({"metadata": {}}, share_teams) is None
_TERMLESS_SHARED: Final = {
"model_name": "gpt-4.1-ptu",
"litellm_params": {"model": "azure/gpt-4.1"},
"model_info": {"id": "shared", "ptu_shares": _SHARES},
}
def test_a_declared_split_is_the_access_rule_even_without_pricing_terms():
"""A map that reached the router without the count, rate and start registration now demands
still names who the deployment is for, so it is served to those teams alone and sized from
their shares; only the flat cost, which needs the terms, goes unattributed."""
assert filter_ptu_shared_deployments([_TERMLESS_SHARED], "team-z").withheld is True
assert filter_ptu_shared_deployments([_TERMLESS_SHARED], "team-a").withheld is False
ceiling: Final = _unaliased_ceiling([_TERMLESS_SHARED], "team-a", "gpt-4.1-ptu")
assert ceiling is not None
assert ceiling.tpm_limit == 30 * _GPT41.input_tpm_per_ptu
@pytest.mark.parametrize("team_id", ["team-a", "team-b"])
def test_a_team_holding_a_share_is_served_from_it_and_nothing_else_in_the_group(team_id: str):
servable: Final = team_servable_deployments([_single_team(), _OPEN, _shared()], team_id)
assert [d["model_info"]["id"] for d in servable] == ["shared"]
def test_a_team_holding_no_share_gets_the_open_deployment_and_never_another_teams_reservation():
servable: Final = team_servable_deployments([_single_team(), _OPEN, _shared()], "team-c")
assert [d["model_info"]["id"] for d in servable] == ["open"]
def test_a_share_converts_to_the_models_input_tpm_per_ptu():
ceiling: Final = _unaliased_ceiling([_shared()], "team-a", "gpt-4.1-ptu")
assert ceiling == PTUTeamCeiling(
model_group="gpt-4.1-ptu",
tpm_limit=30 * _GPT41.input_tpm_per_ptu,
output_to_input_ratio=_GPT41.output_to_input_ratio,
cached_input_ratio=_GPT41.cached_input_ratio,
)
def test_shares_across_deployments_add_up_and_the_larger_output_ratio_wins():
gpt4o: Final = _shared(model="azure/gpt-4o", shares={"team-a": 10}, deployment_id="shared-4o")
ceiling: Final = _unaliased_ceiling([_shared(), gpt4o, _OPEN], "team-a", "gpt-4.1-ptu")
assert ceiling is not None
assert ceiling.tpm_limit == 30 * _GPT41.input_tpm_per_ptu + 10 * _GPT4O.input_tpm_per_ptu
assert ceiling.output_to_input_ratio == max(_GPT41.output_to_input_ratio, _GPT4O.output_to_input_ratio)
def test_the_larger_cached_input_ratio_wins_across_deployments():
"""A team sharing two models is weighted by the one that charges more for cache reads,
whichever order the deployments come in."""
gpt6sol: Final = _shared(model="azure/gpt-6-sol", shares={"team-a": 10}, deployment_id="shared-6")
ceiling: Final = _unaliased_ceiling([_shared(), gpt6sol], "team-a", "gpt-4.1-ptu")
assert ceiling is not None
assert _GPT41.cached_input_ratio < _GPT6SOL.cached_input_ratio
assert ceiling.cached_input_ratio == _GPT6SOL.cached_input_ratio
def test_a_group_is_served_by_name_or_by_a_team_scoped_deployments_public_name():
"""A deployment registered for one team is renamed to a unique internal name and keeps
the name callers use in ``team_public_model_name``."""
team_scoped: Final = {
"model_name": "gpt-4.1-ptu-3f9c1b",
"litellm_params": {"model": "azure/gpt-4.1"},
"model_info": {"id": "team-scoped", "team_id": "team-a", "team_public_model_name": "gpt-4.1-ptu"},
}
other: Final = {"model_name": "other", "litellm_params": {"model": "azure/gpt-4o"}, "model_info": {"id": "other"}}
deployments: Final = [team_scoped, _OPEN, other]
served: Final = model_group_deployments(deployments, "gpt-4.1-ptu")
assert [d["model_info"]["id"] for d in served] == ["team-scoped", "open"]
assert model_group_deployments(deployments, "gpt-4.1-ptu-3f9c1b") == (team_scoped,)
assert model_group_deployments(deployments, "missing") == ()
def test_no_share_or_no_sizing_row_sets_no_ceiling():
assert _unaliased_ceiling([_shared()], "team-c", "gpt-4.1-ptu") is None
assert _unaliased_ceiling([_shared(model="azure/unknown-deployment")], "team-a", "gpt-4.1-ptu") is None
assert _unaliased_ceiling([_single_team(), _OPEN], "team-a", "gpt-4.1-ptu") is None
def test_naming_a_shared_deployment_by_id_or_provider_model_draws_on_its_groups_ceiling():
"""The router serves a deployment id or a provider model string when no group has that
name, so those names share the group's ceiling instead of bypassing it."""
payg: Final = {"model_name": "gpt-4.1-payg", "litellm_params": {"model": "azure/gpt-4.1"}, "model_info": {"id": "payg"}}
deployments: Final = [payg, _shared(), _OPEN]
by_group: Final = _unaliased_ceiling(deployments, "team-a", "gpt-4.1-ptu")
assert by_group is not None
assert by_group.model_group == "gpt-4.1-ptu"
assert _unaliased_ceiling(deployments, "team-a", "shared") == by_group
assert _unaliased_ceiling(deployments, "team-a", "azure/gpt-4.1") == by_group
assert _unaliased_ceiling(deployments, "team-a", "payg") is None
assert _unaliased_ceiling(deployments, "team-a", "missing") is None
def test_a_group_name_wins_over_a_deployment_id_it_collides_with():
"""The router routes a name that is both a group and a deployment id to the group."""
colliding: Final = {
"model_name": "shared",
"litellm_params": {"model": "azure/gpt-4o"},
"model_info": {"id": "colliding"},
}
assert _unaliased_ceiling([_shared(), colliding], "team-a", "shared") is None
def test_a_team_scoped_deployment_named_by_id_draws_on_its_public_groups_ceiling():
team_scoped: Final = {
"model_name": "gpt-4.1-ptu-3f9c1b",
"litellm_params": {"model": "azure/gpt-4.1"},
"model_info": {**_shared()["model_info"], "id": "team-scoped", "team_public_model_name": "gpt-4.1-ptu"},
}
ceiling: Final = _unaliased_ceiling([team_scoped], "team-a", "team-scoped")
assert ceiling is not None
assert ceiling.model_group == "gpt-4.1-ptu"
assert ceiling == _unaliased_ceiling([team_scoped], "team-a", "gpt-4.1-ptu")
def test_alias_and_routing_group_copies_do_not_split_a_deployments_ceiling():
"""The router lists alias and routing-group copies of a deployment under their own names
ahead of its own rows, so every name still resolves to the deployment's group."""
shared: Final = _shared()
listed: Final = [{**shared, "model_name": "ptu-alias"}, {**shared, "model_name": "ptu-routing-group"}, shared]
by_group: Final = team_ptu_ceiling(listed, [shared], "team-a", "gpt-4.1-ptu", _no_wildcard_route)
assert by_group is not None
assert by_group.model_group == "gpt-4.1-ptu"
for name in ("shared", "azure/gpt-4.1", "ptu-routing-group"):
assert team_ptu_ceiling(listed, [shared], "team-a", name, _no_wildcard_route) == by_group
def test_a_provider_model_draws_on_the_group_where_the_team_holds_its_share():
"""Two groups share deployments of one provider model among different teams, and the router
serves each team only the one it holds a share of."""
east: Final = _shared(shares={"team-a": 30}, deployment_id="east")
west: Final = {**_shared(shares={"team-b": 20}, deployment_id="west"), "model_name": "gpt-4.1-ptu-west"}
by_provider_model: Final = _unaliased_ceiling([east, west], "team-b", "azure/gpt-4.1")
assert by_provider_model is not None
assert by_provider_model.model_group == "gpt-4.1-ptu-west"
assert by_provider_model == _unaliased_ceiling([east, west], "team-b", "gpt-4.1-ptu-west")
assert _unaliased_ceiling([east, west], "team-a", "azure/gpt-4.1") == _unaliased_ceiling(
[east, west], "team-a", "gpt-4.1-ptu"
)
def test_a_groups_capacity_comes_from_its_first_reserved_deployment_with_a_row():
day: Final = date(2026, 9, 23)
assert model_group_ptu_capacity([_OPEN, _single_team()], day) is _GPT41
assert model_group_ptu_capacity([_shared(model="azure/unknown"), _shared(model="azure/gpt-4o")], day) is _GPT4O
assert model_group_ptu_capacity([_OPEN], day) is None
assert model_group_ptu_capacity([], day) is None
def test_a_groups_capacity_on_a_day_comes_from_the_reservation_covering_that_day():
ended: Final = {**_shared(), "model_info": {**_shared()["model_info"], "ptu_effective_to": "2026-06-01T00:00:00Z"}}
current: Final = {
**_shared(model="azure/gpt-4o", deployment_id="current"),
"model_info": {**_shared(deployment_id="current")["model_info"], "ptu_effective_from": "2026-06-01T00:00:00Z"},
}
assert model_group_ptu_capacity([ended, current], date(2026, 5, 31)) is _GPT41
assert model_group_ptu_capacity([ended, current], date(2026, 6, 1)) is _GPT4O
assert model_group_ptu_capacity([ended], date(2026, 6, 1)) is None
assert model_group_ptu_capacity([current], date(2025, 12, 31)) is None
def test_a_reserved_deployment_without_a_sizing_row_is_warned_about_by_name():
warning: Final = ptu_capacity_warning("gpt-4.1-ptu", _shared(model="azure/my-ptu-deployment"))
assert warning is not None
assert "gpt-4.1-ptu" in warning
assert "base_model" in warning
def test_a_sized_reservation_and_an_unreserved_deployment_raise_no_warning():
assert ptu_capacity_warning("gpt-4.1-ptu", _shared()) is None
assert ptu_capacity_warning("gpt-4.1-ptu", _single_team()) is None
unsized_open: Final = {**_OPEN, "litellm_params": {"model": "azure/my-ptu-deployment"}}
assert ptu_capacity_warning("gpt-4.1-ptu", unsized_open) is None
def test_an_unsized_single_team_azure_reservation_is_warned_about_its_ptu_hours_only():
warning: Final = ptu_capacity_warning("gpt-4.1-ptu", _single_team(model="azure/my-ptu-deployment"))
assert warning is not None
assert "PTU hours" in warning
assert "ceiling" not in warning
def test_an_unsized_single_team_reservation_on_another_provider_is_not_warned_about():
assert ptu_capacity_warning("claude-ptu", _single_team(model="anthropic/claude-sonnet-4-5")) is None
def test_a_bare_model_name_counts_as_azure_through_custom_llm_provider():
deployment: Final = {
**_single_team(model="my-ptu-deployment"),
"litellm_params": {"model": "my-ptu-deployment", "custom_llm_provider": "azure"},
}
warning: Final = ptu_capacity_warning("gpt-4.1-ptu", deployment)
assert warning is not None
assert "PTU hours" in warning

View file

@ -7362,6 +7362,48 @@ def test_get_wildcard_deployment_usable_by_team_prefers_the_team_pattern():
assert router.get_credential_deployment(model_id=ocr, team_id="team-b").model_info.id == "shared-wildcard"
@pytest.mark.parametrize(
("with_catch_all", "requested", "team_id", "expected_ids"),
[
pytest.param(False, "rwild-chat", "team-r", ("team-r-wildcard",), id="team-pattern-when-no-proxy-wide-match"),
pytest.param(False, "rwild-chat", "team-b", (), id="another-teams-pattern-never-serves"),
pytest.param(False, "rwild-chat", None, (), id="no-team-no-team-pattern"),
pytest.param(True, "rwild-chat", "team-r", ("catch-all",), id="proxy-wide-pattern-beats-the-team-pattern"),
pytest.param(True, "ptu-chat", "team-r", ("ptu-wildcard",), id="only-the-most-specific-proxy-wide-pattern"),
],
)
def test_wildcard_route_deployments_is_the_one_pattern_routing_serves_from(
with_catch_all: bool, requested: str, team_id: str | None, expected_ids: tuple[str, ...]
):
catch_all: Final = (
({"model_name": "*", "litellm_params": {"model": "openai/*", "api_key": "sk-open"}, "model_info": {"id": "catch-all"}},)
if with_catch_all
else ()
)
router: Final = litellm.Router(
model_list=[
{
"model_name": "ptu-*",
"litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "sk-ptu"},
"model_info": {"id": "ptu-wildcard"},
},
{
"model_name": "team-r-internal",
"litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "sk-team-r"},
"model_info": {"id": "team-r-wildcard", "team_id": "team-r", "team_public_model_name": "rwild-*"},
},
*catch_all,
]
)
served: Final = router.wildcard_route_deployments(requested, team_id)
early: Final = router._try_early_resolve_deployments_for_model_not_in_names(model=requested, request_team_id=team_id)
early_deployments: Final = early[1] if early is not None and isinstance(early[1], list) else []
assert tuple(Deployment.model_validate(d).model_info.id for d in served) == expected_ids
assert tuple(Deployment.model_validate(d).model_info.id for d in early_deployments) == expected_ids
def test_get_deployment_credentials_with_provider_aws_bedrock_runtime_endpoint():
"""
Test that get_deployment_credentials_with_provider correctly copies
@ -8899,6 +8941,44 @@ async def test_team_scoped_model_fallback_cross_team_blocked():
)
@pytest.mark.asyncio
def _shared_ptu_deployment(model_name: str, deployment_id: str, **litellm_params: str) -> dict[str, object]:
"""A one-PTU Azure deployment whose whole reservation is team-a's share."""
return {
"model_name": model_name,
"litellm_params": {"model": "azure/gpt-4.1", "api_key": "fake", "api_base": "https://ptu.example", **litellm_params},
"model_info": {
"id": deployment_id,
"base_model": "azure/gpt-4.1",
"ptu_count": 1,
"cost_per_ptu_per_hour": 1.0,
"ptu_effective_from": "2026-01-01T00:00:00Z",
"ptu_shares": {"team-a": 1},
},
}
async def test_a_ptu_share_holders_failed_call_is_retried_on_its_shared_deployment(monkeypatch):
"""A shared PTU deployment is only servable to the teams holding a share, so the retry loop
has to look it up as the calling team; looked up with no team it finds no healthy deployment
and gives up after the first failure."""
monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "true")
router = litellm.Router(
model_list=[_shared_ptu_deployment("gpt-4.1-ptu", "shared-ptu", mock_response="litellm.InternalServerError")],
num_retries=2,
retry_after=0,
)
with pytest.raises(litellm.InternalServerError) as failed:
await router.acompletion(
model="gpt-4.1-ptu",
messages=[{"role": "user", "content": "Hello"}],
metadata={"user_api_key_team_id": "team-a"},
)
assert failed.value.num_retries == 2
def test_get_all_deployments_with_team_id():
"""
Test that _get_all_deployments with team_id can find deployments
@ -15958,6 +16038,28 @@ async def test_anthropic_messages_stream_dropped_before_content_is_retried_withi
assert [(attempted, budget) for _, attempted, budget in provider.calls] == [(0, 2), (1, 2)]
@pytest.mark.asyncio
async def test_a_ptu_share_holders_messages_stream_dropped_before_content_is_retried_on_its_shared_group(monkeypatch):
"""The same-group retry has to look the group's health up as the calling team; looked up with no
team a shared PTU group looks empty, so the share holder's dropped stream is never retried."""
monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "true")
router: Final = Router(
model_list=[_shared_ptu_deployment("glm", f"shared-ptu-{i}") for i in (1, 2)],
num_retries=2,
)
provider: Final = _AnthropicMessagesScriptedProvider(
_anthropic_messages_dropped_before_content, _anthropic_messages_retried_stream
)
stream: Final = await _anthropic_messages_stream_through_router(
router, provider, litellm_metadata={"user_api_key_team_id": "team-a"}
)
body: Final = [chunk async for chunk in stream]
assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")]
assert len(provider.calls) == 2
@pytest.mark.asyncio
async def test_anthropic_messages_stream_dropped_after_content_keeps_the_error_and_is_not_retried():
"""A drop once content reached the client cannot be retried without a second overlapping message
@ -20516,6 +20618,27 @@ async def test_prioritized_request_leaves_queue_when_it_stops_waiting(stop_waiti
assert await router.scheduler.get_queue("sched-model") == [(0, "head-of-queue")]
@pytest.mark.asyncio
async def test_a_ptu_share_holders_prioritized_request_is_admitted_while_its_shared_deployment_is_healthy(monkeypatch):
"""The scheduler has to look the group's health up as the calling team; looked up with no team a
shared PTU deployment looks unhealthy, so the share holder waits behind the queue until it times out."""
monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "true")
router: Final = Router(
model_list=[_shared_ptu_deployment("sched-ptu", "shared-ptu", mock_response="hi")],
timeout=0.5,
)
await router.scheduler.add_request(FlowItem(priority=0, request_id="head-of-queue", model_name="sched-ptu"))
response: Final = await router.acompletion(
model="sched-ptu",
messages=[{"role": "user", "content": "hi"}],
priority=5,
metadata={"user_api_key_team_id": "team-a"},
)
assert response.choices[0].message.content == "hi"
@pytest.mark.asyncio
async def test_bare_model_group_served_by_wildcard_deployment_uses_provider_prefixed_fallback_key() -> None:
"""Claude Code sends the bare "claude-sonnet-4-6" to /v1/messages; routing serves it through the
@ -20608,6 +20731,128 @@ def test_bare_model_group_served_by_wildcard_deployment_has_provider_prefixed_co
assert router._has_content_policy_fallback("claude-haiku-4-5", {}) is False
def _shared_ptu_model_list() -> list:
return [
{
"model_name": "gpt-4.1-ptu",
"litellm_params": {"model": "gpt-4.1", "mock_response": "shared"},
"model_info": {
"id": "shared-deployment",
"base_model": "azure/gpt-4.1",
"ptu_count": 50,
"cost_per_ptu_per_hour": 1.0,
"ptu_effective_from": "2026-01-01T00:00:00Z",
"ptu_shares": {"team-a": 30, "team-b": 20},
},
},
{
"model_name": "gpt-4.1-ptu",
"litellm_params": {"model": "gpt-4.1", "mock_response": "open"},
"model_info": {"id": "open-deployment"},
},
]
def test_ptu_shares_hide_the_shared_deployment_from_teams_holding_no_share(monkeypatch):
monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True")
router = Router(model_list=_shared_ptu_model_list())
for request_kwargs in ({"metadata": {"user_api_key_team_id": "team-c"}}, {"metadata": {}}, {}):
_, deployments = router._common_checks_available_deployment(model="gpt-4.1-ptu", request_kwargs=request_kwargs)
assert [d["model_info"]["id"] for d in deployments] == ["open-deployment"]
for team_id in ("team-a", "team-b"):
_, deployments = router._common_checks_available_deployment(
model="gpt-4.1-ptu",
request_kwargs={"metadata": {"user_api_key_team_id": team_id}},
)
assert [d["model_info"]["id"] for d in deployments] == ["shared-deployment"]
def test_ptu_shares_raise_when_only_shared_deployments_remain(monkeypatch):
monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True")
router = Router(model_list=_shared_ptu_model_list()[:1])
with pytest.raises(litellm.BadRequestError, match="reserved for the teams holding a PTU share"):
router._common_checks_available_deployment(
model="gpt-4.1-ptu",
request_kwargs={"metadata": {"user_api_key_team_id": "team-c"}},
)
_, deployments = router._common_checks_available_deployment(
model="gpt-4.1-ptu",
request_kwargs={"metadata": {"user_api_key_team_id": "team-b"}},
)
assert [d["model_info"]["id"] for d in deployments] == ["shared-deployment"]
def test_team_ptu_ceiling_follows_shares_moved_on_the_live_model_list():
router: Final = Router(model_list=_shared_ptu_model_list()[:1])
team_a_before: Final = router.team_ptu_ceiling("team-a", "gpt-4.1-ptu")
assert router.team_ptu_ceiling("team-c", "gpt-4.1-ptu") is None
moved: Final = _shared_ptu_model_list()[0]
moved["model_info"]["ptu_shares"] = {"team-a": 10, "team-c": 40}
router.upsert_deployment(Deployment(**moved))
team_a_after: Final = router.team_ptu_ceiling("team-a", "gpt-4.1-ptu")
team_c_after: Final = router.team_ptu_ceiling("team-c", "gpt-4.1-ptu")
assert team_a_before is not None and team_a_after is not None and team_c_after is not None
assert team_a_before.tpm_limit == 3 * team_a_after.tpm_limit
assert team_c_after.tpm_limit == 4 * team_a_after.tpm_limit
router.delete_deployment(id="shared-deployment")
assert router.team_ptu_ceiling("team-a", "gpt-4.1-ptu") is None
def test_ptu_shares_hide_the_shared_deployment_even_while_the_feature_is_off(monkeypatch):
"""The flag switches cost attribution on; a declared split is an access rule and holds
without it, so a team never reaches a deployment reserved for others while the flag is
off."""
monkeypatch.delenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", raising=False)
router = Router(model_list=_shared_ptu_model_list())
_, deployments = router._common_checks_available_deployment(
model="gpt-4.1-ptu",
request_kwargs={"metadata": {"user_api_key_team_id": "team-c"}},
)
assert [d["model_info"]["id"] for d in deployments] == ["open-deployment"]
def test_a_shared_ptu_deployment_whose_shares_do_not_add_up_is_refused_at_registration(monkeypatch):
monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True")
model_list = _shared_ptu_model_list()[:1]
model_list[0]["model_info"]["ptu_shares"] = {"team-a": 30}
with pytest.raises(ValueError, match=r"gpt-4\.1-ptu.*30 of 50 allocated"):
Router(model_list=model_list)
def test_a_shared_ptu_deployment_with_a_fractional_count_is_refused_at_registration(monkeypatch):
monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True")
model_list = _shared_ptu_model_list()[:1]
model_list[0]["model_info"]["ptu_count"] = 50.5
with pytest.raises(ValueError, match="ptu_count"):
Router(model_list=model_list)
@pytest.mark.parametrize(("flag", "warned"), [("True", True), ("", False)])
def test_an_unsized_shared_ptu_deployment_is_warned_about_at_registration_only_with_the_feature_on(
monkeypatch, caplog, flag, warned
):
monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", flag)
model_list = _shared_ptu_model_list()[:1]
model_list[0]["litellm_params"]["model"] = "azure/my-ptu-deployment"
del model_list[0]["model_info"]["base_model"]
with caplog.at_level(logging.WARNING, logger="LiteLLM Router"):
Router(model_list=model_list)
assert ("'gpt-4.1-ptu'" in caplog.text and "base_model" in caplog.text) is warned
def test_a_config_entry_declaring_shares_without_terms_is_refused_even_while_the_feature_is_off(monkeypatch):
"""The split is enforced with the flag off, so its shape is checked with the flag off too: a
malformed one is refused at boot instead of being honoured or ignored without a word."""
monkeypatch.delenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", raising=False)
model_list = _shared_ptu_model_list()[:1]
model_list[0]["model_info"] = {"id": "shared-deployment", "ptu_shares": {"team-a": 30, "team-b": 20}}
with pytest.raises(ValueError, match=r"ptu_count and cost_per_ptu_per_hour are required when ptu_shares is set"):
Router(model_list=model_list)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"verdict",

View file

@ -2659,7 +2659,7 @@ def test_an_incomplete_reservation_is_refused_rather_than_served(dropped):
@pytest.mark.parametrize(
"dropped, expected",
[
("team_id", "team_id is required when PTU fields are set (one model maps to one team)"),
("team_id", "team_id or ptu_shares is required when PTU fields are set"),
("cost_per_ptu_per_hour", "ptu_count and cost_per_ptu_per_hour must be set together"),
],
ids=["no team_id", "count without rate"],

View file

@ -2,7 +2,13 @@ import useTeams from "@/app/(dashboard)/hooks/useTeams";
import { StackedUsageChart, type StackedUsageScale } from "@/components/shared/charts";
import { DataTable } from "@/components/shared/DataTable";
import { getProviderSpend, getTopAgents, getTopAPIKeys, getTopModels } from "./entityUsageAggregations";
import { buildCostBreakdownTiles, buildSummaryTiles, hasFlatCost, type SummaryTile } from "./entityUsageSummary";
import {
buildCostBreakdownTiles,
buildSummaryTiles,
hasFlatCost,
hasPtuHours,
type SummaryTile,
} from "./entityUsageSummary";
import { MoneyCell } from "@/components/shared/table_cells";
import { hasCapability, type Capability } from "@/utils/capabilities";
import type { DateRangePickerValue } from "@/components/shared/date_picker_types";
@ -340,6 +346,7 @@ const EntityUsage: React.FC<EntityUsageProps> = ({
const capitalizedEntityLabel = entityType.charAt(0).toUpperCase() + entityType.slice(1);
const showFlatCost = entityType === "team" && hasFlatCost(spendData.metadata);
const showPtuHours = entityType === "team" && hasPtuHours(spendData.metadata);
const userSpendTeamIds = useMemo(
() =>
selectedTags.length > 0
@ -449,7 +456,7 @@ const EntityUsage: React.FC<EntityUsageProps> = ({
};
const breakdownTiles = showFlatCost && showCostBreakdown ? buildCostBreakdownTiles(spendData.metadata) : [];
const summaryTiles = buildSummaryTiles(spendData.metadata, showFlatCost);
const summaryTiles = buildSummaryTiles(spendData.metadata, showFlatCost, showPtuHours);
const modelViewTitle = modelViewType === "groups" ? "Top Public Model Names" : "Top Litellm Models";
const chartFormat = (value: number) => formatMetricValue(value, breakdown.metric);
@ -460,7 +467,12 @@ const EntityUsage: React.FC<EntityUsageProps> = ({
<div className="grid gap-3">
<section className="grid gap-2">
<h2 className="text-sm font-medium text-foreground">{capitalizedEntityLabel} Spend Overview</h2>
<div className="grid grid-cols-2 gap-px overflow-hidden rounded-xl border bg-border lg:grid-cols-5">
<div
className={cn(
"grid grid-cols-2 gap-px overflow-hidden rounded-xl border bg-border",
showPtuHours ? "lg:grid-cols-6" : "lg:grid-cols-5",
)}
>
{summaryTiles.map(renderSummaryTile)}
</div>
{breakdownTiles.length > 0 && (

View file

@ -1,5 +1,5 @@
import { describe, expect, it } from "vitest";
import { buildCostBreakdownTiles, buildSummaryTiles, hasFlatCost } from "./entityUsageSummary";
import { buildCostBreakdownTiles, buildSummaryTiles, hasFlatCost, hasPtuHours } from "./entityUsageSummary";
const metadata = {
total_spend: 100,
@ -10,6 +10,17 @@ const metadata = {
total_tokens: 3456,
};
describe("hasPtuHours", () => {
it("is false when no PTU hours were consumed or the backend predates the field", () => {
expect(hasPtuHours({ ...metadata, total_ptu_hours: 0 })).toBe(false);
expect(hasPtuHours(metadata)).toBe(false);
});
it("is true once a PTU model group has been used", () => {
expect(hasPtuHours({ ...metadata, total_ptu_hours: 0.25 })).toBe(true);
});
});
describe("hasFlatCost", () => {
it("is false when there is no flat cost to report", () => {
expect(hasFlatCost({ ...metadata, total_flat_cost: 0 })).toBe(false);
@ -28,6 +39,23 @@ describe("buildSummaryTiles", () => {
expect(buildSummaryTiles(metadata, true)).toHaveLength(5);
});
it("adds a sixth PTU Hours tile only when asked, after Total Tokens", () => {
const withPtu = { ...metadata, total_ptu_hours: 1234.5678 };
expect(buildSummaryTiles(withPtu, true)).toHaveLength(5);
const tiles = buildSummaryTiles(withPtu, true, true);
expect(tiles).toHaveLength(6);
const last = tiles[tiles.length - 1];
expect(last.title).toBe("PTU Hours");
expect(last.value).toBe("1,234.57");
expect(last.tooltip).toContain("PTU-hours");
expect(tiles[tiles.length - 2].title).toBe("Total Tokens");
});
it("treats missing PTU hours as zero when the tile is forced on", () => {
const tiles = buildSummaryTiles(metadata, false, true);
expect(tiles[tiles.length - 1].value).toBe("0.00");
});
it("shows request-only spend under the original title when there is no flat cost", () => {
const [first] = buildSummaryTiles(metadata, false);
expect(first.title).toBe("Total Spend");

View file

@ -11,6 +11,7 @@ export interface SummaryTile {
interface SpendSummaryMetadata {
total_spend: number;
total_flat_cost?: number;
total_ptu_hours?: number;
total_api_requests: number;
total_successful_requests: number;
total_failed_requests: number;
@ -26,10 +27,28 @@ export const REQUEST_COST_TOOLTIP =
export const FLAT_COST_TOOLTIP =
"Reserved provisioned throughput, billed per hour whether or not requests are sent. Reported here only; it does not count toward team, key, user, or organization budgets.";
export const PTU_HOURS_TOOLTIP =
"Provisioned throughput consumed, in PTU-hours: the tokens sent to PTU deployments converted through the Azure sizing table for each model. Divide it by the PTUs reserved times the hours in the period to see how much of the reservation was used.";
export const hasFlatCost = (metadata: SpendSummaryMetadata): boolean => (metadata.total_flat_cost ?? 0) > 0;
export const buildSummaryTiles = (metadata: SpendSummaryMetadata, showFlatCost: boolean): SummaryTile[] => {
export const hasPtuHours = (metadata: SpendSummaryMetadata): boolean => (metadata.total_ptu_hours ?? 0) > 0;
export const buildSummaryTiles = (
metadata: SpendSummaryMetadata,
showFlatCost: boolean,
showPtuHours: boolean = false,
): SummaryTile[] => {
const flatCost = metadata.total_flat_cost ?? 0;
const ptuHoursTile: SummaryTile[] = showPtuHours
? [
{
title: "PTU Hours",
value: formatNumberWithCommas(metadata.total_ptu_hours ?? 0, 2),
tooltip: PTU_HOURS_TOOLTIP,
},
]
: [];
return [
showFlatCost
? {
@ -47,6 +66,7 @@ export const buildSummaryTiles = (metadata: SpendSummaryMetadata, showFlatCost:
},
{ title: "Failed Requests", value: metadata.total_failed_requests.toLocaleString(), className: "text-destructive" },
{ title: "Total Tokens", value: metadata.total_tokens.toLocaleString() },
...ptuHoursTile,
];
};

View file

@ -34,6 +34,7 @@ export const EMPTY_DAILY_ACTIVITY_METADATA: DailyActivityMetadata = {
total_pages: 1,
total_spend: 0,
total_flat_cost: 0,
total_ptu_hours: 0,
total_api_requests: 0,
total_successful_requests: 0,
total_failed_requests: 0,

View file

@ -1,6 +1,7 @@
export interface SpendMetrics {
spend: number;
flat_cost?: number;
ptu_hours?: number;
prompt_tokens: number;
completion_tokens: number;
total_tokens: number;

View file

@ -31337,6 +31337,11 @@ export interface components {
* @default 0
*/
total_prompt_tokens: number;
/**
* Total Ptu Hours
* @default 0
*/
total_ptu_hours: number;
/**
* Total Response Time Ms
* @default 0
@ -45425,6 +45430,11 @@ export interface components {
* @default 0
*/
prompt_tokens: number;
/**
* Ptu Hours
* @default 0
*/
ptu_hours: number;
/**
* Spend
* @default 0
@ -49929,6 +49939,10 @@ export interface components {
ptu_effective_from?: string | null;
/** Ptu Effective To */
ptu_effective_to?: string | null;
/** Ptu Shares */
ptu_shares?: {
[key: string]: number;
} | null;
/** Team Id */
team_id?: string | null;
/** Team Public Model Name */