From eb2be3758a683ccf8d80fb174df187c9cbb5fa28 Mon Sep 17 00:00:00 2001 From: kerry Date: Thu, 17 Sep 2026 22:22:24 +0000 Subject: [PATCH] test: read cost expectations from the catalog row the code bills against Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../llms/parallel_ai/test_parallel_ai_search.py | 7 ++++--- tests/test_litellm/test_cost_calculator.py | 2 +- 2 files changed, 5 insertions(+), 4 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..03fda270b6f 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 @@ -465,9 +465,10 @@ class TestParallelAISearch: max_results=max_results, ) - rate: Final = litellm.model_cost[ - "parallel_ai/search-fast" if mode in ("fast", "turbo") else "parallel_ai/search" - ]["input_cost_per_query"] + pricing_model: Final = {"fast": "parallel_ai/search-fast", "turbo": "parallel_ai/search-turbo"}.get( + mode, "parallel_ai/search" + ) + rate: Final = litellm.model_cost[pricing_model]["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/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"]