mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-16 04:29:31 +00:00
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:
parent
3151c47960
commit
e344dbd8db
50 changed files with 6186 additions and 256 deletions
|
|
@ -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,
|
||||
|
|
|
|||
139
litellm/llms/azure/ptu_capacity.py
Normal file
139
litellm/llms/azure/ptu_capacity.py
Normal 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
|
||||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
145
litellm/proxy/management_endpoints/ptu_consumption.py
Normal file
145
litellm/proxy/management_endpoints/ptu_consumption.py
Normal 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))
|
||||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
30
litellm/proxy/management_helpers/ptu_share_guard.py
Normal file
30
litellm/proxy/management_helpers/ptu_share_guard.py
Normal 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."
|
||||
},
|
||||
)
|
||||
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
275
litellm/router_utils/ptu_shares.py
Normal file
275
litellm/router_utils/ptu_shares.py
Normal 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
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
412
tests/integration/management/test_ptu_share_guards.py
Normal file
412
tests/integration/management/test_ptu_share_guards.py
Normal 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}
|
||||
706
tests/integration/routing/test_ptu_shares.py
Normal file
706
tests/integration/routing/test_ptu_shares.py
Normal 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:]
|
||||
143
tests/integration/routing/test_ptu_shares_regex_timing.py
Normal file
143
tests/integration/routing/test_ptu_shares_regex_timing.py
Normal 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
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
362
tests/integration/spend/test_ptu_share_attribution.py
Normal file
362
tests/integration/spend/test_ptu_share_attribution.py
Normal 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}
|
||||
|
|
@ -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
|
||||
|
|
|
|||
106
tests/unit/llms/azure/test_azure_ptu_capacity.py
Normal file
106
tests/unit/llms/azure/test_azure_ptu_capacity.py
Normal 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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
[
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
381
tests/unit/proxy/management_endpoints/test_ptu_consumption.py
Normal file
381
tests/unit/proxy/management_endpoints/test_ptu_consumption.py
Normal 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)
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
301
tests/unit/router_utils/test_ptu_shares.py
Normal file
301
tests/unit/router_utils/test_ptu_shares.py
Normal 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
|
||||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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"],
|
||||
|
|
|
|||
|
|
@ -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 && (
|
||||
|
|
|
|||
|
|
@ -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");
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
];
|
||||
};
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
export interface SpendMetrics {
|
||||
spend: number;
|
||||
flat_cost?: number;
|
||||
ptu_hours?: number;
|
||||
prompt_tokens: number;
|
||||
completion_tokens: number;
|
||||
total_tokens: number;
|
||||
|
|
|
|||
14
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
14
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -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 */
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue