Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions tests/e2e/coverage_registry/quota_management.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -44,3 +44,10 @@
- {id: quota_management.spend_tracking.failure.writes_failure_row, module: quota_management, tier: P1, behavior: spend_tracking, variant: failure, assertions: [writes_failure_row], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_log_error_logger.py", rationale: "A failed call writes a failure-status spend row"}
- {id: quota_management.spend_tracking.spend_calculate.returns_cost, module: quota_management, tier: P2, behavior: spend_tracking, variant: spend_calculate, assertions: [returns_cost], exercised_on: [spend_calculate], source: "proxy/spend_tracking/spend_management_endpoints.py", rationale: "/spend/calculate prices a hypothetical request at nonzero cost"}
- {id: quota_management.spend_tracking.pagination.keeps_total, module: quota_management, tier: P2, behavior: spend_tracking, variant: pagination, assertions: [keeps_total], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_management_endpoints.py", rationale: "Spend-logs v2 pagination caps page size without losing the total"}
- {id: quota_management.spend_tracking.cache_write.bills_cache_creation_rate, module: quota_management, tier: P1, behavior: spend_tracking, variant: cache_write, assertions: [bills_cache_creation_rate], exercised_on: [chat_completions], source: "litellm_core_utils/llm_cost_calc/utils.py", rationale: "OpenAI cache-write tokens land on the spend row as cache-creation tokens billed at the cache-creation rate, not silently at the input rate (#34046)"}
- {id: quota_management.spend_tracking.cost_breakdown.reports_component_costs, module: quota_management, tier: P1, behavior: spend_tracking, variant: cost_breakdown, assertions: [reports_component_costs], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "The spend row's metadata.cost_breakdown itemizes cache-read, cache-creation, output, and reasoning costs at the deployment's own rates and they sum to the row's spend (#31686)"}
- {id: quota_management.spend_tracking.stream_cache_read.bills_cache_read_rate, module: quota_management, tier: P1, behavior: spend_tracking, variant: stream_cache_read, assertions: [bills_cache_read_rate], exercised_on: [chat_completions], source: "litellm_core_utils/streaming_chunk_builder_utils.py", rationale: "A streamed call's reassembled usage keeps the cached-token detail so cache reads bill at the cache-read discount, not full input price (#34812)"}
- {id: quota_management.spend_tracking.messages_bridge.keeps_cache_tokens, module: quota_management, tier: P1, behavior: spend_tracking, variant: messages_bridge, assertions: [keeps_cache_tokens], exercised_on: [messages], source: "llms/anthropic/experimental_pass_through/responses_adapters/handler.py", rationale: "A /v1/messages request served by a Responses-only OpenAI model keeps its cache-read tokens and their discounted billing across the bridge (#34957)"}
- {id: quota_management.spend_tracking.service_tier.bills_tier_rates, module: quota_management, tier: P1, behavior: spend_tracking, variant: service_tier, assertions: [bills_tier_rates], exercised_on: [chat_completions], source: "cost_calculator.py", rationale: "A priority service_tier call bills input, output, and reasoning at the deployment's *_priority rates and records the tier on the row (#35923, #35925)"}
- {id: quota_management.spend_tracking.cost_headers.additive_components, module: quota_management, tier: P1, behavior: spend_tracking, variant: cost_headers, assertions: [additive_components], exercised_on: [chat_completions], source: "proxy/common_request_processing.py", rationale: "The x-litellm-response-cost-* component headers sum to the total, input covers only fresh tokens, and reasoning stays a subset of output (#36965)"}
- {id: quota_management.spend_tracking.passthrough_stream.injects_usage_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: passthrough_stream, assertions: [injects_usage_cost], exercised_on: [openai_passthrough], source: "proxy/pass_through_endpoints/streaming_handler.py", rationale: "With include_cost_in_streaming_usage on, the /openai passthrough's final streaming usage frame carries the proxy-computed cost (#36503). Uncovered: the flag is only settable in litellm_settings, and the shared e2e stack does not turn it on yet"}
11 changes: 9 additions & 2 deletions tests/e2e/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -452,6 +452,7 @@ class AnthropicMessagesResponse(BaseModel):
model: str | None = None
content: list[AnthropicContentBlock] | None = None
choices: list[ChatChoice] | None = None
usage: Usage | None = None


class CountTokensResponse(BaseModel):
Expand Down Expand Up @@ -716,8 +717,10 @@ class FineTuningJobsResponse(BaseModel):
class LiteLLMParamsBody(BaseModel):
"""POST /model/new litellm_params: `model` is the only required field; `api_key`
et al may be an `os.environ/FOO` reference the proxy resolves at call time.
`input_cost_per_token`/`output_cost_per_token` register a per-deployment custom
pricing override; left None (and dropped from the body) the deployment keeps the
The `*_cost_per_token` / `*_token_cost` fields register a per-deployment custom
pricing override (the cache and `_priority` rates only apply when both base
rates are set, which is what makes the proxy register the deployment's full
pricing entry); left None (and dropped from the body) the deployment keeps the
backend's canonical rate."""

model: str
Expand All @@ -744,6 +747,10 @@ class LiteLLMParamsBody(BaseModel):
aws_external_id: str | None = None
input_cost_per_token: float | None = None
output_cost_per_token: float | None = None
cache_read_input_token_cost: float | None = None
cache_creation_input_token_cost: float | None = None
input_cost_per_token_priority: float | None = None
output_cost_per_token_priority: float | None = None
extra_headers: dict[str, str] | None = None
use_in_pass_through: bool | None = None
complexity_router_config: dict[str, object] | None = None
Expand Down
204 changes: 204 additions & 0 deletions tests/e2e/quota_management/spend_tracking/cost_rows.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,204 @@
"""Cost-accounting helpers for the spend-tracking suite: the /spend/logs row shape
that carries the per-component cost breakdown, a poll that waits for it, and the
builders the cache-pricing tests share.

The shared SpendLogRow deliberately stays thin (most tests only read totals), so
the component-cost tests model the metadata they assert on here instead:
`metadata.cost_breakdown` (input/output/cache-read/cache-creation/reasoning costs
plus the service-tier pricing basis) and `metadata.additional_usage_values` (the
cache token counts the biller derived from the provider's usage).

Determinism strategy: every test registers its own deployment with explicit custom
rates for each component it asserts on (`register_priced_model`), so expected cost
is exactly tokens-on-the-row times configured rate, immune to provider price
changes. The rates are chosen ~100x above canonical and distinct from one another,
so a component billed at the wrong rate can never accidentally match.

OpenAI prompt caching is implicit and keyed on the exact token prefix, with a
1024-token minimum. `cacheable_prefix` builds a prefix whose first word is the
run's unique marker: unique marker = the whole prefix is novel (a fresh cache
write), same marker + different question = a cache read that still misses the
proxy's own response cache. How long the prefix has to be before the provider
actually reports a read varies by model, so callers pass `words` to suit theirs.

Two facts about the recorded bill that the assertions here encode, because the
two surfaces disagree on purpose. On the spend row, `input_cost` is gross: it
already contains the cache-read and cache-creation costs, so the row's total is
input + output + tool-usage and the fresh-token cost is input minus the two cache
components. In the response headers, `x-litellm-response-cost-input` is net of
cache, which is what makes the component headers sum to the total.
"""

import time
from collections.abc import Callable

from pydantic import BaseModel, RootModel

from e2e_config import unique_marker
from e2e_http import Success
from lifecycle import ResourceManager
from models import LiteLLMParamsBody, SpendLogsParams
from proxy_client import ProxyClient


class CostBreakdownRow(BaseModel):
input_cost: float | None = None
output_cost: float | None = None
cache_read_cost: float | None = None
cache_creation_cost: float | None = None
reasoning_cost: float | None = None
tool_usage_cost: float | None = None
total_cost: float | None = None
service_tier: str | None = None


class AdditionalUsageValues(BaseModel):
cache_read_input_tokens: int | None = None
cache_creation_input_tokens: int | None = None


class CostRowMetadata(BaseModel):
cost_breakdown: CostBreakdownRow | None = None
additional_usage_values: AdditionalUsageValues | None = None


class CostRow(BaseModel):
request_id: str | None = None
spend: float | None = None
prompt_tokens: int | None = None
completion_tokens: int | None = None
metadata: CostRowMetadata | None = None

@property
def breakdown(self) -> CostBreakdownRow:
assert self.metadata and self.metadata.cost_breakdown, (
f"spend row {self.request_id} landed without a cost breakdown"
)
return self.metadata.cost_breakdown

@property
def cache_read_tokens(self) -> int:
if self.metadata and self.metadata.additional_usage_values:
return self.metadata.additional_usage_values.cache_read_input_tokens or 0
return 0

@property
def cache_creation_tokens(self) -> int:
if self.metadata and self.metadata.additional_usage_values:
return self.metadata.additional_usage_values.cache_creation_input_tokens or 0
return 0


class CostRows(RootModel[list[CostRow]]):
pass


def approx_equal(actual: float, expected: float) -> bool:
"""Within 1% or 1e-9 absolute - spend math, not exact float identity."""
return abs(actual - expected) <= max(1e-9, abs(expected) * 1e-2)


def assert_total_is_sum_of_components(row: CostRow) -> None:
"""The row's total is input + output + tool usage. The cache components are
already inside the gross input cost, so adding them again would double-bill."""
breakdown = row.breakdown
components = sum(
cost or 0.0
for cost in (breakdown.input_cost, breakdown.output_cost, breakdown.tool_usage_cost)
)
assert breakdown.total_cost is not None and approx_equal(breakdown.total_cost, components), (
f"total_cost {breakdown.total_cost} != input + output + tool usage ({components}): {breakdown}"
)
assert row.spend is not None and approx_equal(row.spend, breakdown.total_cost), (
f"row spend {row.spend} != breakdown total {breakdown.total_cost}"
)


def assert_fresh_tokens_billed_at(row: CostRow, input_rate: float) -> None:
"""Strip the cache components out of the gross input cost and what is left must
be the freshly-read tokens at the deployment's input rate."""
breakdown = row.breakdown
fresh_tokens = (row.prompt_tokens or 0) - row.cache_read_tokens - row.cache_creation_tokens
fresh_cost = (
(breakdown.input_cost or 0.0)
- (breakdown.cache_read_cost or 0.0)
- (breakdown.cache_creation_cost or 0.0)
)
assert breakdown.input_cost is not None and approx_equal(fresh_cost, fresh_tokens * input_rate), (
f"input_cost {breakdown.input_cost} less cache read {breakdown.cache_read_cost} and "
f"cache creation {breakdown.cache_creation_cost} leaves {fresh_cost}, not "
f"{fresh_tokens} fresh tokens * {input_rate} (prompt {row.prompt_tokens}, "
f"cache read {row.cache_read_tokens}, cache creation {row.cache_creation_tokens}); "
"cached tokens are being billed at the input rate"
)


def poll_cost_row(proxy: ProxyClient, request_id: str) -> CostRow | None:
"""Poll /spend/logs for the call's row until it lands with a cost breakdown
(rows flush ~60s behind the call via proxy_batch_write_at); None on timeout."""
deadline = time.monotonic() + proxy.poll_timeout
while time.monotonic() < deadline:
result = proxy.transport.get(
"/spend/logs",
headers=proxy.transport.master,
params=SpendLogsParams(request_id=request_id),
response_type=CostRows,
)
match result:
case Success(data=data):
rows = data.root
case _:
rows = []
for row in rows:
if row.metadata and row.metadata.cost_breakdown:
return row
time.sleep(proxy.poll_interval)
return None


def poll_cost_row_where(
proxy: ProxyClient, api_key: str, predicate: Callable[[CostRow], bool]
) -> CostRow | None:
"""Poll the key's own /spend/logs until one of its rows carries a cost breakdown
the predicate accepts; None on timeout. For calls whose response id is not the
id the bill is filed under, which is how a user finds the row in the UI anyway."""
deadline = time.monotonic() + proxy.poll_timeout
while time.monotonic() < deadline:
result = proxy.transport.get(
"/spend/logs",
headers=proxy.transport.master,
params=SpendLogsParams(api_key=api_key),
response_type=CostRows,
)
match result:
case Success(data=data):
rows = data.root
case _:
rows = []
for row in rows:
if row.metadata and row.metadata.cost_breakdown and predicate(row):
return row
time.sleep(proxy.poll_interval)
return None


def register_priced_model(
proxy: ProxyClient,
resources: ResourceManager,
name_prefix: str,
litellm_params: LiteLLMParamsBody,
) -> str:
"""Register a deployment with explicit custom rates (deleted on teardown) and
return its unique model name."""
model_name = f"{name_prefix}-{unique_marker()}"
model_id = proxy.create_model(model_name, litellm_params)
resources.defer(lambda: proxy.delete_model(model_id))
return model_name


def cacheable_prefix(marker: str, *, words: int = 1200) -> str:
"""A prompt prefix above OpenAI's 1024-token caching minimum whose identity is
fully determined by `marker` (it is the first word, and prefix caching matches
from token zero). Raise `words` for models that only report a cache read on a
substantially longer prefix."""
return " ".join(marker if i == 0 else f"token{i:04d}" for i in range(words))
Loading
Loading