mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
feat(ptu): surface ptu_flat_cost via /team/daily/activity response
Extends the shared daily-activity read path so team daily activity carries prorated PTU flat cost alongside per-request spend. All changes are additive to the response shape. - SpendMetrics gains flat_cost (float, default 0.0) - DailySpendMetadata gains total_flat_cost (float, default 0.0) - update_metrics + _record_to_spend_metrics read ptu_flat_cost from rollup rows via getattr, so non-team daily tables (user, org, tag, agent, end_user) pass through with flat_cost defaulting to 0 - Grouping-sets SQL query selects SUM(ptu_flat_cost) only for litellm_dailyteamspend; every other daily table emits a 0::float shim so the response shape stays uniform - Sentinel PTU rows (api_key = PTU_SENTINEL_API_KEY) contribute to per-day, per-model, per-provider, per-endpoint, and per-entity totals, but never appear in api_keys / api_key_breakdown maps at any level - get_api_key_metadata skips the sentinel in its Prisma lookup PTU_SENTINEL_API_KEY and PTU_ROLLUP_JOB_ID moved to litellm/constants.py. The rollup module re-exports both so external callers continue working. Addresses Greptile P2 on #33137. Regenerated ui/litellm-dashboard/src/lib/http/schema.d.ts. 15 new unit tests, 8 kill mutations that either drop the ptu_flat_cost accumulation or turn off the sentinel filter.
This commit is contained in:
parent
1a9899091a
commit
48d46cc354
6 changed files with 437 additions and 89 deletions
|
|
@ -1628,3 +1628,9 @@ ADVISOR_TOOL_DESCRIPTION: str = (
|
|||
"want to verify your reasoning, or face a complex decision. "
|
||||
"Describe your question or challenge clearly in the 'question' field."
|
||||
)
|
||||
|
||||
# PTU reservation rollup writes rows to LiteLLM_DailyTeamSpend with this
|
||||
# sentinel api_key so PTU flat cost stays distinguishable from real per-request
|
||||
# spend under the table's composite unique constraint.
|
||||
PTU_SENTINEL_API_KEY: str = "__ptu_reservation__"
|
||||
PTU_ROLLUP_JOB_ID: str = "ptu_reservation_rollup_job"
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ from typing import Any, Awaitable, Callable, Dict, List, Optional, Set, Tuple, U
|
|||
from fastapi import HTTPException, status
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import PTU_SENTINEL_API_KEY
|
||||
from litellm.proxy._types import CommonProxyErrors
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
from litellm.repositories.table_repositories import DeletedVerificationTokenRepository
|
||||
|
|
@ -44,6 +45,7 @@ def update_metrics(existing_metrics: SpendMetrics, record: Any) -> SpendMetrics:
|
|||
prompt_tokens = record.prompt_tokens or 0
|
||||
completion_tokens = record.completion_tokens or 0
|
||||
existing_metrics.spend += record.spend or 0.0
|
||||
existing_metrics.flat_cost += getattr(record, "ptu_flat_cost", None) or 0.0
|
||||
existing_metrics.prompt_tokens += prompt_tokens
|
||||
existing_metrics.completion_tokens += completion_tokens
|
||||
existing_metrics.total_tokens += prompt_tokens + completion_tokens
|
||||
|
|
@ -98,7 +100,15 @@ def update_breakdown_metrics(
|
|||
entity_id_field: Optional[str] = None,
|
||||
entity_metadata_field: Optional[Dict[str, dict]] = None,
|
||||
) -> BreakdownMetrics:
|
||||
"""Updates breakdown metrics for a single record using the existing update_metrics function"""
|
||||
"""Updates breakdown metrics for a single record using the existing update_metrics function.
|
||||
|
||||
PTU sentinel rows (``api_key == PTU_SENTINEL_API_KEY``) contribute their
|
||||
``ptu_flat_cost`` to every parent breakdown (per-model, per-provider,
|
||||
per-endpoint, per-entity, per-day totals) but never appear as a row in
|
||||
any ``api_keys`` / ``api_key_breakdown`` map — the sentinel string is not
|
||||
a real key alias.
|
||||
"""
|
||||
is_ptu_sentinel = record.api_key == PTU_SENTINEL_API_KEY
|
||||
|
||||
# Update model breakdown
|
||||
if record.model and record.model not in breakdown.models:
|
||||
|
|
@ -110,18 +120,19 @@ def update_breakdown_metrics(
|
|||
breakdown.models[record.model].metrics = update_metrics(breakdown.models[record.model].metrics, record)
|
||||
|
||||
# Update API key breakdown for this model
|
||||
if record.api_key not in breakdown.models[record.model].api_key_breakdown:
|
||||
breakdown.models[record.model].api_key_breakdown[record.api_key] = KeyMetricWithMetadata(
|
||||
metrics=SpendMetrics(),
|
||||
metadata=KeyMetadata(
|
||||
key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None),
|
||||
team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None),
|
||||
),
|
||||
if not is_ptu_sentinel:
|
||||
if record.api_key not in breakdown.models[record.model].api_key_breakdown:
|
||||
breakdown.models[record.model].api_key_breakdown[record.api_key] = KeyMetricWithMetadata(
|
||||
metrics=SpendMetrics(),
|
||||
metadata=KeyMetadata(
|
||||
key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None),
|
||||
team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None),
|
||||
),
|
||||
)
|
||||
breakdown.models[record.model].api_key_breakdown[record.api_key].metrics = update_metrics(
|
||||
breakdown.models[record.model].api_key_breakdown[record.api_key].metrics,
|
||||
record,
|
||||
)
|
||||
breakdown.models[record.model].api_key_breakdown[record.api_key].metrics = update_metrics(
|
||||
breakdown.models[record.model].api_key_breakdown[record.api_key].metrics,
|
||||
record,
|
||||
)
|
||||
|
||||
# Update model group breakdown
|
||||
if record.model_group and record.model_group not in breakdown.model_groups:
|
||||
|
|
@ -135,18 +146,19 @@ def update_breakdown_metrics(
|
|||
)
|
||||
|
||||
# Update API key breakdown for this model
|
||||
if record.api_key not in breakdown.model_groups[record.model_group].api_key_breakdown:
|
||||
breakdown.model_groups[record.model_group].api_key_breakdown[record.api_key] = KeyMetricWithMetadata(
|
||||
metrics=SpendMetrics(),
|
||||
metadata=KeyMetadata(
|
||||
key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None),
|
||||
team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None),
|
||||
),
|
||||
if not is_ptu_sentinel:
|
||||
if record.api_key not in breakdown.model_groups[record.model_group].api_key_breakdown:
|
||||
breakdown.model_groups[record.model_group].api_key_breakdown[record.api_key] = KeyMetricWithMetadata(
|
||||
metrics=SpendMetrics(),
|
||||
metadata=KeyMetadata(
|
||||
key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None),
|
||||
team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None),
|
||||
),
|
||||
)
|
||||
breakdown.model_groups[record.model_group].api_key_breakdown[record.api_key].metrics = update_metrics(
|
||||
breakdown.model_groups[record.model_group].api_key_breakdown[record.api_key].metrics,
|
||||
record,
|
||||
)
|
||||
breakdown.model_groups[record.model_group].api_key_breakdown[record.api_key].metrics = update_metrics(
|
||||
breakdown.model_groups[record.model_group].api_key_breakdown[record.api_key].metrics,
|
||||
record,
|
||||
)
|
||||
|
||||
if record.mcp_namespaced_tool_name:
|
||||
if record.mcp_namespaced_tool_name not in breakdown.mcp_servers:
|
||||
|
|
@ -159,23 +171,24 @@ def update_breakdown_metrics(
|
|||
)
|
||||
|
||||
# Update API key breakdown for this MCP server
|
||||
if record.api_key not in breakdown.mcp_servers[record.mcp_namespaced_tool_name].api_key_breakdown:
|
||||
breakdown.mcp_servers[record.mcp_namespaced_tool_name].api_key_breakdown[record.api_key] = (
|
||||
KeyMetricWithMetadata(
|
||||
metrics=SpendMetrics(),
|
||||
metadata=KeyMetadata(
|
||||
key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None),
|
||||
team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None),
|
||||
),
|
||||
if not is_ptu_sentinel:
|
||||
if record.api_key not in breakdown.mcp_servers[record.mcp_namespaced_tool_name].api_key_breakdown:
|
||||
breakdown.mcp_servers[record.mcp_namespaced_tool_name].api_key_breakdown[record.api_key] = (
|
||||
KeyMetricWithMetadata(
|
||||
metrics=SpendMetrics(),
|
||||
metadata=KeyMetadata(
|
||||
key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None),
|
||||
team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None),
|
||||
),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
breakdown.mcp_servers[record.mcp_namespaced_tool_name].api_key_breakdown[
|
||||
record.api_key
|
||||
].metrics = update_metrics(
|
||||
breakdown.mcp_servers[record.mcp_namespaced_tool_name].api_key_breakdown[record.api_key].metrics,
|
||||
record,
|
||||
)
|
||||
breakdown.mcp_servers[record.mcp_namespaced_tool_name].api_key_breakdown[
|
||||
record.api_key
|
||||
].metrics = update_metrics(
|
||||
breakdown.mcp_servers[record.mcp_namespaced_tool_name].api_key_breakdown[record.api_key].metrics,
|
||||
record,
|
||||
)
|
||||
|
||||
# Update provider breakdown
|
||||
provider = record.custom_llm_provider or "unknown"
|
||||
|
|
@ -187,18 +200,19 @@ def update_breakdown_metrics(
|
|||
breakdown.providers[provider].metrics = update_metrics(breakdown.providers[provider].metrics, record)
|
||||
|
||||
# Update API key breakdown for this provider
|
||||
if record.api_key not in breakdown.providers[provider].api_key_breakdown:
|
||||
breakdown.providers[provider].api_key_breakdown[record.api_key] = KeyMetricWithMetadata(
|
||||
metrics=SpendMetrics(),
|
||||
metadata=KeyMetadata(
|
||||
key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None),
|
||||
team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None),
|
||||
),
|
||||
if not is_ptu_sentinel:
|
||||
if record.api_key not in breakdown.providers[provider].api_key_breakdown:
|
||||
breakdown.providers[provider].api_key_breakdown[record.api_key] = KeyMetricWithMetadata(
|
||||
metrics=SpendMetrics(),
|
||||
metadata=KeyMetadata(
|
||||
key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None),
|
||||
team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None),
|
||||
),
|
||||
)
|
||||
breakdown.providers[provider].api_key_breakdown[record.api_key].metrics = update_metrics(
|
||||
breakdown.providers[provider].api_key_breakdown[record.api_key].metrics,
|
||||
record,
|
||||
)
|
||||
breakdown.providers[provider].api_key_breakdown[record.api_key].metrics = update_metrics(
|
||||
breakdown.providers[provider].api_key_breakdown[record.api_key].metrics,
|
||||
record,
|
||||
)
|
||||
|
||||
# Update endpoint breakdown
|
||||
if record.endpoint:
|
||||
|
|
@ -212,29 +226,31 @@ def update_breakdown_metrics(
|
|||
)
|
||||
|
||||
# Update API key breakdown for this endpoint
|
||||
if record.api_key not in breakdown.endpoints[record.endpoint].api_key_breakdown:
|
||||
breakdown.endpoints[record.endpoint].api_key_breakdown[record.api_key] = KeyMetricWithMetadata(
|
||||
if not is_ptu_sentinel:
|
||||
if record.api_key not in breakdown.endpoints[record.endpoint].api_key_breakdown:
|
||||
breakdown.endpoints[record.endpoint].api_key_breakdown[record.api_key] = KeyMetricWithMetadata(
|
||||
metrics=SpendMetrics(),
|
||||
metadata=KeyMetadata(
|
||||
key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None),
|
||||
team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None),
|
||||
),
|
||||
)
|
||||
breakdown.endpoints[record.endpoint].api_key_breakdown[record.api_key].metrics = update_metrics(
|
||||
breakdown.endpoints[record.endpoint].api_key_breakdown[record.api_key].metrics,
|
||||
record,
|
||||
)
|
||||
|
||||
# Update api key breakdown
|
||||
if not is_ptu_sentinel:
|
||||
if record.api_key not in breakdown.api_keys:
|
||||
breakdown.api_keys[record.api_key] = KeyMetricWithMetadata(
|
||||
metrics=SpendMetrics(),
|
||||
metadata=KeyMetadata(
|
||||
key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None),
|
||||
team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None),
|
||||
),
|
||||
), # Add any api_key-specific metadata here
|
||||
)
|
||||
breakdown.endpoints[record.endpoint].api_key_breakdown[record.api_key].metrics = update_metrics(
|
||||
breakdown.endpoints[record.endpoint].api_key_breakdown[record.api_key].metrics,
|
||||
record,
|
||||
)
|
||||
|
||||
# Update api key breakdown
|
||||
if record.api_key not in breakdown.api_keys:
|
||||
breakdown.api_keys[record.api_key] = KeyMetricWithMetadata(
|
||||
metrics=SpendMetrics(),
|
||||
metadata=KeyMetadata(
|
||||
key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None),
|
||||
team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None),
|
||||
), # Add any api_key-specific metadata here
|
||||
)
|
||||
breakdown.api_keys[record.api_key].metrics = update_metrics(breakdown.api_keys[record.api_key].metrics, record)
|
||||
breakdown.api_keys[record.api_key].metrics = update_metrics(breakdown.api_keys[record.api_key].metrics, record)
|
||||
|
||||
# Update entity-specific metrics if entity_id_field is provided
|
||||
if entity_id_field:
|
||||
|
|
@ -248,18 +264,19 @@ def update_breakdown_metrics(
|
|||
breakdown.entities[entity_value].metrics = update_metrics(breakdown.entities[entity_value].metrics, record)
|
||||
|
||||
# Update API key breakdown for this entity
|
||||
if record.api_key not in breakdown.entities[entity_value].api_key_breakdown:
|
||||
breakdown.entities[entity_value].api_key_breakdown[record.api_key] = KeyMetricWithMetadata(
|
||||
metrics=SpendMetrics(),
|
||||
metadata=KeyMetadata(
|
||||
key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None),
|
||||
team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None),
|
||||
),
|
||||
if not is_ptu_sentinel:
|
||||
if record.api_key not in breakdown.entities[entity_value].api_key_breakdown:
|
||||
breakdown.entities[entity_value].api_key_breakdown[record.api_key] = KeyMetricWithMetadata(
|
||||
metrics=SpendMetrics(),
|
||||
metadata=KeyMetadata(
|
||||
key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None),
|
||||
team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None),
|
||||
),
|
||||
)
|
||||
breakdown.entities[entity_value].api_key_breakdown[record.api_key].metrics = update_metrics(
|
||||
breakdown.entities[entity_value].api_key_breakdown[record.api_key].metrics,
|
||||
record,
|
||||
)
|
||||
breakdown.entities[entity_value].api_key_breakdown[record.api_key].metrics = update_metrics(
|
||||
breakdown.entities[entity_value].api_key_breakdown[record.api_key].metrics,
|
||||
record,
|
||||
)
|
||||
|
||||
return breakdown
|
||||
|
||||
|
|
@ -401,6 +418,14 @@ def _build_aggregated_sql_query(
|
|||
if pg_table is None:
|
||||
raise ValueError(f"Unknown table name: {table_name}")
|
||||
|
||||
# Only LiteLLM_DailyTeamSpend carries ptu_flat_cost today. Other daily tables
|
||||
# emit a constant zero so the response shape (SpendMetrics.flat_cost) stays
|
||||
# uniform for every entity.
|
||||
has_ptu_flat_cost = table_name == "litellm_dailyteamspend"
|
||||
ptu_flat_cost_select = (
|
||||
"SUM(ptu_flat_cost)::float AS ptu_flat_cost" if has_ptu_flat_cost else "0::float AS ptu_flat_cost"
|
||||
)
|
||||
|
||||
adjusted_start, adjusted_end = _adjust_dates_for_timezone(start_date, end_date, timezone_offset_minutes)
|
||||
|
||||
sql_conditions: List[str] = []
|
||||
|
|
@ -469,6 +494,7 @@ def _build_aggregated_sql_query(
|
|||
custom_llm_provider, mcp_namespaced_tool_name,
|
||||
endpoint) AS group_level,
|
||||
SUM(spend)::float AS spend,
|
||||
{ptu_flat_cost_select},
|
||||
SUM(prompt_tokens)::bigint AS prompt_tokens,
|
||||
SUM(completion_tokens)::bigint AS completion_tokens,
|
||||
SUM(cache_read_input_tokens)::bigint AS cache_read_input_tokens,
|
||||
|
|
@ -560,7 +586,9 @@ async def _aggregate_spend_records(
|
|||
The per-row loop is offloaded to a worker thread via asyncio.to_thread so
|
||||
a large result set doesn't peg the event loop.
|
||||
"""
|
||||
api_keys: Set[str] = {record.api_key for record in records if record.api_key}
|
||||
api_keys: Set[str] = {
|
||||
record.api_key for record in records if record.api_key and record.api_key != PTU_SENTINEL_API_KEY
|
||||
}
|
||||
|
||||
api_key_metadata: Dict[str, Dict[str, Any]] = {}
|
||||
if api_keys:
|
||||
|
|
@ -607,6 +635,7 @@ def _record_to_spend_metrics(record: Any) -> SpendMetrics:
|
|||
completion_tokens = record.completion_tokens or 0
|
||||
return SpendMetrics(
|
||||
spend=record.spend or 0.0,
|
||||
flat_cost=getattr(record, "ptu_flat_cost", None) or 0.0,
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
total_tokens=prompt_tokens + completion_tokens,
|
||||
|
|
@ -669,6 +698,10 @@ def _aggregate_grouping_sets_records_sync(
|
|||
for record in records:
|
||||
level = record.group_level
|
||||
metrics = _record_to_spend_metrics(record)
|
||||
# Sentinel PTU rows contribute to grand-total / per-day / per-model /
|
||||
# per-provider / per-endpoint buckets, but must never appear in any
|
||||
# api_key or api_key_breakdown map.
|
||||
real_api_key = record.api_key and record.api_key != PTU_SENTINEL_API_KEY
|
||||
|
||||
if level == _GROUP_GRAND_TOTAL:
|
||||
total_metrics = metrics
|
||||
|
|
@ -681,7 +714,7 @@ def _aggregate_grouping_sets_records_sync(
|
|||
breakdown = ensure_date(record.date)["breakdown"]
|
||||
|
||||
if level == _GROUP_DATE_API_KEY:
|
||||
if record.api_key:
|
||||
if real_api_key:
|
||||
breakdown.api_keys[record.api_key] = KeyMetricWithMetadata(
|
||||
metrics=metrics,
|
||||
metadata=_key_metadata(api_key_metadata, record.api_key),
|
||||
|
|
@ -690,13 +723,13 @@ def _aggregate_grouping_sets_records_sync(
|
|||
if record.model:
|
||||
assign_metric_with_metadata(breakdown.models, record.model, metrics)
|
||||
elif level == _GROUP_DATE_MODEL_API_KEY:
|
||||
if record.model and record.api_key:
|
||||
if record.model and real_api_key:
|
||||
assign_api_key_breakdown(breakdown.models, record.model, record.api_key, metrics)
|
||||
elif level == _GROUP_DATE_MODEL_GROUP:
|
||||
if record.model_group:
|
||||
assign_metric_with_metadata(breakdown.model_groups, record.model_group, metrics)
|
||||
elif level == _GROUP_DATE_MODEL_GROUP_API_KEY:
|
||||
if record.model_group and record.api_key:
|
||||
if record.model_group and real_api_key:
|
||||
assign_api_key_breakdown(
|
||||
breakdown.model_groups,
|
||||
record.model_group,
|
||||
|
|
@ -707,14 +740,14 @@ def _aggregate_grouping_sets_records_sync(
|
|||
provider = record.custom_llm_provider or "unknown"
|
||||
assign_metric_with_metadata(breakdown.providers, provider, metrics)
|
||||
elif level == _GROUP_DATE_PROVIDER_API_KEY:
|
||||
if record.api_key:
|
||||
if real_api_key:
|
||||
provider = record.custom_llm_provider or "unknown"
|
||||
assign_api_key_breakdown(breakdown.providers, provider, record.api_key, metrics)
|
||||
elif level == _GROUP_DATE_MCP:
|
||||
if record.mcp_namespaced_tool_name:
|
||||
assign_metric_with_metadata(breakdown.mcp_servers, record.mcp_namespaced_tool_name, metrics)
|
||||
elif level == _GROUP_DATE_MCP_API_KEY:
|
||||
if record.mcp_namespaced_tool_name and record.api_key:
|
||||
if record.mcp_namespaced_tool_name and real_api_key:
|
||||
assign_api_key_breakdown(
|
||||
breakdown.mcp_servers,
|
||||
record.mcp_namespaced_tool_name,
|
||||
|
|
@ -725,7 +758,7 @@ def _aggregate_grouping_sets_records_sync(
|
|||
if record.endpoint:
|
||||
assign_metric_with_metadata(breakdown.endpoints, record.endpoint, metrics)
|
||||
elif level == _GROUP_DATE_ENDPOINT_API_KEY:
|
||||
if record.endpoint and record.api_key:
|
||||
if record.endpoint and real_api_key:
|
||||
assign_api_key_breakdown(breakdown.endpoints, record.endpoint, record.api_key, metrics)
|
||||
|
||||
results = [
|
||||
|
|
@ -747,7 +780,7 @@ async def _aggregate_grouping_sets_records(
|
|||
records: List[Any],
|
||||
) -> Dict[str, Any]:
|
||||
"""Async wrapper: fetch api_key_metadata, then dispatch on a worker thread."""
|
||||
api_keys: Set[str] = {r.api_key for r in records if r.api_key}
|
||||
api_keys: Set[str] = {r.api_key for r in records if r.api_key and r.api_key != PTU_SENTINEL_API_KEY}
|
||||
|
||||
api_key_metadata: Dict[str, Dict[str, Any]] = {}
|
||||
if api_keys:
|
||||
|
|
@ -854,6 +887,7 @@ async def get_daily_activity(
|
|||
results=aggregated["results"],
|
||||
metadata=DailySpendMetadata(
|
||||
total_spend=metadata_metrics.spend,
|
||||
total_flat_cost=metadata_metrics.flat_cost,
|
||||
total_prompt_tokens=metadata_metrics.prompt_tokens,
|
||||
total_completion_tokens=metadata_metrics.completion_tokens,
|
||||
total_tokens=metadata_metrics.total_tokens,
|
||||
|
|
@ -940,6 +974,7 @@ async def get_daily_activity_aggregated(
|
|||
results=aggregated["results"],
|
||||
metadata=DailySpendMetadata(
|
||||
total_spend=aggregated["totals"].spend,
|
||||
total_flat_cost=aggregated["totals"].flat_cost,
|
||||
total_prompt_tokens=aggregated["totals"].prompt_tokens,
|
||||
total_completion_tokens=aggregated["totals"].completion_tokens,
|
||||
total_tokens=aggregated["totals"].total_tokens,
|
||||
|
|
|
|||
|
|
@ -12,11 +12,9 @@ from datetime import date, datetime, time, timedelta, timezone
|
|||
from typing import Any
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import PTU_ROLLUP_JOB_ID, PTU_SENTINEL_API_KEY
|
||||
from litellm.repositories.ptu_reservation_repository import PTUReservationRepository
|
||||
|
||||
PTU_SENTINEL_API_KEY = "__ptu_reservation__"
|
||||
PTU_ROLLUP_JOB_ID = "ptu_reservation_rollup_job"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RollupResult:
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ class GroupByDimension(str, Enum):
|
|||
|
||||
class SpendMetrics(BaseModel):
|
||||
spend: float = Field(default=0.0)
|
||||
flat_cost: 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)
|
||||
|
|
@ -71,6 +72,7 @@ class DailySpendData(BaseModel):
|
|||
|
||||
class DailySpendMetadata(BaseModel):
|
||||
total_spend: float = Field(default=0.0)
|
||||
total_flat_cost: 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)
|
||||
|
|
|
|||
|
|
@ -936,3 +936,300 @@ def test_update_metrics_handles_none_values():
|
|||
assert metrics.failed_requests == 0
|
||||
assert metrics.cache_read_input_tokens == 0
|
||||
assert metrics.cache_creation_input_tokens == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Stage 3: ptu_flat_cost surfacing on team daily activity read path
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
from litellm.constants import PTU_SENTINEL_API_KEY # noqa: E402
|
||||
from litellm.proxy.management_endpoints.common_daily_activity import ( # noqa: E402
|
||||
update_breakdown_metrics,
|
||||
)
|
||||
from litellm.types.proxy.management_endpoints.common_daily_activity import ( # noqa: E402
|
||||
BreakdownMetrics,
|
||||
)
|
||||
|
||||
|
||||
def _team_row(
|
||||
*,
|
||||
date="2026-08-01",
|
||||
api_key="sk-real",
|
||||
model="gpt-4",
|
||||
model_group=None,
|
||||
custom_llm_provider="azure",
|
||||
endpoint=None,
|
||||
mcp_namespaced_tool_name=None,
|
||||
spend=0.0,
|
||||
ptu_flat_cost=0.0,
|
||||
prompt_tokens=0,
|
||||
completion_tokens=0,
|
||||
cache_read_input_tokens=0,
|
||||
cache_creation_input_tokens=0,
|
||||
api_requests=0,
|
||||
successful_requests=0,
|
||||
failed_requests=0,
|
||||
team_id="team_x",
|
||||
):
|
||||
return SimpleNamespace(
|
||||
date=date,
|
||||
api_key=api_key,
|
||||
model=model,
|
||||
model_group=model_group,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
endpoint=endpoint,
|
||||
mcp_namespaced_tool_name=mcp_namespaced_tool_name,
|
||||
spend=spend,
|
||||
ptu_flat_cost=ptu_flat_cost,
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
cache_read_input_tokens=cache_read_input_tokens,
|
||||
cache_creation_input_tokens=cache_creation_input_tokens,
|
||||
api_requests=api_requests,
|
||||
successful_requests=successful_requests,
|
||||
failed_requests=failed_requests,
|
||||
team_id=team_id,
|
||||
)
|
||||
|
||||
|
||||
def test_spend_metrics_flat_cost_defaults_zero():
|
||||
assert SpendMetrics().flat_cost == 0.0
|
||||
|
||||
|
||||
def test_update_metrics_adds_ptu_flat_cost():
|
||||
metrics = SpendMetrics()
|
||||
update_metrics(metrics, _team_row(ptu_flat_cost=6.45))
|
||||
update_metrics(metrics, _team_row(ptu_flat_cost=3.55))
|
||||
assert metrics.flat_cost == pytest.approx(10.0)
|
||||
|
||||
|
||||
def test_update_metrics_ignores_missing_ptu_flat_cost_attr():
|
||||
"""User/org/tag daily rows don't carry ptu_flat_cost — must not raise."""
|
||||
row_without_field = SimpleNamespace(
|
||||
spend=1.0,
|
||||
prompt_tokens=0,
|
||||
completion_tokens=0,
|
||||
cache_read_input_tokens=0,
|
||||
cache_creation_input_tokens=0,
|
||||
api_requests=0,
|
||||
successful_requests=0,
|
||||
failed_requests=0,
|
||||
)
|
||||
metrics = SpendMetrics()
|
||||
update_metrics(metrics, row_without_field)
|
||||
assert metrics.spend == pytest.approx(1.0)
|
||||
assert metrics.flat_cost == 0.0
|
||||
|
||||
|
||||
def test_update_metrics_treats_none_ptu_flat_cost_as_zero():
|
||||
row = _team_row(ptu_flat_cost=None)
|
||||
metrics = SpendMetrics()
|
||||
update_metrics(metrics, row)
|
||||
assert metrics.flat_cost == 0.0
|
||||
|
||||
|
||||
def test_record_to_spend_metrics_reads_ptu_flat_cost():
|
||||
row = _team_row(ptu_flat_cost=42.0)
|
||||
m = _record_to_spend_metrics(row)
|
||||
assert m.flat_cost == pytest.approx(42.0)
|
||||
|
||||
|
||||
def test_record_to_spend_metrics_defaults_ptu_flat_cost_when_absent():
|
||||
row_without = SimpleNamespace(
|
||||
spend=1.0,
|
||||
prompt_tokens=0,
|
||||
completion_tokens=0,
|
||||
cache_read_input_tokens=0,
|
||||
cache_creation_input_tokens=0,
|
||||
api_requests=0,
|
||||
successful_requests=0,
|
||||
failed_requests=0,
|
||||
)
|
||||
m = _record_to_spend_metrics(row_without)
|
||||
assert m.flat_cost == 0.0
|
||||
|
||||
|
||||
def test_sentinel_row_skipped_from_top_level_api_keys_breakdown():
|
||||
breakdown = BreakdownMetrics()
|
||||
sentinel_row = _team_row(
|
||||
api_key=PTU_SENTINEL_API_KEY,
|
||||
ptu_flat_cost=6.45,
|
||||
custom_llm_provider=None,
|
||||
)
|
||||
update_breakdown_metrics(breakdown, sentinel_row, {}, {}, {})
|
||||
assert PTU_SENTINEL_API_KEY not in breakdown.api_keys
|
||||
assert breakdown.api_keys == {}
|
||||
|
||||
|
||||
def test_sentinel_row_contributes_to_model_breakdown_flat_cost():
|
||||
breakdown = BreakdownMetrics()
|
||||
sentinel_row = _team_row(
|
||||
api_key=PTU_SENTINEL_API_KEY,
|
||||
model="gpt-4",
|
||||
ptu_flat_cost=6.45,
|
||||
)
|
||||
update_breakdown_metrics(breakdown, sentinel_row, {}, {}, {})
|
||||
assert "gpt-4" in breakdown.models
|
||||
assert breakdown.models["gpt-4"].metrics.flat_cost == pytest.approx(6.45)
|
||||
# And it must not leak into api_key_breakdown under the model
|
||||
assert PTU_SENTINEL_API_KEY not in breakdown.models["gpt-4"].api_key_breakdown
|
||||
|
||||
|
||||
def test_real_key_and_sentinel_row_share_a_model_bucket():
|
||||
breakdown = BreakdownMetrics()
|
||||
real_row = _team_row(api_key="sk-real", model="gpt-4", spend=1.5)
|
||||
sentinel_row = _team_row(api_key=PTU_SENTINEL_API_KEY, model="gpt-4", ptu_flat_cost=6.45)
|
||||
update_breakdown_metrics(breakdown, real_row, {}, {}, {})
|
||||
update_breakdown_metrics(breakdown, sentinel_row, {}, {}, {})
|
||||
|
||||
model_bucket = breakdown.models["gpt-4"]
|
||||
assert model_bucket.metrics.spend == pytest.approx(1.5)
|
||||
assert model_bucket.metrics.flat_cost == pytest.approx(6.45)
|
||||
|
||||
# api_keys only surfaces the real key
|
||||
assert list(breakdown.api_keys.keys()) == ["sk-real"]
|
||||
# and the api_key_breakdown under the model also only holds the real key
|
||||
assert list(model_bucket.api_key_breakdown.keys()) == ["sk-real"]
|
||||
|
||||
|
||||
def test_sentinel_row_contributes_to_provider_and_endpoint_breakdowns():
|
||||
breakdown = BreakdownMetrics()
|
||||
sentinel_row = _team_row(
|
||||
api_key=PTU_SENTINEL_API_KEY,
|
||||
model="gpt-4",
|
||||
custom_llm_provider="azure",
|
||||
endpoint="/v1/chat/completions",
|
||||
ptu_flat_cost=6.45,
|
||||
)
|
||||
update_breakdown_metrics(breakdown, sentinel_row, {}, {}, {})
|
||||
assert breakdown.providers["azure"].metrics.flat_cost == pytest.approx(6.45)
|
||||
assert PTU_SENTINEL_API_KEY not in breakdown.providers["azure"].api_key_breakdown
|
||||
assert breakdown.endpoints["/v1/chat/completions"].metrics.flat_cost == pytest.approx(6.45)
|
||||
assert PTU_SENTINEL_API_KEY not in breakdown.endpoints["/v1/chat/completions"].api_key_breakdown
|
||||
|
||||
|
||||
def test_sentinel_row_contributes_to_entity_breakdown():
|
||||
breakdown = BreakdownMetrics()
|
||||
sentinel_row = _team_row(
|
||||
api_key=PTU_SENTINEL_API_KEY,
|
||||
team_id="team_x",
|
||||
ptu_flat_cost=6.45,
|
||||
)
|
||||
update_breakdown_metrics(
|
||||
breakdown,
|
||||
sentinel_row,
|
||||
{},
|
||||
{},
|
||||
{},
|
||||
entity_id_field="team_id",
|
||||
)
|
||||
assert breakdown.entities["team_x"].metrics.flat_cost == pytest.approx(6.45)
|
||||
assert PTU_SENTINEL_API_KEY not in breakdown.entities["team_x"].api_key_breakdown
|
||||
|
||||
|
||||
def test_get_api_key_metadata_excludes_sentinel_from_lookup():
|
||||
"""The sentinel api_key is not a real hashed token; skip the Prisma lookup for it."""
|
||||
from litellm.proxy.management_endpoints.common_daily_activity import (
|
||||
_aggregate_spend_records,
|
||||
)
|
||||
import asyncio
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_vt_table = MagicMock()
|
||||
mock_vt_table.find_many = AsyncMock(return_value=[])
|
||||
mock_deleted_table = MagicMock()
|
||||
mock_deleted_table.find_many = AsyncMock(return_value=[])
|
||||
mock_prisma.db = MagicMock()
|
||||
mock_prisma.db.litellm_verificationtoken = mock_vt_table
|
||||
mock_prisma.db.litellm_deletedverificationtoken = mock_deleted_table
|
||||
|
||||
real_row = _team_row(api_key="sk-real")
|
||||
sentinel_row = _team_row(api_key=PTU_SENTINEL_API_KEY, ptu_flat_cost=6.45)
|
||||
|
||||
asyncio.run(
|
||||
_aggregate_spend_records(
|
||||
prisma_client=mock_prisma,
|
||||
records=[real_row, sentinel_row],
|
||||
entity_id_field="team_id",
|
||||
entity_metadata_field=None,
|
||||
)
|
||||
)
|
||||
|
||||
mock_vt_table.find_many.assert_called_once()
|
||||
where = mock_vt_table.find_many.await_args.kwargs["where"]
|
||||
queried = set(where["token"]["in"])
|
||||
assert "sk-real" in queried
|
||||
assert PTU_SENTINEL_API_KEY not in queried
|
||||
|
||||
|
||||
def test_build_aggregated_sql_query_selects_ptu_flat_cost_only_for_team_table():
|
||||
sql_team, _ = _build_aggregated_sql_query(
|
||||
table_name="litellm_dailyteamspend",
|
||||
entity_id_field="team_id",
|
||||
entity_id=None,
|
||||
start_date="2026-08-01",
|
||||
end_date="2026-08-31",
|
||||
model=None,
|
||||
api_key=None,
|
||||
)
|
||||
assert "SUM(ptu_flat_cost)::float AS ptu_flat_cost" in sql_team
|
||||
|
||||
sql_user, _ = _build_aggregated_sql_query(
|
||||
table_name="litellm_dailyuserspend",
|
||||
entity_id_field="user_id",
|
||||
entity_id=None,
|
||||
start_date="2026-08-01",
|
||||
end_date="2026-08-31",
|
||||
model=None,
|
||||
api_key=None,
|
||||
)
|
||||
assert "SUM(ptu_flat_cost)" not in sql_user
|
||||
assert "0::float AS ptu_flat_cost" in sql_user
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_daily_activity_returns_total_flat_cost_for_team():
|
||||
"""End-to-end: seed team daily rows (one real + one sentinel), assert response shape."""
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db = MagicMock()
|
||||
|
||||
mock_table = MagicMock()
|
||||
mock_table.count = AsyncMock(return_value=2)
|
||||
mock_table.find_many = AsyncMock(
|
||||
return_value=[
|
||||
_team_row(api_key="sk-real", spend=1.5, ptu_flat_cost=0.0),
|
||||
_team_row(api_key=PTU_SENTINEL_API_KEY, spend=0.0, ptu_flat_cost=6.45),
|
||||
]
|
||||
)
|
||||
mock_prisma.db.litellm_dailyteamspend = mock_table
|
||||
|
||||
mock_vt_table = MagicMock()
|
||||
mock_vt_table.find_many = AsyncMock(return_value=[])
|
||||
mock_prisma.db.litellm_verificationtoken = mock_vt_table
|
||||
mock_deleted_table = MagicMock()
|
||||
mock_deleted_table.find_many = AsyncMock(return_value=[])
|
||||
mock_prisma.db.litellm_deletedverificationtoken = mock_deleted_table
|
||||
|
||||
result = await get_daily_activity(
|
||||
prisma_client=mock_prisma,
|
||||
table_name="litellm_dailyteamspend",
|
||||
entity_id_field="team_id",
|
||||
entity_id=None,
|
||||
entity_metadata_field=None,
|
||||
start_date="2026-08-01",
|
||||
end_date="2026-08-01",
|
||||
model=None,
|
||||
api_key=None,
|
||||
page=1,
|
||||
page_size=10,
|
||||
)
|
||||
|
||||
assert result.metadata.total_spend == pytest.approx(1.5)
|
||||
assert result.metadata.total_flat_cost == pytest.approx(6.45)
|
||||
assert len(result.results) == 1
|
||||
day = result.results[0]
|
||||
assert day.metrics.spend == pytest.approx(1.5)
|
||||
assert day.metrics.flat_cost == pytest.approx(6.45)
|
||||
# api_keys breakdown must not contain the sentinel
|
||||
assert PTU_SENTINEL_API_KEY not in day.breakdown.api_keys
|
||||
assert "sk-real" in day.breakdown.api_keys
|
||||
|
|
|
|||
10
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
10
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -23332,6 +23332,11 @@ export interface components {
|
|||
* @default 0
|
||||
*/
|
||||
total_failed_requests: number;
|
||||
/**
|
||||
* Total Flat Cost
|
||||
* @default 0
|
||||
*/
|
||||
total_flat_cost: number;
|
||||
/**
|
||||
* Total Pages
|
||||
* @default 1
|
||||
|
|
@ -31022,6 +31027,11 @@ export interface components {
|
|||
* @default 0
|
||||
*/
|
||||
failed_requests: number;
|
||||
/**
|
||||
* Flat Cost
|
||||
* @default 0
|
||||
*/
|
||||
flat_cost: number;
|
||||
/**
|
||||
* Prompt Tokens
|
||||
* @default 0
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue