From 52ca688c9d0970ab1e6c746f8d0e693ee9609992 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Fri, 31 Jul 2026 21:57:39 -0700 Subject: [PATCH] refactor(spend): let the types carry the facts the tests already assert Three review points, one theme: state that was true in the tests but not expressed in the code. `autorouter_savings_spend` was a required key on `BaseDailySpendTransaction` while every reader coalesces a missing value to zero and the aggregation is explicitly tested against rows that omit it. Rows queued by a pod on the previous release, or replayed from the Redis buffer across an upgrade, carry no such key, so the honest declaration is NotRequired. Requiring it also made every construction site outside this PR a type error for no gain. The derived baseline was cached behind a value plus a separate "have we derived yet" flag, so a baseline of None was indistinguishable from an uncomputed one until the flag was consulted. `cached_property` already stores None as a real cached value, so both attributes and the flag collapse into one declaration and the ambiguity is gone by construction rather than by discipline. `_baseline_usage` moves the cache-creation tokens into the cached count and drops the creation charge, which is the whole point of the counterfactual and was readable only to someone who had absorbed the PR description. It now says so where it happens. The pairing test that enumerates additive metrics now unwraps NotRequired, so a metric declared that way is still covered rather than silently skipped. --- litellm/proxy/_types.py | 8 +++++-- litellm/proxy/spend_tracking/savings.py | 3 +++ .../auto_router/auto_router.py | 23 ++++++++++--------- .../test_daily_spend_update_queue.py | 10 +++++++- .../router_strategy/test_auto_router.py | 2 -- 5 files changed, 30 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index b24e075dd68..8c4fc54797c 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -13,7 +13,7 @@ from pydantic import ( field_validator, model_validator, ) -from typing_extensions import Required, TypedDict +from typing_extensions import NotRequired, Required, TypedDict from litellm._uuid import uuid from litellm.constants import MCP_STDIO_ALLOWED_COMMANDS @@ -4560,7 +4560,11 @@ class BaseDailySpendTransaction(TypedDict): # cost-savings metrics (dollars, priced per request before aggregation) compression_savings_spend: float prompt_caching_savings_spend: float - autorouter_savings_spend: float + # Not required: rows queued by a pod running the previous release, or replayed from + # the Redis buffer across an upgrade, carry no such key. Every reader coalesces a + # missing value to zero, so requiring it here would describe a shape the aggregation + # is explicitly tested against. + autorouter_savings_spend: NotRequired[float] # request level metrics spend: float diff --git a/litellm/proxy/spend_tracking/savings.py b/litellm/proxy/spend_tracking/savings.py index b291ac6cf02..a611dc4301d 100644 --- a/litellm/proxy/spend_tracking/savings.py +++ b/litellm/proxy/spend_tracking/savings.py @@ -115,6 +115,9 @@ def _baseline_usage(usage: Usage) -> Usage: total_tokens=usage.total_tokens, completion_tokens_details=usage.completion_tokens_details, prompt_tokens_details=PromptTokensDetailsWrapper( + # The tokens this request paid to write are moved into the cached count and + # the creation charge is dropped: on one model that cache was already warm, + # so the baseline would have read them rather than paying to create them. cached_tokens=cache_read + cache_creation, cache_creation_tokens=0, text_tokens=max(usage.prompt_tokens - cache_read - cache_creation, 0), diff --git a/litellm/router_strategy/auto_router/auto_router.py b/litellm/router_strategy/auto_router/auto_router.py index 737cbeedb50..2fd31dca541 100644 --- a/litellm/router_strategy/auto_router/auto_router.py +++ b/litellm/router_strategy/auto_router/auto_router.py @@ -2,6 +2,7 @@ Auto-Routing Strategy that works with a Semantic Router Config """ +from functools import cached_property from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union from litellm._logging import verbose_router_logger @@ -54,8 +55,6 @@ class AutoRouter(CustomLogger): self.embedding_model: str = embedding_model self.litellm_router_instance: "Router" = litellm_router_instance self.configured_savings_baseline_model: str | None = savings_baseline_model - self._derived_savings_baseline_model: str | None = None - self._savings_baseline_derived = False @staticmethod def _canonical_model(model: str, custom_llm_provider: str | None) -> str | None: @@ -136,6 +135,14 @@ class AutoRouter(CustomLogger): return None return max(priced)[2] + @cached_property + def _derived_savings_baseline_model(self) -> str | None: + """Resolved on first use, not at construction, because the parent router's + deployments are still being assembled while this router is built. Caching + through `cached_property` keeps "derived to nothing" distinct from "not derived + yet" without a second flag to hold them apart.""" + return self._most_expensive_candidate() + @property def savings_baseline_model(self) -> str | None: """The model this router's savings are measured against. @@ -146,16 +153,10 @@ class AutoRouter(CustomLogger): honest: a fixed flagship credits savings against a model the operator would never have run, and drifts the moment the routes change. - Resolved lazily and cached, because the parent router's deployments are still - being assembled while this router is constructed. ``None`` when nothing can be - priced, which zeroes the driver rather than inventing a baseline. + ``None`` when nothing can be priced, which zeroes the driver rather than + inventing a baseline. """ - if self.configured_savings_baseline_model: - return self.configured_savings_baseline_model - if not self._savings_baseline_derived: - self._derived_savings_baseline_model = self._most_expensive_candidate() - self._savings_baseline_derived = True - return self._derived_savings_baseline_model + return self.configured_savings_baseline_model or self._derived_savings_baseline_model def _load_semantic_routing_routes(self) -> List[Route]: from semantic_router.routers import SemanticRouter diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_daily_spend_update_queue.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_daily_spend_update_queue.py index 2ac22a6d1fd..a55d4f0dcfd 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_daily_spend_update_queue.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_daily_spend_update_queue.py @@ -16,6 +16,8 @@ from litellm.proxy._types import ( Litellm_EntityType, SpendUpdateQueueItem, ) +from typing import get_args + from litellm.proxy._types import BaseDailySpendTransaction from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( DailySpendUpdateQueue, @@ -542,8 +544,14 @@ async def test_every_optional_daily_metric_aggregates(daily_spend_update_queue): paths, so the driver reads as zero on the dashboard however much it saved. """ test_key = "user1_2023-01-01_key123_claude-haiku-4-5_anthropic" + def _numeric(annotation): + # additive metrics may be declared NotRequired[float] for rows queued by a pod + # running the previous release, so unwrap before matching + args = get_args(annotation) + return (args[0] if args else annotation) in (int, float) + numeric_fields = [ - name for name, annotation in BaseDailySpendTransaction.__annotations__.items() if annotation in (int, float) + name for name, annotation in BaseDailySpendTransaction.__annotations__.items() if _numeric(annotation) ] assert "autorouter_savings_spend" in numeric_fields increments = {field: index + 1 for index, field in enumerate(numeric_fields)} diff --git a/tests/test_litellm/router_strategy/test_auto_router.py b/tests/test_litellm/router_strategy/test_auto_router.py index 52c39fff6c3..72ecee53bc5 100644 --- a/tests/test_litellm/router_strategy/test_auto_router.py +++ b/tests/test_litellm/router_strategy/test_auto_router.py @@ -412,8 +412,6 @@ class TestSavingsBaselineModel: auto_router.default_model = default_model auto_router.litellm_router_instance = parent auto_router.configured_savings_baseline_model = configured - auto_router._derived_savings_baseline_model = None - auto_router._savings_baseline_derived = False return auto_router def test_route_names_resolve_through_the_parent_router_to_pricable_models(self):