From 5a8762b6a1e02a8062a96dc326c795facd9df5ce Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 16 Jul 2025 08:55:07 -0700 Subject: [PATCH] test: update tests --- tests/llm_translation/test_gemini.py | 56 ++++++++----------- .../llm_cost_calc/test_llm_cost_calc_utils.py | 2 +- tests/test_litellm/test_cost_calculator.py | 25 ++++----- 3 files changed, 36 insertions(+), 47 deletions(-) diff --git a/tests/llm_translation/test_gemini.py b/tests/llm_translation/test_gemini.py index ae279bbde4..46efd738b6 100644 --- a/tests/llm_translation/test_gemini.py +++ b/tests/llm_translation/test_gemini.py @@ -24,7 +24,7 @@ class TestGoogleAIStudioGemini(BaseLLMChatTest): return {"model": "gemini/gemini-2.0-flash"} def get_base_completion_call_args_with_reasoning_model(self) -> dict: - return {"model": "gemini/gemini-2.5-flash-preview-04-17"} + return {"model": "gemini/gemini-2.5-flash"} def test_tool_call_no_arguments(self, tool_call_no_arguments): """Test that tool calls with no arguments is translated correctly. Relevant issue: https://github.com/BerriAI/litellm/issues/6833""" @@ -35,7 +35,6 @@ class TestGoogleAIStudioGemini(BaseLLMChatTest): result = convert_to_gemini_tool_call_invoke(tool_call_no_arguments) print(result) - @pytest.mark.flaky(retries=3, delay=2) def test_url_context(self): from litellm.utils import supports_url_context @@ -70,7 +69,7 @@ class TestGoogleAIStudioGemini(BaseLLMChatTest): def test_gemini_context_caching_with_ttl(): """Test Gemini context caching with TTL support""" - + # Test case 1: Basic TTL functionality messages_with_ttl = [ { @@ -92,22 +91,22 @@ def test_gemini_context_caching_with_ttl(): "cache_control": {"type": "ephemeral", "ttl": "7200s"}, } ], - } + }, ] - + # Test the transformation function directly result = transform_openai_messages_to_gemini_context_caching( model="gemini-1.5-pro", messages=messages_with_ttl, - cache_key="test-ttl-cache-key" + cache_key="test-ttl-cache-key", ) - + # Verify TTL is properly included in the result assert "ttl" in result assert result["ttl"] == "3600s" # Should use the first valid TTL found assert result["model"] == "models/gemini-1.5-pro" assert result["displayName"] == "test-ttl-cache-key" - + # Test case 2: Invalid TTL should be ignored messages_invalid_ttl = [ { @@ -121,18 +120,18 @@ def test_gemini_context_caching_with_ttl(): ], } ] - + result_invalid = transform_openai_messages_to_gemini_context_caching( model="gemini-1.5-pro", messages=messages_invalid_ttl, - cache_key="test-invalid-ttl" + cache_key="test-invalid-ttl", ) - + # Verify invalid TTL is not included assert "ttl" not in result_invalid assert result_invalid["model"] == "models/gemini-1.5-pro" assert result_invalid["displayName"] == "test-invalid-ttl" - + # Test case 3: Messages without TTL should work normally messages_no_ttl = [ { @@ -146,18 +145,16 @@ def test_gemini_context_caching_with_ttl(): ], } ] - + result_no_ttl = transform_openai_messages_to_gemini_context_caching( - model="gemini-1.5-pro", - messages=messages_no_ttl, - cache_key="test-no-ttl" + model="gemini-1.5-pro", messages=messages_no_ttl, cache_key="test-no-ttl" ) - + # Verify no TTL field is present when not specified assert "ttl" not in result_no_ttl assert result_no_ttl["model"] == "models/gemini-1.5-pro" assert result_no_ttl["displayName"] == "test-no-ttl" - + # Test case 4: Mixed messages with some having TTL messages_mixed = [ { @@ -180,10 +177,7 @@ def test_gemini_context_caching_with_ttl(): } ], }, - { - "role": "assistant", - "content": "Assistant response without cache control" - }, + {"role": "assistant", "content": "Assistant response without cache control"}, { "role": "user", "content": [ @@ -193,21 +187,19 @@ def test_gemini_context_caching_with_ttl(): "cache_control": {"type": "ephemeral", "ttl": "900s"}, } ], - } + }, ] - + # Test separation of cached messages cached_messages, non_cached_messages = separate_cached_messages(messages_mixed) assert len(cached_messages) > 0 assert len(non_cached_messages) > 0 - + # Test transformation with mixed messages result_mixed = transform_openai_messages_to_gemini_context_caching( - model="gemini-1.5-pro", - messages=messages_mixed, - cache_key="test-mixed-ttl" + model="gemini-1.5-pro", messages=messages_mixed, cache_key="test-mixed-ttl" ) - + # Should pick up the first valid TTL assert "ttl" in result_mixed assert result_mixed["ttl"] == "1800s" @@ -299,13 +291,13 @@ def test_gemini_thinking(): raw_request = return_raw_request( endpoint=CallTypes.completion, kwargs={ - "model": "gemini/gemini-2.5-flash-preview-04-17", + "model": "gemini/gemini-2.5-flash", "messages": messages, }, ) assert reasoning_content in json.dumps(raw_request) response = completion( - model="gemini/gemini-2.5-flash-preview-04-17", + model="gemini/gemini-2.5-flash", messages=messages, # make sure call works ) print(response.choices[0].message) @@ -321,7 +313,7 @@ def test_gemini_thinking_budget_0(): raw_request = return_raw_request( endpoint=CallTypes.completion, kwargs={ - "model": "gemini/gemini-2.5-flash-preview-04-17", + "model": "gemini/gemini-2.5-flash", "messages": [ { "role": "user", 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 612662fd2c..79707359c1 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 @@ -68,7 +68,7 @@ def test_reasoning_tokens_no_price_set(): def test_reasoning_tokens_gemini(): - model = "gemini-2.5-flash-preview-04-17" + model = "gemini-2.5-flash" custom_llm_provider = "gemini" os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" litellm.model_cost = litellm.get_model_cost_map(url="") diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 314c3069b4..ef9b6c1889 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -330,12 +330,7 @@ def test_cost_calculator_with_cache_creation(): def test_bedrock_cost_calculator_comparison_with_without_cache(): """Test that Bedrock caching reduces costs compared to non-cached requests""" from litellm import completion_cost - from litellm.types.utils import ( - Choices, - Message, - PromptTokensDetailsWrapper, - Usage, - ) + from litellm.types.utils import Choices, Message, PromptTokensDetailsWrapper, Usage # Response WITHOUT caching response_no_cache = ModelResponse( @@ -416,7 +411,7 @@ def test_bedrock_cost_calculator_comparison_with_without_cache(): def test_gemini_25_implicit_caching_cost(): """ Test that Gemini 2.5 models correctly calculate costs with implicit caching. - + This test reproduces the issue from #11156 where cached tokens should receive a 75% discount. """ @@ -433,7 +428,7 @@ def test_gemini_25_implicit_caching_cost(): litellm_model_response = ModelResponse( id="test-response", created=1750733889, - model="gemini/gemini-2.5-flash-preview-04-17", + model="gemini/gemini-2.5-flash", object="chat.completion", system_fingerprint=None, choices=[ @@ -463,23 +458,25 @@ def test_gemini_25_implicit_caching_cost(): # Calculate the cost result = completion_cost( completion_response=litellm_model_response, - model="gemini/gemini-2.5-flash-preview-04-17", + model="gemini/gemini-2.5-flash", ) # From the issue: # input: $0.15 / 1000000 tokens # output: $0.60 / 1000000 tokens # With caching: 0.15*0.25*(14316/1000000)+0.15*((15033-14316)/1000000)+0.6*(17/1000000) = 0.0006546 - + # Breakdown: # - Cached tokens: 14316 * 0.15/1M * 0.25 = 0.00053685 # - Non-cached tokens: (15033-14316) * 0.15/1M = 717 * 0.15/1M = 0.00010755 # - Output tokens: 17 * 0.6/1M = 0.00001020 # Total: 0.00053685 + 0.00010755 + 0.00001020 = 0.0006546 - + expected_cost = 0.0006546 - + # Allow for small floating point differences - assert abs(result - expected_cost) < 1e-8, f"Expected cost {expected_cost}, but got {result}" - + assert ( + abs(result - expected_cost) < 1e-8 + ), f"Expected cost {expected_cost}, but got {result}" + print(f"✓ Gemini 2.5 implicit caching cost calculation is correct: ${result:.8f}")