From 0e1a3babf047f9142a1423022cea45a5d960cebc Mon Sep 17 00:00:00 2001
From: Yassin Kortam <yassin@berri.ai>
Date: Fri, 26 Jun 2026 19:38:58 +0300
Subject: [PATCH] perf(cost-calc): precompute service-tier cost-key suffixes
 (#31431)

_get_token_base_cost rebuilt f"_{st.value}" for every ServiceTier while
scanning every model_info key on each request, and _get_cost_per_unit rebuilt
the same f-strings in its fallback loop. The suffixes are constant, so compute
them once at module level (matching the existing _IMAGE_RESPONSE_CALL_TYPES /
_VALID_DATA_RESIDENCIES pattern) and use str.endswith(tuple) for the threshold
check. Behavior is identical; ~3.5 us/call to ~2.2 us/call on the threshold scan.
---
 .../litellm_core_utils/llm_cost_calc/utils.py | 12 +++--
 .../llm_cost_calc/test_llm_cost_calc_utils.py | 48 +++++++++++++++++++
 2 files changed, 56 insertions(+), 4 deletions(-)

diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py
index d2c97a5d6e7..e407dd70bf0 100644
--- a/litellm/litellm_core_utils/llm_cost_calc/utils.py
+++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py
@@ -33,6 +33,11 @@ _IMAGE_RESPONSE_CALL_TYPES = frozenset(
 # Pre-resolved DataResidency enum values for fast membership checks
 _VALID_DATA_RESIDENCIES = frozenset(r.value for r in DataResidency)
 
+# Pre-resolved service-tier cost-key suffixes (e.g. "_priority"). Used per
+# request in the cost-calc path, so the f-strings are built once here instead
+# of being rebuilt for every model_info key on every call.
+_SERVICE_TIER_SUFFIXES: tuple[str, ...] = tuple(f"_{st.value}" for st in ServiceTier)
+
 
 def _get_token_detail_value(details: object, key: str) -> Optional[int]:
     if isinstance(details, dict):
@@ -252,7 +257,7 @@ def _get_token_base_cost(
         k
         for k in model_info
         if k.startswith("input_cost_per_token_above_")
-        and not any(k.endswith(f"_{st.value}") for st in ServiceTier)
+        and not k.endswith(_SERVICE_TIER_SUFFIXES)
     ]
     if not threshold_keys:
         return (
@@ -418,9 +423,8 @@ def _get_cost_per_unit(
 
     # If the service tier key doesn't exist or is None, try to fall back to the standard key
     if cost_per_unit is None:
-        # Check if any service tier suffix exists in the cost key using ServiceTier enum
-        for service_tier in ServiceTier:
-            suffix = f"_{service_tier.value}"
+        # Check if any service tier suffix exists in the cost key
+        for suffix in _SERVICE_TIER_SUFFIXES:
             if suffix in cost_key:
                 # Extract the base key by removing the matched suffix
                 base_key = cost_key.replace(suffix, "")
diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
index 72fee9bfb40..b75b6d78090 100644
--- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
+++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
@@ -1720,3 +1720,51 @@ def test_priority_service_tier_above_threshold_falls_back_to_standard_for_cache_
     expected_completion = 1_000 * 2.25e-5
     assert prompt_cost == pytest.approx(expected_prompt, rel=1e-9)
     assert completion_cost == pytest.approx(expected_completion, rel=1e-9)
+
+
+def test_service_tier_suffixes_constant_in_sync_with_enum():
+    from litellm.litellm_core_utils.llm_cost_calc.utils import _SERVICE_TIER_SUFFIXES
+    from litellm.types.utils import ServiceTier
+
+    assert _SERVICE_TIER_SUFFIXES == tuple(f"_{st.value}" for st in ServiceTier)
+
+
+def test_get_cost_per_unit_falls_back_from_service_tier_key_to_base():
+    from litellm.litellm_core_utils.llm_cost_calc.utils import _get_cost_per_unit
+
+    model_info = {"input_cost_per_token": 2e-6}
+    # service-tier key is absent -> falls back to the base key
+    assert _get_cost_per_unit(model_info, "input_cost_per_token_priority") == 2e-6
+    # service-tier key present -> used directly, no fallback
+    model_info_direct = {
+        "input_cost_per_token_priority": 5e-6,
+        "input_cost_per_token": 2e-6,
+    }
+    assert (
+        _get_cost_per_unit(model_info_direct, "input_cost_per_token_priority") == 5e-6
+    )
+
+
+def test_threshold_keys_exclude_service_tier_variants():
+    from typing import cast
+
+    from litellm.litellm_core_utils.llm_cost_calc.utils import _get_token_base_cost
+    from litellm.types.utils import ModelInfo, Usage
+
+    # The service-tier-suffixed above-threshold key must be excluded from
+    # threshold detection. The _priority variant has a higher threshold (300k),
+    # so if it were not excluded it would sort first and drive a 9e-6 rate for
+    # this non-tier request. With the exclusion only the standard 200k key
+    # applies, giving 3e-6.
+    model_info = cast(
+        ModelInfo,
+        {
+            "input_cost_per_token": 1e-6,
+            "input_cost_per_token_above_200k_tokens": 3e-6,
+            "input_cost_per_token_above_300k_tokens_priority": 9e-6,
+            "output_cost_per_token": 2e-6,
+        },
+    )
+    usage = Usage(prompt_tokens=350_000, completion_tokens=1_000, total_tokens=351_000)
+    prompt_base, *_ = _get_token_base_cost(model_info=model_info, usage=usage)
+    assert prompt_base == 3e-6