From 9eb6fbc5727f89217fc6e8e4eb1477c66d2414b3 Mon Sep 17 00:00:00 2001 From: kerry Date: Thu, 17 Sep 2026 22:32:39 +0000 Subject: [PATCH] test: read cost-map keys the implementation resolves and isolate the tariff test's model_cost copy Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../parallel_ai/test_parallel_ai_search.py | 2 +- .../common_utils/test_prompt_cache_pricing.py | 21 ++++++++----------- tests/test_litellm/proxy/test_proxy_utils.py | 5 +++-- tests/test_litellm/test_cost_calculator.py | 2 +- 4 files changed, 14 insertions(+), 16 deletions(-) diff --git a/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search.py b/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search.py index 51fe5cea4d3..f08689464a7 100644 --- a/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search.py +++ b/tests/test_litellm/llms/parallel_ai/test_parallel_ai_search.py @@ -466,7 +466,7 @@ class TestParallelAISearch: ) rate: Final = litellm.model_cost[ - "parallel_ai/search-fast" if mode in ("fast", "turbo") else "parallel_ai/search" + {"fast": "parallel_ai/search-fast", "turbo": "parallel_ai/search-turbo"}.get(mode, "parallel_ai/search") ]["input_cost_per_query"] request_count: Final = ( sum(item["count"] for item in usage if item["name"] == "sku_search") if usage is not None else 1 diff --git a/tests/test_litellm/proxy/common_utils/test_prompt_cache_pricing.py b/tests/test_litellm/proxy/common_utils/test_prompt_cache_pricing.py index b6bffaf79af..52d388fbad0 100644 --- a/tests/test_litellm/proxy/common_utils/test_prompt_cache_pricing.py +++ b/tests/test_litellm/proxy/common_utils/test_prompt_cache_pricing.py @@ -1,4 +1,5 @@ from collections.abc import Mapping +from copy import deepcopy from typing import Final import pytest @@ -8,27 +9,23 @@ from litellm.proxy.common_utils.prompt_cache_pricing import price_cache_tokens from litellm.types.management_endpoints.prompt_cache_prediction import CacheTokenBuckets -def _tiered_rate(entry: Mapping[str, float], field: str, total: int) -> float: - above_field: Final = f"{field}_above_200k_tokens" - if total > 200_000 and above_field in entry: - return entry[above_field] - return entry[field] +def _tiered_rate(entry: Mapping[str, float | None], field: str, total: int) -> float: + above_rate: Final = entry.get(f"{field}_above_200k_tokens") if total > 200_000 else None + rate: Final = above_rate if above_rate is not None else entry[field] + assert rate is not None + return rate def _expected_cache_cost(model: str, tokens: CacheTokenBuckets) -> float: key: Final = litellm.get_model_info(model=model, custom_llm_provider="anthropic")["key"] entry: Final = litellm.model_cost[key] total: Final = tokens.total_tokens - one_hour_field: Final = ( - "cache_creation_input_token_cost_above_1hr_above_200k_tokens" - if total > 200_000 and "cache_creation_input_token_cost_above_1hr_above_200k_tokens" in entry - else "cache_creation_input_token_cost_above_1hr" - ) return ( tokens.uncached_input_tokens * _tiered_rate(entry, "input_cost_per_token", total) + tokens.cache_read_input_tokens * _tiered_rate(entry, "cache_read_input_token_cost", total) + tokens.cache_creation_5m_input_tokens * _tiered_rate(entry, "cache_creation_input_token_cost", total) - + tokens.cache_creation_1h_input_tokens * entry[one_hour_field] + + tokens.cache_creation_1h_input_tokens + * _tiered_rate(entry, "cache_creation_input_token_cost_above_1hr", total) ) @@ -58,7 +55,7 @@ def test_long_context_tier_starts_above_threshold(total: int) -> None: def test_deployment_tariff_wins_without_proxy_discounts_or_margins(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(litellm, "model_cost", litellm.model_cost.copy()) + monkeypatch.setattr(litellm, "model_cost", deepcopy(litellm.model_cost)) litellm.Router( model_list=[ { diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index 9f3d03ce515..cd8b5ba8844 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -2239,8 +2239,9 @@ def test_create_model_info_response_falls_back_to_alias_for_opaque_deployment_na litellm.model_cost.clear() litellm.model_cost.update(saved_model_cost) - assert response["max_input_tokens"] == 128000 - assert response["max_output_tokens"] == 16384 + entry: Final = litellm.model_cost["gpt-4o"] + assert response["max_input_tokens"] == entry["max_input_tokens"] + assert response["max_output_tokens"] == entry["max_output_tokens"] def test_create_model_info_response_resolves_mode_through_deployment_model(): diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 80f2a9903a5..03e4ef3b2c3 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -1346,7 +1346,7 @@ def test_gemini_25_implicit_caching_cost(): model="gemini/gemini-2.5-flash", ) - model_info: Final = litellm.model_cost["gemini-2.5-flash"] + model_info: Final = litellm.model_cost["gemini/gemini-2.5-flash"] expected_cost = ( 14316 * model_info["cache_read_input_token_cost"] + (15033 - 14316) * model_info["input_cost_per_token"]