From bc593ccda3e8ef39adf07489a6e46f63cfde39e6 Mon Sep 17 00:00:00 2001 From: liming Date: Mon, 7 Sep 2026 21:41:12 +0800 Subject: [PATCH] refactor(cost): register provider cache rate resolver --- .../litellm_core_utils/llm_cost_calc/utils.py | 51 ++++++++++++++----- litellm/llms/dashscope/cost_calculator.py | 20 ++++++++ 2 files changed, 59 insertions(+), 12 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 2c4f96c02ee..eca65e542e0 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -2,7 +2,7 @@ ## Helper utilities for cost_per_token() import re -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timezone, tzinfo from types import MappingProxyType @@ -428,6 +428,36 @@ class TokenRates: return self.output_rate if self.reasoning_rate is None else self.reasoning_rate +CacheReadCostRateResolver = Callable[ + [ModelInfo, Usage, datetime | None], # mutable-ok: Callable parameter syntax + float, +] +_CACHE_READ_COST_RATE_RESOLVERS: dict[ + str, CacheReadCostRateResolver +] = {} # mutable-ok: provider adapters register one process-wide resolver at import time + + +def register_cache_read_cost_rate_resolver( + providers: Sequence[str], + resolver: CacheReadCostRateResolver, +) -> None: + """Register a provider-owned cache-read rate resolver for cost breakdowns.""" + for provider in providers: + _CACHE_READ_COST_RATE_RESOLVERS[provider] = resolver + + +def _resolve_cache_read_cost_rate( + custom_llm_provider: str | None, + model_info: ModelInfo, + usage: Usage, + current_time: datetime | None, +) -> float | None: + resolver: Final = _CACHE_READ_COST_RATE_RESOLVERS.get(custom_llm_provider or "") + if resolver is None: + return None + return resolver(model_info, usage, current_time) + + def _parse_off_peak_rate(value: object) -> float | None: if isinstance(value, bool): return None @@ -1356,17 +1386,14 @@ def get_token_type_cost_breakdown( threshold_is_inclusive=_uses_inclusive_token_thresholds(custom_llm_provider), ) - if custom_llm_provider in ("dashscope", "qwencloud", "qwen_ai_platform"): - from litellm.llms.dashscope.cost_calculator import ( - get_token_breakdown_and_rates, - ) - - _, dashscope_rates = get_token_breakdown_and_rates( - model_info=model_info, - usage=usage, - current_time=billing_time, - ) - cache_read_cost_rate = dashscope_rates.cache_read_rate + provider_cache_read_cost_rate: Final = _resolve_cache_read_cost_rate( + custom_llm_provider=custom_llm_provider, + model_info=model_info, + usage=usage, + current_time=billing_time, + ) + if provider_cache_read_cost_rate is not None: + cache_read_cost_rate = provider_cache_read_cost_rate reasoning_tokens = ( parse_completion_tokens_details(usage)["reasoning_tokens"] if usage.completion_tokens_details is not None else 0 diff --git a/litellm/llms/dashscope/cost_calculator.py b/litellm/llms/dashscope/cost_calculator.py index 35a0fdc6edf..7d75be4fe0d 100644 --- a/litellm/llms/dashscope/cost_calculator.py +++ b/litellm/llms/dashscope/cost_calculator.py @@ -21,6 +21,7 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import ( apply_off_peak_pricing, parse_completion_tokens_details, parse_prompt_tokens_details, + register_cache_read_cost_rate_resolver, ) from litellm.types.utils import ModelInfo, Usage from litellm.utils import get_model_info @@ -204,6 +205,25 @@ def get_token_breakdown_and_rates( return breakdown, rates +def _resolve_cache_read_cost_rate( + model_info: ModelInfo, + usage: Usage, + current_time: datetime | None, +) -> float: + _, rates = get_token_breakdown_and_rates( + model_info=model_info, + usage=usage, + current_time=current_time, + ) + return rates.cache_read_rate + + +register_cache_read_cost_rate_resolver( + providers=("dashscope", "qwencloud", "qwen_ai_platform"), + resolver=_resolve_cache_read_cost_rate, +) + + def cost_per_token( model: str, usage: Usage,