test: update tests

This commit is contained in:
Krrish Dholakia 2025-07-16 08:55:07 -07:00 • committed by Ishaan Jaff
parent ae6d98d221
commit 49e0ddecdb
3 changed files with 7 additions and 12 deletions

View file

@ -23,7 +23,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"""
@ -34,7 +34,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
@ -67,6 +66,7 @@ class TestGoogleAIStudioGemini(BaseLLMChatTest):
print(f"response={response}")
def test_gemini_context_caching_separate_messages():
messages = [
# System Message
@ -151,13 +151,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)
@ -173,7 +173,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",

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="")

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(