tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-10 16:22:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

test: update tests

Browse Source
This commit is contained in:
Krrish Dholakia
2025-07-16 08:56:41 -07:00
parent c0935b9d58
commit 5a8762b6a1
3 changed files with 36 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/llm_translation/test_gemini.py
+24 -32
View File
@@ -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",
tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
+1 -1
View File
@@ -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="")
tests/test_litellm/test_cost_calculator.py
+11 -14
View File
@@ -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}")